diff --git a/Cargo.lock b/Cargo.lock index 3f177d98..51372dcc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1367,6 +1367,7 @@ name = "obstore" version = "0.10.0" dependencies = [ "arrow", + "async-trait", "bytes", "cargo-lock", "chrono", diff --git a/Cargo.toml b/Cargo.toml index 3e894fc9..dd4ef8c0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,6 +21,7 @@ categories = [] rust-version = "1.75" [workspace.dependencies] +async-trait = "0.1.85" bytes = "1.10.1" chrono = "0.4.44" futures = "0.3.31" diff --git a/obstore/Cargo.toml b/obstore/Cargo.toml index 49aba300..e06b7706 100644 --- a/obstore/Cargo.toml +++ b/obstore/Cargo.toml @@ -23,6 +23,7 @@ abi3-py311 = ["pyo3/abi3-py311"] generate-import-lib = ["pyo3/generate-import-lib"] [dependencies] +async-trait = { workspace = true } arrow = "58" bytes = { workspace = true } chrono = { workspace = true } diff --git a/obstore/src/get.rs b/obstore/src/get.rs index afdcc978..acad905e 100644 --- a/obstore/src/get.rs +++ b/obstore/src/get.rs @@ -7,8 +7,8 @@ use chrono::{DateTime, Utc}; use futures::stream::{BoxStream, Fuse}; use futures::StreamExt; use object_store::{ - coalesce_ranges, Attributes, GetOptions, GetRange, GetResult, ObjectMeta, ObjectStore, - ObjectStoreExt, OBJECT_STORE_COALESCE_DEFAULT, + coalesce_ranges, Attributes, GetOptions, GetRange, GetResult, ObjectMeta, ObjectStoreExt, + OBJECT_STORE_COALESCE_DEFAULT, }; use pyo3::exceptions::{PyStopAsyncIteration, PyStopIteration, PyValueError}; use pyo3::prelude::*; diff --git a/obstore/src/list.rs b/obstore/src/list.rs index 253787c2..0dc96ebe 100644 --- a/obstore/src/list.rs +++ b/obstore/src/list.rs @@ -1,3 +1,4 @@ +use async_trait::async_trait; use std::ops::AddAssign; use std::sync::Arc; @@ -8,6 +9,7 @@ use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; use futures::stream::{BoxStream, Fuse}; use futures::StreamExt; use indexmap::IndexMap; +use object_store::list::{PaginatedListOptions, PaginatedListResult, PaginatedListStore}; use object_store::{ListResult, ObjectMeta, ObjectStore}; use pyo3::exceptions::{PyImportError, PyStopAsyncIteration, PyStopIteration}; use pyo3::prelude::*; @@ -347,13 +349,92 @@ impl<'py> IntoPyObject<'py> for PyListResult { } } +enum MaybePaginatedStore { + /// Stores that natively support pagination + Native(Arc), + /// Paginated stores emulated by collecting all results and filtering prefix in memory + Emulated(Arc), +} + +impl From for MaybePaginatedStore { + fn from(store: PyObjectStore) -> Self { + match store { + PyObjectStore::S3(store) => Self::Native(store.into_inner()), + PyObjectStore::Azure(store) => Self::Native(store.into_inner()), + PyObjectStore::Gcs(store) => Self::Native(store.into_inner()), + PyObjectStore::Http(store) => Self::Emulated(store.into_inner()), + PyObjectStore::Local(store) => Self::Emulated(store.into_inner()), + PyObjectStore::Memory(store) => Self::Emulated(store.into_inner()), + } + } +} + +/// A custom implementation of PaginatedListStore for stores that don't natively support +/// pagination, like HttpStore and LocalStore. +/// +/// PaginatedListStore is not implemented in upstream for LocalFileSystem because there's no way to +/// get a stable offset in local FS APIs. +/// https://github.com/apache/arrow-rs-object-store/issues/388 +/// +/// Instead, we collect _all_ results and filter them in memory with the provided substring. +async fn emulate_paginated_list( + store: &Arc, + prefix: Option<&str>, +) -> object_store::Result { + // `PaginatedListStore` treats `prefix` as a raw string prefix (substring-style), + // whereas `ObjectStore::list` matches on whole path segments. To emulate the former + // with the latter, we list recursively under the last complete path segment and then + // keep only the keys whose full location starts with the requested prefix string. + // + // e.g. prefix "data/tes" -> list recursively under "data/", then keep keys starting + // with "data/tes" (including nested keys like "data/test/deep/file.txt"). + let list_path = match prefix.and_then(|prefix| prefix.rsplit_once('/')) { + // List recursively under the path portion before the final '/'. + Some((dir, _)) => Some(object_store::path::Path::parse(dir)?), + // No '/' in the prefix (or no prefix at all): list from the root. + None => None, + }; + + let mut stream = store.list(list_path.as_ref()); + let mut objects = Vec::new(); + while let Some(meta) = stream.next().await.transpose()? { + match prefix { + Some(prefix) if !meta.location.as_ref().starts_with(prefix) => continue, + _ => objects.push(meta), + } + } + + Ok(PaginatedListResult { + result: ListResult { + common_prefixes: Vec::new(), + objects, + }, + // Emulated stores return everything in a single page. + page_token: None, + }) +} + +#[async_trait] +impl PaginatedListStore for MaybePaginatedStore { + async fn list_paginated( + &self, + prefix: Option<&str>, + opts: PaginatedListOptions, + ) -> object_store::Result { + match self { + Self::Native(store) => store.list_paginated(prefix, opts).await, + Self::Emulated(store) => emulate_paginated_list(store, prefix).await, + } + } +} + #[pyfunction] #[pyo3(signature = (store, prefix=None, *, offset=None, chunk_size=50, return_arrow=false))] pub(crate) fn list( py: Python, store: PyObjectStore, - prefix: Option, - offset: Option, + prefix: Option, + offset: Option, chunk_size: usize, return_arrow: bool, ) -> PyObjectStoreResult { @@ -369,13 +450,8 @@ 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.as_ref()) - } else { - store.list(prefix.as_ref()) - }; + let maybe_paginated_store = Arc::new(MaybePaginatedStore::from(store)); + let stream = create_paginated_stream(maybe_paginated_store, prefix, offset, chunk_size); Ok(PyListStream::new(stream, chunk_size, return_arrow)) } @@ -424,3 +500,89 @@ async fn list_with_delimiter_materialize( .await?; Ok(PyListResult::new(list_result, return_arrow)) } + +/// Internal stream state +struct StreamState { + store: Arc, + prefix: Option, + offset: Option, + page_token: Option, + has_more: bool, +} + +async fn stream_step( + state: StreamState, + chunk_size: usize, +) -> Option<(Vec>, StreamState)> { + let StreamState { + store, + prefix, + offset, + page_token, + has_more, + } = state; + + 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: Vec> = + result.result.objects.into_iter().map(Ok).collect(); + + let next_state = StreamState { + store, + prefix, + offset, + page_token: next_page_token, + has_more: next_has_more, + }; + Some((objects, next_state)) + } + // Surface the error to the consumer as a stream item, then stop the stream + // so we don't silently truncate results on a failed page. + Err(e) => Some(( + vec![Err(e)], + StreamState { + store, + prefix, + offset, + page_token: None, + has_more: false, + }, + )), + } +} + +fn create_paginated_stream( + store: Arc, + prefix: Option, + offset: Option, + chunk_size: usize, +) -> BoxStream<'static, object_store::Result> { + // Create a stream that will fetch from the paginated store + let stream = futures::stream::unfold( + StreamState { + store, + prefix, + offset, + page_token: None, + has_more: true, + }, + move |state| stream_step(state, chunk_size), + ) + .flat_map(futures::stream::iter); + + Box::pin(stream) +} diff --git a/pyo3-object_store/src/aws/store.rs b/pyo3-object_store/src/aws/store.rs index f8041c4a..56826bb7 100644 --- a/pyo3-object_store/src/aws/store.rs +++ b/pyo3-object_store/src/aws/store.rs @@ -65,7 +65,7 @@ impl S3Config { pub struct PyS3Store { store: Arc>, /// A config used for pickling. This must stay in sync with the underlying store's config. - config: S3Config, + config: Arc, } impl AsRef>> for PyS3Store { @@ -128,13 +128,13 @@ impl PyS3Store { Ok(Self { store: Arc::new(MaybePrefixedStore::new(builder.build()?, prefix.clone())), - config: S3Config { + config: Arc::new(S3Config { prefix, config: combined_config, client_options, retry_config, credential_provider, - }, + }), }) } diff --git a/pyo3-object_store/src/azure/store.rs b/pyo3-object_store/src/azure/store.rs index 97420416..eb596fd8 100644 --- a/pyo3-object_store/src/azure/store.rs +++ b/pyo3-object_store/src/azure/store.rs @@ -71,7 +71,7 @@ impl AzureConfig { pub struct PyAzureStore { store: Arc>, /// A config used for pickling. This must stay in sync with the underlying store's config. - config: AzureConfig, + config: Arc, } impl AsRef>> for PyAzureStore { @@ -145,13 +145,13 @@ impl PyAzureStore { Ok(Self { store: Arc::new(MaybePrefixedStore::new(builder.build()?, prefix.clone())), - config: AzureConfig { + config: Arc::new(AzureConfig { prefix, config: combined_config, client_options, retry_config, credential_provider, - }, + }), }) } diff --git a/pyo3-object_store/src/gcp/store.rs b/pyo3-object_store/src/gcp/store.rs index 260708df..8620bf53 100644 --- a/pyo3-object_store/src/gcp/store.rs +++ b/pyo3-object_store/src/gcp/store.rs @@ -63,7 +63,7 @@ impl GCSConfig { pub struct PyGCSStore { store: Arc>, /// A config used for pickling. This must stay in sync with the underlying store's config. - config: GCSConfig, + config: Arc, } impl AsRef>> for PyGCSStore { @@ -113,13 +113,13 @@ impl PyGCSStore { } Ok(Self { store: Arc::new(MaybePrefixedStore::new(builder.build()?, prefix.clone())), - config: GCSConfig { + config: Arc::new(GCSConfig { prefix, config: combined_config, client_options, retry_config, credential_provider, - }, + }), }) } diff --git a/pyo3-object_store/src/http.rs b/pyo3-object_store/src/http.rs index 10d4a7f8..1bf61f8b 100644 --- a/pyo3-object_store/src/http.rs +++ b/pyo3-object_store/src/http.rs @@ -40,7 +40,7 @@ pub struct PyHttpStore { // own prefix. store: Arc, /// A config used for pickling. This must stay in sync with the underlying store's config. - config: HTTPConfig, + config: Arc, } impl AsRef> for PyHttpStore { @@ -74,11 +74,11 @@ impl PyHttpStore { } Ok(Self { store: Arc::new(builder.build()?), - config: HTTPConfig { + config: Arc::new(HTTPConfig { url, client_options, retry_config, - }, + }), }) } diff --git a/pyo3-object_store/src/local.rs b/pyo3-object_store/src/local.rs index fc36df91..96bdce6f 100644 --- a/pyo3-object_store/src/local.rs +++ b/pyo3-object_store/src/local.rs @@ -33,7 +33,7 @@ impl LocalConfig { #[pyclass(name = "LocalStore", frozen, subclass, from_py_object)] pub struct PyLocalStore { store: Arc, - config: LocalConfig, + config: Arc, } impl AsRef> for PyLocalStore { @@ -69,11 +69,11 @@ impl PyLocalStore { let fs = fs.with_automatic_cleanup(automatic_cleanup); Ok(Self { store: Arc::new(fs), - config: LocalConfig { + config: Arc::new(LocalConfig { prefix, automatic_cleanup, mkdir, - }, + }), }) } diff --git a/pyo3-object_store/src/prefix.rs b/pyo3-object_store/src/prefix.rs index db84807a..497083b6 100644 --- a/pyo3-object_store/src/prefix.rs +++ b/pyo3-object_store/src/prefix.rs @@ -6,6 +6,7 @@ use bytes::Bytes; use futures::{stream::BoxStream, StreamExt, TryStreamExt}; use http::Method; +use object_store::list::{PaginatedListOptions, PaginatedListResult, PaginatedListStore}; use object_store::signer::Signer; use std::borrow::Cow; use std::future::Future; @@ -62,30 +63,6 @@ impl MaybePrefixedStore { Cow::Borrowed(location) } } - - /// Strip the constant prefix from a given path - fn strip_prefix(&self, path: Path) -> Path { - if let Some(prefix) = &self.prefix { - // Note cannot use match because of borrow checker - if let Some(suffix) = path.prefix_match(prefix) { - return suffix.collect(); - } - path - } else { - path - } - } - - /// Strip the constant prefix from a given ObjectMeta - fn strip_meta(&self, meta: ObjectMeta) -> ObjectMeta { - ObjectMeta { - last_modified: meta.last_modified, - size: meta.size, - location: self.strip_prefix(meta.location), - e_tag: meta.e_tag, - version: None, - } - } } // Note: This is a relative hack to move these two functions to pure functions so they don't rely @@ -123,6 +100,22 @@ fn strip_meta(prefix: Option<&Path>, meta: ObjectMeta) -> ObjectMeta { version: None, } } + +fn strip_list_result(prefix: Option<&Path>, lst: ListResult) -> ListResult { + ListResult { + common_prefixes: lst + .common_prefixes + .into_iter() + .map(|p| strip_prefix(prefix, p)) + .collect(), + objects: lst + .objects + .into_iter() + .map(|meta| strip_meta(prefix, meta)) + .collect(), + } +} + #[async_trait::async_trait] impl ObjectStore for MaybePrefixedStore { async fn put_opts( @@ -182,18 +175,7 @@ impl ObjectStore for MaybePrefixedStore { self.inner .list_with_delimiter(Some(&prefix)) .await - .map(|lst| ListResult { - common_prefixes: lst - .common_prefixes - .into_iter() - .map(|p| self.strip_prefix(p)) - .collect(), - objects: lst - .objects - .into_iter() - .map(|meta| self.strip_meta(meta)) - .collect(), - }) + .map(|lst| strip_list_result(self.prefix.as_ref(), lst)) } async fn copy_opts(&self, from: &Path, to: &Path, options: CopyOptions) -> Result<()> { @@ -258,3 +240,41 @@ impl Signer for MaybePrefixedStore { }) } } + +fn create_paginated_list_prefix<'a>( + store_prefix: Option<&'a Path>, + list_prefix: Option<&'a str>, + delimiter: Option<&Cow<'static, str>>, +) -> Option> { + match (store_prefix, list_prefix) { + (None, None) => None, + (Some(store_prefix), None) => Some(Cow::Borrowed(store_prefix.as_ref())), + (None, Some(list_prefix)) => Some(Cow::Borrowed(list_prefix)), + (Some(store_prefix), Some(list_prefix)) => { + let delimiter = delimiter.map(|x| x.as_ref()).unwrap_or("/"); + let combined = format!("{}{delimiter}{list_prefix}", store_prefix.as_ref()); + Some(Cow::Owned(combined)) + } + } +} + +#[async_trait::async_trait] +impl PaginatedListStore for MaybePrefixedStore { + async fn list_paginated( + &self, + prefix: Option<&str>, + opts: PaginatedListOptions, + ) -> Result { + let store_prefix = self.prefix.as_ref(); + let list_prefix = + create_paginated_list_prefix(store_prefix, prefix, opts.delimiter.as_ref()); + let lst = self + .inner + .list_paginated(list_prefix.as_deref(), opts) + .await?; + Ok(PaginatedListResult { + result: strip_list_result(store_prefix, lst.result), + page_token: lst.page_token, + }) + } +} diff --git a/pyo3-object_store/src/store.rs b/pyo3-object_store/src/store.rs index 2b802f1e..810859c3 100644 --- a/pyo3-object_store/src/store.rs +++ b/pyo3-object_store/src/store.rs @@ -15,24 +15,37 @@ use crate::{PyAzureStore, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, /// This will only accept ObjectStore instances created from the same library. See /// [register_store_module][crate::register_store_module]. #[derive(Debug, Clone)] -pub struct PyObjectStore(Arc); +pub enum PyObjectStore { + /// A wrapper around a [`PyS3Store`]. + S3(PyS3Store), + /// A wrapper around a [`PyAzureStore`]. + Azure(PyAzureStore), + /// A wrapper around a [`PyGCSStore`]. + Gcs(PyGCSStore), + /// A wrapper around a [`PyHttpStore`]. + Http(PyHttpStore), + /// A wrapper around a [`PyLocalStore`]. + Local(PyLocalStore), + /// A wrapper around a [`PyMemoryStore`]. + Memory(PyMemoryStore), +} impl<'py> FromPyObject<'_, 'py> for PyObjectStore { type Error = PyErr; fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { if let Ok(store) = obj.cast::() { - Ok(Self(store.get().as_ref().clone())) + Ok(Self::S3(store.get().clone())) } else if let Ok(store) = obj.cast::() { - Ok(Self(store.get().as_ref().clone())) + Ok(Self::Azure(store.get().clone())) } else if let Ok(store) = obj.cast::() { - Ok(Self(store.get().as_ref().clone())) + Ok(Self::Gcs(store.get().clone())) } else if let Ok(store) = obj.cast::() { - Ok(Self(store.get().as_ref().clone())) + Ok(Self::Http(store.get().clone())) } else if let Ok(store) = obj.cast::() { - Ok(Self(store.get().as_ref().clone())) + Ok(Self::Local(store.get().clone())) } else if let Ok(store) = obj.cast::() { - Ok(Self(store.get().as_ref().clone())) + Ok(Self::Memory(store.get().clone())) } else { let py = obj.py(); // Check for object-store instance from other library @@ -61,27 +74,45 @@ impl<'py> FromPyObject<'_, 'py> for PyObjectStore { } } -impl AsRef> for PyObjectStore { - fn as_ref(&self) -> &Arc { - &self.0 +impl AsRef for PyObjectStore { + fn as_ref(&self) -> &dyn ObjectStore { + match self { + PyObjectStore::S3(store) => store.as_ref(), + PyObjectStore::Azure(store) => store.as_ref(), + PyObjectStore::Gcs(store) => store.as_ref(), + PyObjectStore::Http(store) => store.as_ref(), + PyObjectStore::Local(store) => store.as_ref(), + PyObjectStore::Memory(store) => store.as_ref(), + } } } impl From for Arc { fn from(value: PyObjectStore) -> Self { - value.0 + value.into_inner() } } impl PyObjectStore { /// Consume self and return the underlying [`ObjectStore`]. + /// + /// This is an alias for [Self::into_dyn]. pub fn into_inner(self) -> Arc { - self.0 + match self { + PyObjectStore::S3(store) => store.into_inner(), + PyObjectStore::Azure(store) => store.into_inner(), + PyObjectStore::Gcs(store) => store.into_inner(), + PyObjectStore::Http(store) => store.into_inner(), + PyObjectStore::Local(store) => store.into_inner(), + PyObjectStore::Memory(store) => store.into_inner(), + } } /// Consume self and return a reference-counted [`ObjectStore`]. + /// + /// This is an alias for [Self::into_inner]. pub fn into_dyn(self) -> Arc { - self.0 + self.into_inner() } } diff --git a/tests/store/test_s3.py b/tests/store/test_s3.py index dd9061f2..b402b4e9 100644 --- a/tests/store/test_s3.py +++ b/tests/store/test_s3.py @@ -17,6 +17,50 @@ async def test_list_async(minio_store: S3Store): assert any("afile" in x["path"] for x in list_result) +def test_list_paginates_across_pages(minio_store: S3Store): + """`create_paginated_stream` should auto-paginate under the hood. + + We upload more objects than fit in a single page and use a small ``chunk_size`` + (which is also used as the underlying ``max_keys`` page size) so the store must + follow several continuation tokens. The stream should stitch every page together + rather than stopping after the first page. + """ + n = 120 + for i in range(n): + minio_store.put(f"obj/{i:04d}.txt", b"x") + + result = minio_store.list("obj/", chunk_size=10).collect() + assert len(result) == n + paths = sorted(item["path"] for item in result) + assert paths[0] == "obj/0000.txt" + assert paths[-1] == f"obj/{n - 1:04d}.txt" + + +def test_list_substring_prefix(minio_store: S3Store): + """Prefixes are treated as raw string prefixes, not whole path segments. + + This is the core behavior the paginated-list rework enables: `2025/log` matches + `2025/log_a.txt` even though `log` is only part of the final path segment. + """ + minio_store.put("2025/log_a.txt", b"x") + minio_store.put("2025/log_b.txt", b"x") + minio_store.put("2025/data_c.txt", b"x") + + result = minio_store.list("2025/log").collect() + paths = sorted(item["path"] for item in result) + assert paths == ["2025/log_a.txt", "2025/log_b.txt"] + + +def test_list_offset(minio_store: S3Store): + """`offset` is forwarded through pagination and is exclusive.""" + for i in range(10): + minio_store.put(f"item/{i:02d}.txt", b"x") + + result = minio_store.list("item/", offset="item/05.txt").collect() + paths = sorted(item["path"] for item in result) + assert paths == [f"item/{i:02d}.txt" for i in range(6, 10)] + + @pytest.mark.asyncio async def test_get_async(minio_store: S3Store): await minio_store.put_async("afile", b"hello world") diff --git a/tests/test_list.py b/tests/test_list.py index dcec45a2..cee0365e 100644 --- a/tests/test_list.py +++ b/tests/test_list.py @@ -1,7 +1,11 @@ +import tempfile +from pathlib import Path + +import pyarrow as pa import pytest from arro3.core import RecordBatch, Table -from obstore.store import MemoryStore +from obstore.store import LocalStore, MemoryStore, ObjectStore, S3Store def test_list(): @@ -159,3 +163,104 @@ async def test_list_with_delimiter_async(): assert objects.num_rows == 2 assert objects["path"][0].as_py() == "a/file1.txt" assert objects["path"][1].as_py() == "a/file2.txt" + + +def test_list_substring_filtering_emulated(): + 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/test/deep/file.txt", b"bar") + store.put("data/another/2.csv", b"baz") + store.put("data/test_data.json", b"qux") + store.put("logs/test_log.txt", b"log") + + # The prefix is a raw string prefix (not a whole path segment), and matching is + # recursive: every key starting with "data/tes" is returned, including nested ones. + result = store.list("data/tes").collect() + paths = {item["path"] for item in result} + assert paths == { + "data/test/file.txt", + "data/test/deep/file.txt", + "data/test_data.json", + } + + # Same filter, returned as arrow. + batch = store.list("data/tes", return_arrow=True).collect() + assert isinstance(batch, RecordBatch) + assert batch.num_rows == 3 + + +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, including a nested directory under a matching + # prefix to verify recursive listing. + data_dir = temp_dir_path / "data" + (data_dir / "test_sub").mkdir(parents=True, exist_ok=True) + + (data_dir / "file1.txt").write_text("foo") + (data_dir / "test_file.txt").write_text("bar") + (data_dir / "another.csv").write_text("baz") + (data_dir / "test_data.json").write_text("qux") + (data_dir / "test_sub" / "deep.txt").write_text("deep") + + # Matching is recursive and on a raw string prefix, so the nested + # "data/test_sub/deep.txt" is included. + result = store.list("data/test").collect() + paths = {item["path"] for item in result} + assert paths == { + "data/test_file.txt", + "data/test_data.json", + "data/test_sub/deep.txt", + } + + +def _assert_substring_prefix_listing(store: ObjectStore): + """Test `list` substring-prefix behavior across native and emulated backends.""" + for path in [ + "data/file1.txt", + "data/test/file.txt", + "data/test/deep/file.txt", + "data/another/2.csv", + "data/test_data.json", + "logs/test_log.txt", + ]: + store.put(path, b"x") + + paths = {item["path"] for item in store.list("data/tes").collect()} + assert paths == { + "data/test/file.txt", + "data/test/deep/file.txt", + "data/test_data.json", + } + + +def test_list_substring_prefix_emulated(): + _assert_substring_prefix_listing(MemoryStore()) + + +def test_list_substring_prefix_native(minio_store: S3Store): + # `minio_store` is an S3Store backed by a real (paginating) MinIO container. + _assert_substring_prefix_listing(minio_store) + + +def test_list_as_arrow_to_pyarrow(): + store = MemoryStore() + + for i in range(100): + store.put(f"file{i}.txt", b"foo") + + stream = store.list(return_arrow=True, chunk_size=10) + + # The RecordBatch yielded by the stream implements the Arrow PyCapsule interface, + # so external Arrow libraries can consume it zero-copy. + batch = next(stream) + assert isinstance(batch, RecordBatch) + + pa_batch = pa.record_batch(batch) + assert pa_batch.num_rows == 10 + assert "path" in pa_batch.column_names