Speedup list_prefix and getsize_prefix when prefix is not empty (#1245)

* Speedup `list_prefix` and `getsize_prefix` when prefix is not empty

These two operations were very inefficient in repositories with a large
number of arrays. This PR implements pre-filtering of chunks, so we only
need to open manifests for the arrays that match the prefix.

Also in this PR: limit to 50 the number of preloaded array manifests.
Making this number configurable is an exercise left to a future PR.

* Fail listing chunks if prefix is not a node

* Add `X-Tigris-Consistent:true` header to the Tigris Storage

* Enable debug logs in GHA test run

* Back to old tigris endpoint

* Back to new tigris endpoint

Author: paraseba

.github/workflows/rust-ci.yaml

@@ -90,6 +90,7 @@ jobs:
       - name: Run tests with Docker services
         if: matrix.os == 'ubuntu-latest' || matrix.os == 'ubuntu-24.04-arm'
         env:
+          RUST_LOG: trace
           R2_BUCKET: ${{ secrets.R2_BUCKET }}
           R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
           R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}

icechunk/src/repository.rs

@@ -13,6 +13,7 @@ use futures::{
     Stream, StreamExt, TryStreamExt,
     stream::{FuturesOrdered, FuturesUnordered},
 };
+use itertools::Itertools;
 use regex::bytes::Regex;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
@@ -28,7 +29,8 @@ use crate::{
         IcechunkFormatError, IcechunkFormatErrorKind, ManifestId, NodeId, Path,
         SnapshotId,
         snapshot::{
-            ManifestFileInfo, NodeData, Snapshot, SnapshotInfo, SnapshotProperties,
+            ManifestFileInfo, NodeData, NodeType, Snapshot, SnapshotInfo,
+            SnapshotProperties,
         },
         transaction_log::{Diff, DiffBuilder},
     },
@@ -857,7 +859,12 @@ impl Repository {
             // TODO: unnest this code
             if let Ok(snap) = asset_manager.fetch_snapshot(&snapshot_id).await {
                 let snap_c = Arc::clone(&snap);
-                for node in snap.iter_arc() {
+                for node in snap
+                    .iter_arc()
+                    .filter_ok(|node| node.node_type() == NodeType::Array)
+                    // TODO: make configurable
+                    .take(50)
+                {
                     match node {
                         Err(err) => {
                             error!(error=%err, "Error retrieving snapshot nodes");

icechunk/src/session.rs

@@ -863,6 +863,21 @@ impl Session {
         all_chunks(&self.asset_manager, &self.change_set, self.snapshot_id()).await
     }
 
+    #[instrument(skip(self))]
+    pub async fn all_node_chunks<'a>(
+        &'a self,
+        node: &'a NodeSnapshot,
+    ) -> impl Stream<Item = SessionResult<(Path, ChunkInfo)>> + 'a {
+        updated_node_chunks_iterator(
+            self.asset_manager.as_ref(),
+            &self.change_set,
+            &self.snapshot_id,
+            node.clone(),
+            ManifestExtents::ALL,
+        )
+        .await
+    }
+
     #[instrument(skip(self))]
     pub async fn chunk_coordinates<'a, 'b: 'a>(
         &'a self,
@@ -1423,8 +1438,7 @@ impl From<Session> for ChangeSet {
 }
 
 pub fn is_prefix_match(key: &str, prefix: &str) -> bool {
-    let tomatch =
-        if prefix != String::from('/') { key.strip_prefix(prefix) } else { Some(key) };
+    let tomatch = if prefix != "/" { key.strip_prefix(prefix) } else { Some(key) };
     match tomatch {
         None => false,
         Some(rest) => {

icechunk/src/storage/mod.rs

@@ -907,22 +907,26 @@ pub fn new_tigris_storage(
         ),
         ..config
     };
-    let mut extra_write_headers = Vec::with_capacity(1);
-    let mut extra_read_headers = Vec::with_capacity(2);
+    let mut extra_write_headers = Vec::with_capacity(2);
+    let mut extra_read_headers = Vec::with_capacity(3);
 
     if !use_weak_consistency {
         // TODO: Tigris will need more than this to offer good eventually consistent behavior
         // For example: we should use no-cache for branches and config file
         if let Some(region) = config.region.as_ref() {
-            extra_write_headers.push(("X-Tigris-Region".to_string(), region.clone()));
-            extra_read_headers.push(("X-Tigris-Region".to_string(), region.clone()));
+            extra_write_headers.push(("X-Tigris-Regions".to_string(), region.clone()));
+            extra_write_headers
+                .push(("X-Tigris-Consistent".to_string(), "true".to_string()));
+
+            extra_read_headers.push(("X-Tigris-Regions".to_string(), region.clone()));
             extra_read_headers
                 .push(("Cache-Control".to_string(), "no-cache".to_string()));
+            extra_read_headers
+                .push(("X-Tigris-Consistent".to_string(), "true".to_string()));
         } else {
             return Err(StorageErrorKind::Other("Tigris storage requires a region to provide full consistency. Either set the region for the bucket or use the read-only, eventually consistent storage by passing `use_weak_consistency=True` (experts only)".to_string()).into());
         }
     }
-
     let st = S3Storage::new(
         config,
         bucket,

icechunk/src/store.rs

@@ -24,7 +24,7 @@ use crate::{
     format::{
         ByteRange, ChunkIndices, ChunkOffset, Path, PathError,
         manifest::{ChunkPayload, VirtualChunkRef},
-        snapshot::{ArrayShape, DimensionName, NodeData, NodeSnapshot},
+        snapshot::{ArrayShape, DimensionName, NodeData, NodeSnapshot, NodeType},
     },
     refs::{RefError, RefErrorKind},
     repository::{RepositoryError, RepositoryErrorKind},
@@ -82,8 +82,8 @@ pub enum StoreErrorKind {
     SerializationError(#[from] Box<rmp_serde::encode::Error>),
     #[error("store method `{0}` is not implemented by Icechunk")]
     Unimplemented(&'static str),
-    #[error("bad key prefix: `{0}`")]
-    BadKeyPrefix(String),
+    #[error("bad key prefix ({prefix}): {message}")]
+    BadKeyPrefix { prefix: String, message: String },
     #[error("error during parallel execution of get_partial_values")]
     PartialValuesPanic,
     #[error("cannot write to read-only store")]
@@ -442,11 +442,14 @@ impl Store {
         if self.read_only().await {
             return Err(StoreErrorKind::ReadOnly.into());
         }
-        let prefix = prefix.trim_start_matches("/").trim_end_matches("/");
+        let prefix = prefix.trim_start_matches('/').trim_end_matches('/');
         // TODO: Handling preceding "/" is ugly!
-        let path = format!("/{prefix}")
-            .try_into()
-            .map_err(|_| StoreErrorKind::BadKeyPrefix(prefix.to_owned()))?;
+        let path = format!("/{prefix}").try_into().map_err(|_| {
+            StoreErrorKind::BadKeyPrefix {
+                prefix: prefix.to_owned(),
+                message: "Cannot convert to a path".to_string(),
+            }
+        })?;
 
         let mut guard = self.session.write().await;
         let node = guard.get_node(&path).await;
@@ -557,8 +560,6 @@ impl Store {
         &self,
         prefix: &str,
     ) -> StoreResult<impl Stream<Item = StoreResult<String>> + Send + use<>> {
-        // TODO: this is inefficient because it filters based on the prefix, instead of only
-        // generating items that could potentially match
         let meta = self.list_metadata_prefix(prefix, false).await?;
         let chunks = self.list_chunks_prefix(prefix).await?;
         // FIXME: this is wrong, we are realizing all keys in memory
@@ -638,13 +639,13 @@ impl Store {
                                 {
                                     let trimmed = chunk_key
                                         .trim_start_matches(prefix)
-                                        .trim_start_matches("/");
+                                        .trim_start_matches('/');
                                     if trimmed.is_empty() {
                                         // we were provided with a prefix that is a path to a chunk key
                                         None
                                     } else if let Some((chunk_prefix, _)) =
                                         // if we can split it, this is a valid prefix to return
-                                        trimmed.split_once("/")
+                                        trimmed.split_once('/')
                                     {
                                         Some(ListDirItem::Prefix(
                                             chunk_prefix.to_string(),
@@ -748,13 +749,13 @@ impl Store {
     ) -> StoreResult<impl Stream<Item = StoreResult<String>> + 'a + use<'a>> {
         let prefix = prefix.trim_end_matches('/');
         let res = try_stream! {
-            let repository = Arc::clone(&self.session).read_owned().await;
-            for node in repository.list_nodes().await? {
+            let session = Arc::clone(&self.session).read_owned().await;
+            for node in session.list_nodes().await? {
                 // TODO: handle non-utf8?
                 let meta_key = Key::Metadata { node_path: node?.path }.to_string();
                 if is_prefix_match(&meta_key, prefix) {
                     if strip_prefix {
-                        yield meta_key.trim_start_matches(prefix).trim_start_matches("/").to_string();
+                        yield meta_key.trim_start_matches(prefix).trim_start_matches('/').to_string();
                     } else {
                         yield meta_key;
                     }
@@ -770,19 +771,40 @@ impl Store {
     ) -> StoreResult<impl Stream<Item = StoreResult<String>> + 'a + use<'a>> {
         let prefix = prefix.trim_end_matches('/');
         let res = try_stream! {
-            let repository = Arc::clone(&self.session).read_owned().await;
-            // TODO: this is inefficient because it filters based on the prefix, instead of only
-            // generating items that could potentially match
-            for await maybe_path_chunk in repository.all_chunks().await.map_err(StoreError::from)? {
-                // FIXME: utf8 handling
-                match maybe_path_chunk {
-                    Ok((path, chunk)) => {
-                        let chunk_key = Key::Chunk { node_path: path, coords: chunk.coord }.to_string();
-                        if is_prefix_match(&chunk_key, prefix) {
-                            yield chunk_key;
-                        }
+            let session = Arc::clone(&self.session).read_owned().await;
+
+            // TODO: Handling preceding "/" is ugly!
+            let path = format!("/{}", prefix.trim_start_matches('/') ).try_into().map_err(|_| {
+                StoreErrorKind::BadKeyPrefix {
+                    prefix: prefix.to_owned(),
+                    message: "Cannot convert to a path".to_string(),
+                }
+            })?;
+
+            if path != Path::root() {
+                let _ = session.get_node(&path).await.map_err(|_| {
+                    StoreErrorKind::BadKeyPrefix {
+                        prefix: prefix.to_owned(),
+                        message: "Only prefixes pointing to a group or array are allowed".to_string(),
                     }
-                    Err(err) => Err(err)?
+                })?;
+            }
+
+            for node in session.list_nodes().await? {
+                let node = node?;
+                if node.node_type() == NodeType::Array &&
+                    // FIXME: utf8 handling
+                    // skip the initial / in the path
+                    is_prefix_match(&node.path.to_string()[1..], prefix) {
+                        for await maybe_path_chunk in session.all_node_chunks(&node).await {
+                            match maybe_path_chunk {
+                                Ok((path, chunk)) => {
+                                    let chunk_key = Key::Chunk { node_path: path, coords: chunk.coord }.to_string();
+                                    yield chunk_key;
+                                }
+                                Err(err) => Err(err)?
+                            }
+                        }
                 }
             }
         };
@@ -1227,10 +1249,7 @@ mod tests {
         Ok(version1)
     }
 
-    async fn keys(
-        store: &Store,
-        prefix: &str,
-    ) -> Result<Vec<String>, Box<dyn std::error::Error>> {
+    async fn keys(store: &Store, prefix: &str) -> Result<Vec<String>, StoreError> {
         let mut res = store.list_prefix(prefix).await?.try_collect::<Vec<_>>().await?;
         res.sort();
         Ok(res)
@@ -1685,6 +1704,7 @@ mod tests {
         assert!(!store.exists("zarr.json").await.unwrap());
 
         assert_eq!(all_keys(&store).await.unwrap(), Vec::<String>::new());
+        assert_eq!(all_keys(&store).await.unwrap(), keys(&store, "").await.unwrap());
         store
             .set(
                 "zarr.json",
@@ -1695,6 +1715,7 @@ mod tests {
         assert!(!store.is_empty("").await.unwrap());
         assert!(store.exists("zarr.json").await.unwrap());
         assert_eq!(all_keys(&store).await.unwrap(), vec!["zarr.json".to_string()]);
+        assert_eq!(all_keys(&store).await.unwrap(), keys(&store, "").await.unwrap());
         store
             .set(
                 "group/zarr.json",
@@ -1705,6 +1726,7 @@ mod tests {
             all_keys(&store).await.unwrap(),
             vec!["group/zarr.json".to_string(), "zarr.json".to_string()]
         );
+        assert_eq!(all_keys(&store).await.unwrap(), keys(&store, "").await.unwrap());
         assert_eq!(
             keys(&store, "group/").await.unwrap(),
             vec!["group/zarr.json".to_string()]
@@ -1724,6 +1746,7 @@ mod tests {
                 "zarr.json".to_string()
             ]
         );
+        assert_eq!(all_keys(&store).await.unwrap(), keys(&store, "").await.unwrap());
         assert_eq!(
             keys(&store, "group/").await.unwrap(),
             vec!["group/array/zarr.json".to_string(), "group/zarr.json".to_string()]
@@ -1797,6 +1820,7 @@ mod tests {
                 "zarr.json".to_string()
             ]
         );
+        assert_eq!(all_keys(&store).await.unwrap(), keys(&store, "").await.unwrap());
 
         session.write().await.commit("foo", None).await?;
 
@@ -1825,6 +1849,17 @@ mod tests {
                 "zarr.json".to_string()
             ]
         );
+        assert_eq!(all_keys(&store).await.unwrap(), keys(&store, "").await.unwrap());
+
+        for bad_prefix in ["arr", "/arr", "/arr/", "zarr", "array/c", "array/c/0"] {
+            assert!(matches!(
+                keys(&store, bad_prefix).await,
+                Err(StoreError {
+                    kind: StoreErrorKind::BadKeyPrefix { message, .. },
+                    ..
+                }) if message.contains("group or array")
+            ));
+        }
 
         Ok(())
     }