Speedup list_dir and list_prefix (#1257)

* Speedup list_dir and list_prefix

This avoids going through every node each time `list_dir` is called.
This is important because Zarr calls `list_dir` recursively to build
`tree`.

In flat repositories with ~ 60k nodes, we verified ~15x performance
improvement for `Group.tree`.

* Trace log level for xarray backend tests

* Don't use in memory store for xarray backend tests

Author: paraseba

.github/workflows/python-check.yaml

@@ -178,6 +178,7 @@ jobs:
         working-directory: icechunk/icechunk-python
         env:
           ICECHUNK_XARRAY_BACKENDS_TESTS: 1
+          ICECHUNK_LOG: trace
         run: |
           set -e
           # pass xarray's pyproject.toml so that pytest can find the `flaky` fixture

icechunk-python/tests/run_xarray_backends_tests.py

@@ -16,7 +16,6 @@
 from icechunk import (
     IcechunkStore,
     Repository,
-    in_memory_storage,
     local_filesystem_storage,
     s3_storage,
 )
@@ -105,7 +104,8 @@ def create_repo(self) -> Generator[Repository]:
 class TestIcechunkStoreMemory(IcechunkStoreBase):
     @contextlib.contextmanager
     def create_repo(self) -> Generator[Repository]:
-        yield Repository.create(in_memory_storage())
+        with tempfile.TemporaryDirectory() as tmpdir:
+            yield Repository.create(local_filesystem_storage(tmpdir))
 
     def test_pickle(self) -> None:
         pytest.skip("pickling memory store is not supported")
@@ -137,9 +137,11 @@ class TestIcechunkRegionAuto(ZarrRegionAutoTests):
     def create_zarr_target(self) -> Generator[IcechunkStore]:
         if zarr.config.config["default_zarr_format"] == 2:
             pytest.skip("v2 not supported")
-        repo = Repository.create(in_memory_storage())
-        session = repo.writable_session("main")
-        yield session.store
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            repo = Repository.create(local_filesystem_storage(tmpdir))
+            session = repo.writable_session("main")
+            yield session.store
 
     @contextlib.contextmanager
     def create(self):
@@ -149,11 +151,12 @@ def create(self):
         ds = xr.Dataset(
             {"test": xr.DataArray(data, dims=("x", "y"), coords={"x": x, "y": y})}
         )
-        repo = Repository.create(in_memory_storage())
-        session = repo.writable_session("main")
-        self.save(session.store, ds)
-        session.commit("initial commit")
-        yield repo.writable_session("main").store, ds
+        with tempfile.TemporaryDirectory() as tmpdir:
+            repo = Repository.create(local_filesystem_storage(tmpdir))
+            session = repo.writable_session("main")
+            self.save(session.store, ds)
+            session.commit("initial commit")
+            yield repo.writable_session("main").store, ds
 
     def save(self, target, ds, **kwargs):
         # not really important here

icechunk/examples/low_level_dataset.rs

@@ -237,7 +237,7 @@ let ds = Repository::update(Arc::clone(&storage), ObjectId::from("{v2_id:?}"));
 async fn print_nodes(ds: &Session) -> Result<(), SessionError> {
     println!("### List of nodes");
     let rows = ds
-        .list_nodes()
+        .list_nodes(&Path::root())
         .await?
         .map(|n| n.unwrap())
         .sorted_by_key(|n| n.path.clone())

icechunk/src/conflicts/detector.rs

@@ -63,7 +63,7 @@ impl ConflictSolver for ConflictDetector {
             Ok(None)
         });
 
-        let path_finder = PathFinder::new(current_repo.list_nodes().await?);
+        let path_finder = PathFinder::new(current_repo.list_nodes(&Path::root()).await?);
 
         let updated_arrays_already_updated = current_changes
             .updated_arrays()

icechunk/src/format/mod.rs

@@ -1,5 +1,6 @@
 use core::fmt;
 use std::{
+    cmp::Ordering,
     convert::Infallible,
     fmt::{Debug, Display},
     hash::Hash,
@@ -462,6 +463,37 @@ impl TryFrom<String> for Path {
     }
 }
 
+#[inline(always)]
+pub fn lookup_index_by_key<'a, T: ::flatbuffers::Follow<'a> + 'a, K: Ord>(
+    v: ::flatbuffers::Vector<'a, T>,
+    key: K,
+    f: fn(&<T as ::flatbuffers::Follow<'a>>::Inner, &K) -> Ordering,
+) -> Option<usize> {
+    if v.is_empty() {
+        return None;
+    }
+
+    let mut left: usize = 0;
+    let mut right = v.len() - 1;
+
+    while left <= right {
+        let mid = (left + right) / 2;
+        let value = v.get(mid);
+        match f(&value, &key) {
+            Ordering::Equal => return Some(mid),
+            Ordering::Less => left = mid + 1,
+            Ordering::Greater => {
+                if mid == 0 {
+                    return None;
+                }
+                right = mid - 1;
+            }
+        }
+    }
+
+    None
+}
+
 #[cfg(test)]
 #[allow(clippy::panic, clippy::unwrap_used, clippy::expect_used)]
 mod tests {

icechunk/src/format/snapshot.rs

@@ -8,6 +8,8 @@ use itertools::Itertools as _;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 
+use crate::format::lookup_index_by_key;
+
 use super::{
     AttributesId, ChunkIndices, IcechunkFormatError, IcechunkFormatErrorKind,
     IcechunkResult, ManifestId, NodeId, Path, SnapshotId,
@@ -524,14 +526,27 @@ impl Snapshot {
         res.try_into()
     }
 
+    pub fn get_node_index(&self, path: &Path) -> IcechunkResult<usize> {
+        let path_str = path.to_string();
+        let res =
+            lookup_index_by_key(self.root().nodes(), path_str.as_str(), |node, path| {
+                node.path().cmp(path)
+            })
+            .ok_or(IcechunkFormatError::from(
+                IcechunkFormatErrorKind::NodeNotFound { path: path.clone() },
+            ))?;
+        Ok(res)
+    }
+
     pub fn iter(&self) -> impl Iterator<Item = IcechunkResult<NodeSnapshot>> + '_ {
         self.root().nodes().iter().map(|node| node.try_into().err_into())
     }
 
     pub fn iter_arc(
         self: Arc<Self>,
-    ) -> impl Iterator<Item = IcechunkResult<NodeSnapshot>> {
-        NodeIterator { snapshot: self, last_index: 0 }
+        parent_group: &Path,
+    ) -> impl Iterator<Item = IcechunkResult<NodeSnapshot>> + use<> {
+        NodeIterator::new(self, parent_group)
     }
 
     pub fn len(&self) -> usize {
@@ -554,18 +569,42 @@ impl Snapshot {
 
 struct NodeIterator {
     snapshot: Arc<Snapshot>,
-    last_index: usize,
+    next_index: usize,
+    prefix: String,
+}
+
+impl NodeIterator {
+    fn new(snapshot: Arc<Snapshot>, parent_group: &Path) -> Self {
+        let next_index = snapshot.get_node_index(parent_group).unwrap_or_default();
+        let prefix = parent_group.to_string();
+        let prefix = if prefix == "/" { String::new() } else { prefix };
+        NodeIterator { snapshot, next_index, prefix }
+    }
 }
 
 impl Iterator for NodeIterator {
     type Item = IcechunkResult<NodeSnapshot>;
 
     fn next(&mut self) -> Option<Self::Item> {
         let nodes = self.snapshot.root().nodes();
-        if self.last_index < nodes.len() {
-            let res = Some(nodes.get(self.last_index).try_into().err_into());
-            self.last_index += 1;
-            res
+        if self.next_index < nodes.len() {
+            let node: IcechunkResult<NodeSnapshot> =
+                nodes.get(self.next_index).try_into();
+
+            match node {
+                Ok(res) => {
+                    if let Some(after_prefix) =
+                        res.path.to_string().strip_prefix(self.prefix.as_str())
+                        && (after_prefix.is_empty() || after_prefix.starts_with('/'))
+                    {
+                        self.next_index += 1;
+                        Some(Ok(res))
+                    } else {
+                        None
+                    }
+                }
+                Err(err) => Some(Err(err)),
+            }
         } else {
             None
         }

icechunk/src/format/transaction_log.rs

@@ -260,9 +260,9 @@ impl DiffBuilder {
 
     pub async fn to_diff(self, from: &Session, to: &Session) -> SessionResult<Diff> {
         let nodes: HashMap<NodeId, Path> = from
-            .list_nodes()
+            .list_nodes(&Path::root())
             .await?
-            .chain(to.list_nodes().await?)
+            .chain(to.list_nodes(&Path::root()).await?)
             .map_ok(|n| (n.id, n.path))
             .try_collect()?;
         Ok(Diff::from_diff_builder(self, nodes))

icechunk/src/repository.rs

@@ -860,7 +860,7 @@ impl Repository {
             if let Ok(snap) = asset_manager.fetch_snapshot(&snapshot_id).await {
                 let snap_c = Arc::clone(&snap);
                 for node in snap
-                    .iter_arc()
+                    .iter_arc(&Path::root())
                     .filter_ok(|node| node.node_type() == NodeType::Array)
                     // TODO: make configurable
                     .take(50)

icechunk/src/session.rs

@@ -372,7 +372,7 @@ impl Session {
         match self.get_group(&path).await {
             Ok(parent) => {
                 let nodes_iter: Vec<NodeSnapshot> = self
-                    .list_nodes()
+                    .list_nodes(&path)
                     .await?
                     .filter_ok(|node| node.path.starts_with(&parent.path))
                     .try_collect()?;
@@ -796,7 +796,7 @@ impl Session {
     pub async fn clear(&mut self) -> SessionResult<()> {
         // TODO: can this be a delete_group("/") instead?
         let to_delete: Vec<(NodeType, Path)> = self
-            .list_nodes()
+            .list_nodes(&Path::root())
             .await?
             .map_ok(|node| (node.node_type(), node.path))
             .try_collect()?;
@@ -850,10 +850,17 @@ impl Session {
     }
 
     #[instrument(skip(self))]
-    pub async fn list_nodes(
-        &self,
-    ) -> SessionResult<impl Iterator<Item = SessionResult<NodeSnapshot>> + '_> {
-        updated_nodes(&self.asset_manager, &self.change_set, &self.snapshot_id).await
+    pub async fn list_nodes<'a>(
+        &'a self,
+        parent_group: &Path,
+    ) -> SessionResult<impl Iterator<Item = SessionResult<NodeSnapshot>> + use<'a>> {
+        updated_nodes(
+            parent_group,
+            &self.asset_manager,
+            &self.change_set,
+            &self.snapshot_id,
+        )
+        .await
     }
 
     #[instrument(skip(self))]
@@ -975,7 +982,7 @@ impl Session {
         message: &str,
         properties: Option<SnapshotProperties>,
     ) -> SessionResult<SnapshotId> {
-        let nodes = self.list_nodes().await?.collect::<Vec<_>>();
+        let nodes = self.list_nodes(&Path::root()).await?.collect::<Vec<_>>();
         // We need to populate the `splits` before calling `commit`.
         // In the normal chunk setting workflow, that would've been done by `set_chunk_ref`
         for node in nodes.into_iter().flatten() {
@@ -1259,12 +1266,13 @@ impl Session {
 
 /// Warning: The presence of a single error may mean multiple missing items
 async fn updated_chunk_iterator<'a>(
+    parent_group: &Path,
     asset_manager: &'a AssetManager,
     change_set: &'a ChangeSet,
     snapshot_id: &'a SnapshotId,
-) -> SessionResult<impl Stream<Item = SessionResult<(Path, ChunkInfo)>> + 'a> {
+) -> SessionResult<impl Stream<Item = SessionResult<(Path, ChunkInfo)>> + use<'a>> {
     let snapshot = asset_manager.fetch_snapshot(snapshot_id).await?;
-    let nodes = futures::stream::iter(snapshot.iter_arc());
+    let nodes = futures::stream::iter(snapshot.iter_arc(parent_group));
     let res = nodes.and_then(move |node| async move {
         // Note: Confusingly, these NodeSnapshot instances have the metadata stored in the snapshot.
         // We have not applied any changeset updates. At the moment, the downstream code only
@@ -1477,14 +1485,15 @@ pub async fn get_chunk(
 
 /// Yields nodes in the base snapshot, applying any relevant updates in the changeset
 async fn updated_existing_nodes<'a>(
+    parent_group: &Path,
     asset_manager: &AssetManager,
     change_set: &'a ChangeSet,
     parent_id: &SnapshotId,
-) -> SessionResult<impl Iterator<Item = SessionResult<NodeSnapshot>> + 'a + use<'a>> {
+) -> SessionResult<impl Iterator<Item = SessionResult<NodeSnapshot>> + use<'a>> {
     let updated_nodes = asset_manager
         .fetch_snapshot(parent_id)
         .await?
-        .iter_arc()
+        .iter_arc(parent_group)
         .filter_map_ok(move |node| change_set.update_existing_node(node))
         .map(|n| match n {
             Ok(n) => Ok(n),
@@ -1497,11 +1506,12 @@ async fn updated_existing_nodes<'a>(
 /// Yields nodes with the snapshot, applying any relevant updates in the changeset,
 /// *and* new nodes in the changeset
 async fn updated_nodes<'a>(
+    parent_group: &Path,
     asset_manager: &AssetManager,
     change_set: &'a ChangeSet,
     parent_id: &SnapshotId,
-) -> SessionResult<impl Iterator<Item = SessionResult<NodeSnapshot>> + 'a + use<'a>> {
-    Ok(updated_existing_nodes(asset_manager, change_set, parent_id)
+) -> SessionResult<impl Iterator<Item = SessionResult<NodeSnapshot>> + use<'a>> {
+    Ok(updated_existing_nodes(parent_group, asset_manager, change_set, parent_id)
         .await?
         .chain(change_set.new_nodes_iterator().map(Ok)))
 }
@@ -1586,7 +1596,8 @@ async fn all_chunks<'a>(
     snapshot_id: &'a SnapshotId,
 ) -> SessionResult<impl Stream<Item = SessionResult<(Path, ChunkInfo)>> + 'a> {
     let existing_array_chunks =
-        updated_chunk_iterator(asset_manager, change_set, snapshot_id).await?;
+        updated_chunk_iterator(&Path::root(), asset_manager, change_set, snapshot_id)
+            .await?;
     let new_array_chunks =
         futures::stream::iter(change_set.new_arrays_chunk_iterator().map(Ok));
     Ok(existing_array_chunks.chain(new_array_chunks))
@@ -2014,6 +2025,7 @@ async fn flush(
     // gather and sort nodes:
     // this is a requirement for Snapshot::from_iter
     let mut all_nodes: Vec<_> = updated_nodes(
+        &Path::root(),
         flush_data.asset_manager.as_ref(),
         flush_data.change_set,
         flush_data.parent_id,
@@ -3213,15 +3225,15 @@ mod tests {
         ds.add_group(Path::root(), Bytes::copy_from_slice(b"")).await?;
         ds.add_group("/1".try_into().unwrap(), Bytes::copy_from_slice(b"")).await?;
         ds.delete_group("/1".try_into().unwrap()).await?;
-        assert_eq!(ds.list_nodes().await?.count(), 1);
+        assert_eq!(ds.list_nodes(&Path::root()).await?.count(), 1);
         ds.commit("commit", None).await?;
 
         let ds = repository
             .readonly_session(&VersionInfo::BranchTipRef("main".to_string()))
             .await?;
         assert!(ds.get_group(&Path::root()).await.is_ok());
         assert!(ds.get_group(&"/1".try_into().unwrap()).await.is_err());
-        assert_eq!(ds.list_nodes().await?.count(), 1);
+        assert_eq!(ds.list_nodes(&Path::root()).await?.count(), 1);
         Ok(())
     }
 
@@ -3237,7 +3249,7 @@ mod tests {
         ds.delete_group("/1".try_into().unwrap()).await?;
         assert!(ds.get_group(&Path::root()).await.is_ok());
         assert!(ds.get_group(&"/1".try_into().unwrap()).await.is_err());
-        assert_eq!(ds.list_nodes().await?.count(), 1);
+        assert_eq!(ds.list_nodes(&Path::root()).await?.count(), 1);
         Ok(())
     }
 
@@ -3255,7 +3267,7 @@ mod tests {
         let ds = repository
             .readonly_session(&VersionInfo::BranchTipRef("main".to_string()))
             .await?;
-        assert_eq!(ds.list_nodes().await?.count(), 0);
+        assert_eq!(ds.list_nodes(&Path::root()).await?.count(), 0);
         Ok(())
     }