Claude-created paginated listing

Author: kylebarron

obstore/src/list.rs

@@ -8,7 +8,7 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
 use futures::stream::{BoxStream, Fuse};
 use futures::StreamExt;
 use indexmap::IndexMap;
-use object_store::list::PaginatedListStore;
+use object_store::list::{PaginatedListOptions, PaginatedListStore};
 use object_store::path::Path;
 use object_store::{ListResult, ObjectMeta, ObjectStore};
 use pyo3::exceptions::{PyImportError, PyStopAsyncIteration, PyStopIteration, PyValueError};
@@ -24,7 +24,7 @@ use pyo3_object_store::{
 };
 use tokio::sync::Mutex;
 
-enum MaybePaginatedListStore {
+pub(crate) enum MaybePaginatedListStore {
     SupportsPagination(Arc<dyn PaginatedListStore>),
     NoPagination(Arc<dyn ObjectStore>),
 }
@@ -452,7 +452,7 @@ impl<'py> IntoPyObject<'py> for PyListResult {
 #[pyo3(signature = (store, prefix=None, *, offset=None, chunk_size=50, return_arrow=false))]
 pub(crate) fn list(
     py: Python,
-    store: PyObjectStore,
+    store: MaybePaginatedListStore,
     prefix: Option<String>,
     offset: Option<String>,
     chunk_size: usize,
@@ -470,12 +470,13 @@ pub(crate) fn list(
             .map_err(|err| PyImportError::new_err(format!("{msg}\n\n{err}")))?;
     }
 
-    let store = store.into_inner().clone();
-    let prefix = prefix.map(|s| s.into());
-    let stream = if let Some(offset) = offset {
-        store.list_with_offset(prefix.as_ref(), &offset.into())
-    } else {
-        store.list(prefix.as_ref())
+    let stream = match store {
+        MaybePaginatedListStore::SupportsPagination(paginated_store) => {
+            create_paginated_stream(paginated_store, prefix, offset, chunk_size)
+        }
+        MaybePaginatedListStore::NoPagination(object_store) => {
+            create_filtered_stream(object_store, prefix, offset)
+        }
     };
     Ok(PyListStream::new(stream, chunk_size, return_arrow))
 }
@@ -526,3 +527,102 @@ async fn list_with_delimiter_materialize(
     let list_result = store.list_with_delimiter(prefix).await?;
     Ok(PyListResult::new(list_result, return_arrow))
 }
+
+fn create_paginated_stream(
+    store: Arc<dyn PaginatedListStore>,
+    prefix: Option<String>,
+    offset: Option<String>,
+    chunk_size: usize,
+) -> BoxStream<'static, object_store::Result<ObjectMeta>> {
+    // Create a stream that will fetch from the paginated store
+    let stream = futures::stream::unfold(
+        (store, prefix, offset, None, true),
+        move |(store, prefix, offset, page_token, has_more)| async move {
+            if !has_more {
+                return None;
+            }
+
+            let opts = PaginatedListOptions {
+                offset: offset.clone(),
+                delimiter: None,
+                max_keys: Some(chunk_size),
+                page_token,
+                ..Default::default()
+            };
+
+            match store.list_paginated(prefix.as_deref(), opts).await {
+                Ok(result) => {
+                    let next_has_more = result.page_token.is_some();
+                    let next_page_token = result.page_token;
+                    let objects = result.result.objects;
+
+                    let next_state = (store, prefix, offset, next_page_token, next_has_more);
+                    Some((objects, next_state))
+                }
+                Err(_e) => {
+                    // TODO: propagate error
+                    // For errors, return empty list and stop
+                    Some((Vec::new(), (store, prefix, offset, None, false)))
+                }
+            }
+        },
+    )
+    .flat_map(|objects| futures::stream::iter(objects.into_iter().map(Ok)));
+
+    Box::pin(stream)
+}
+
+fn create_filtered_stream(
+    store: Arc<dyn ObjectStore>,
+    prefix: Option<String>,
+    offset: Option<String>,
+) -> BoxStream<'static, object_store::Result<ObjectMeta>> {
+    // For substring filtering, we need to split the prefix into:
+    // 1. A directory prefix for efficient listing
+    // 2. A substring filter to apply to the results
+    let (list_prefix, substring_filter) = if let Some(prefix_str) = &prefix {
+        if let Some((dir_prefix, substring)) = prefix_str.rsplit_once('/') {
+            (Some(dir_prefix.to_string()), Some(substring.to_string()))
+        } else {
+            (None, Some(prefix_str.clone()))
+        }
+    } else {
+        (None, None)
+    };
+
+    let prefix_path = list_prefix.map(|s| s.into());
+    let base_stream = if let Some(offset) = offset {
+        store.list_with_offset(prefix_path.as_ref(), &offset.into())
+    } else {
+        store.list(prefix_path.as_ref())
+    };
+
+    // Apply substring filtering if needed
+    let filtered_stream = if let Some(substring) = substring_filter {
+        Box::pin(base_stream.filter_map(move |result| {
+            let substring = substring.clone();
+            async move {
+                match result {
+                    Ok(meta) => {
+                        // Extract filename from path for substring matching
+                        let path_str = meta.location.as_ref();
+                        if let Some(filename) = path_str.split('/').last() {
+                            if filename.contains(&substring) {
+                                Some(Ok(meta))
+                            } else {
+                                None
+                            }
+                        } else {
+                            Some(Ok(meta))
+                        }
+                    }
+                    Err(e) => Some(Err(e)),
+                }
+            }
+        }))
+    } else {
+        base_stream
+    };
+
+    filtered_stream
+}

tests/test_list.py

@@ -1,9 +1,12 @@
+import tempfile
+from pathlib import Path
+
 import polars as pl
 import pyarrow as pa
 import pytest
 from arro3.core import RecordBatch, Table
 
-from obstore.store import MemoryStore
+from obstore.store import LocalStore, MemoryStore
 
 
 def test_list():
@@ -130,6 +133,64 @@ async def test_list_with_delimiter_async():
     assert objects["path"][1].as_py() == "a/file2.txt"
 
 
+def test_list_substring_filtering():
+    store = MemoryStore()
+
+    # Add files with various patterns
+    store.put("data/file1.txt", b"foo")
+    store.put("data/test_file.txt", b"bar")
+    store.put("data/another.csv", b"baz")
+    store.put("data/test_data.json", b"qux")
+    store.put("logs/test_log.txt", b"log")
+
+    # Test substring filtering for files containing "test"
+    result = store.list("data/test").collect()
+    paths = [item["path"] for item in result]
+
+    # Should match files with "test" in the filename within data/ directory
+    assert "data/test_file.txt" in paths
+    assert "data/test_data.json" in paths
+    assert "data/file1.txt" not in paths
+    assert "data/another.csv" not in paths
+    assert "logs/test_log.txt" not in paths
+
+    # Test with arrow format
+    stream = store.list("data/test", return_arrow=True)
+    batch = stream.collect()
+    assert isinstance(batch, RecordBatch)
+    assert batch.num_rows == 2
+
+
+def test_list_substring_filtering_local_store():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_dir_path = Path(temp_dir)
+        store = LocalStore(temp_dir_path)
+
+        # Create directory structure
+        data_dir = temp_dir_path / "data"
+        data_dir.mkdir(parents=True, exist_ok=True)
+
+        # Write test files
+        with (data_dir / "file1.txt").open("w") as f:
+            f.write("foo")
+        with (data_dir / "test_file.txt").open("w") as f:
+            f.write("bar")
+        with (data_dir / "another.csv").open("w") as f:
+            f.write("baz")
+        with (data_dir / "test_data.json").open("w") as f:
+            f.write("qux")
+
+        # Test substring filtering for files containing "test"
+        result = store.list("data/test").collect()
+        paths = [item["path"] for item in result]
+
+        # Should match files with "test" in the filename within data/ directory
+        assert "data/test_file.txt" in paths
+        assert "data/test_data.json" in paths
+        assert "data/file1.txt" not in paths
+        assert "data/another.csv" not in paths
+
+
 def test_list_as_arrow_to_polars():
     store = MemoryStore()