From a123ed1e72ff05fed0f4d717a759a6d34e7bb446 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Fri, 29 Aug 2025 17:19:56 -0400 Subject: [PATCH 01/15] WIP: Paginated list support --- pyo3-object_store/src/aws/store.rs | 12 ++ pyo3-object_store/src/azure/store.rs | 12 ++ pyo3-object_store/src/gcp/store.rs | 12 ++ pyo3-object_store/src/local.rs | 66 +++++++++- pyo3-object_store/src/prefix.rs | 172 ++++++++++++++++++--------- 5 files changed, 218 insertions(+), 56 deletions(-) diff --git a/pyo3-object_store/src/aws/store.rs b/pyo3-object_store/src/aws/store.rs index e87c6fc8..75ea79aa 100644 --- a/pyo3-object_store/src/aws/store.rs +++ b/pyo3-object_store/src/aws/store.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use itertools::Itertools; use object_store::aws::{AmazonS3, AmazonS3Builder, AmazonS3ConfigKey}; +use object_store::list::{PaginatedListOptions, PaginatedListResult, PaginatedListStore}; use object_store::ObjectStoreScheme; use pyo3::prelude::*; use pyo3::pybacked::PyBackedStr; @@ -420,3 +421,14 @@ fn parse_url( Ok(config) } + +#[async_trait::async_trait] +impl PaginatedListStore for PyS3Store { + async fn list_paginated( + &self, + prefix: Option<&str>, + opts: PaginatedListOptions, + ) -> object_store::Result { + self.store.list_paginated(prefix, opts).await + } +} diff --git a/pyo3-object_store/src/azure/store.rs b/pyo3-object_store/src/azure/store.rs index f643e92b..1d94aa67 100644 --- a/pyo3-object_store/src/azure/store.rs +++ b/pyo3-object_store/src/azure/store.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use std::sync::Arc; use object_store::azure::{AzureConfigKey, MicrosoftAzure, MicrosoftAzureBuilder}; +use object_store::list::{PaginatedListOptions, PaginatedListResult, PaginatedListStore}; use object_store::ObjectStoreScheme; use pyo3::prelude::*; use pyo3::pybacked::PyBackedStr; @@ -470,3 +471,14 @@ fn parse_url(config: Option, parsed: &Url) -> object_store::Resul Ok(config) } + +#[async_trait::async_trait] +impl PaginatedListStore for PyAzureStore { + async fn list_paginated( + &self, + prefix: Option<&str>, + opts: PaginatedListOptions, + ) -> object_store::Result { + self.store.list_paginated(prefix, opts).await + } +} diff --git a/pyo3-object_store/src/gcp/store.rs b/pyo3-object_store/src/gcp/store.rs index 4c28193b..29e7758e 100644 --- a/pyo3-object_store/src/gcp/store.rs +++ b/pyo3-object_store/src/gcp/store.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use std::sync::Arc; use object_store::gcp::{GoogleCloudStorage, GoogleCloudStorageBuilder, GoogleConfigKey}; +use object_store::list::{PaginatedListOptions, PaginatedListResult, PaginatedListStore}; use object_store::ObjectStoreScheme; use pyo3::prelude::*; use pyo3::pybacked::PyBackedStr; @@ -374,3 +375,14 @@ fn parse_url(config: Option, parsed: &Url) -> object_store::Resu Ok(config) } + +#[async_trait::async_trait] +impl PaginatedListStore for PyGCSStore { + async fn list_paginated( + &self, + prefix: Option<&str>, + opts: PaginatedListOptions, + ) -> object_store::Result { + self.store.list_paginated(prefix, opts).await + } +} diff --git a/pyo3-object_store/src/local.rs b/pyo3-object_store/src/local.rs index 9b416d5c..74abb1e8 100644 --- a/pyo3-object_store/src/local.rs +++ b/pyo3-object_store/src/local.rs @@ -1,8 +1,9 @@ use std::fs::create_dir_all; use std::sync::Arc; +use object_store::list::{PaginatedListOptions, PaginatedListResult, PaginatedListStore}; use object_store::local::LocalFileSystem; -use object_store::ObjectStoreScheme; +use object_store::{ListResult, ObjectStore, ObjectStoreScheme}; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::types::{PyDict, PyTuple, PyType}; @@ -143,3 +144,66 @@ impl PyLocalStore { } } } + +/// A custom implementation of PaginatedListStore for LocalFileSystem +/// +/// 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_trait::async_trait] +impl PaginatedListStore for PyLocalStore { + async fn list_paginated( + &self, + prefix: Option<&str>, + _opts: PaginatedListOptions, + ) -> object_store::Result { + // Split a path like "some/prefix/abc" into (Some(Path("some/prefix")), Some("abc")) + // This allows us to do a substring prefix match after the / delimiter + let (list_path, list_prefix_match): (Option, Option) = + if let Some(list_prefix) = prefix { + if let Some((list_path, list_prefix_match)) = list_prefix.rsplit_once('/') { + // There's a / in the prefix, so we assume the part before the last / is a + // path, and the end is a substring match + ( + Some(object_store::path::Path::parse(list_path)?), + Some(list_prefix_match.to_string()), + ) + } else { + // No / in prefix, so we assume it's a substring + (None, Some(list_prefix.to_string())) + } + } else { + (None, None) + }; + + let list_result = self.store.list_with_delimiter(list_path.as_ref()).await?; + + // Filter list result to include only results with the given prefix after the / delimiter + let filtered_list_result = if let Some(list_prefix_match) = list_prefix_match { + let filtered_common_prefixes = list_result + .common_prefixes + .into_iter() + .filter(|p| p.as_ref().starts_with(&list_prefix_match)) + .collect(); + let filtered_objects = list_result + .objects + .into_iter() + .filter(|obj| obj.location.as_ref().starts_with(&list_prefix_match)) + .collect(); + ListResult { + common_prefixes: filtered_common_prefixes, + objects: filtered_objects, + } + } else { + list_result + }; + + Ok(PaginatedListResult { + result: filtered_list_result, + // Local FS does not support pagination + page_token: None, + }) + } +} diff --git a/pyo3-object_store/src/prefix.rs b/pyo3-object_store/src/prefix.rs index 3c376afa..60e006f5 100644 --- a/pyo3-object_store/src/prefix.rs +++ b/pyo3-object_store/src/prefix.rs @@ -5,6 +5,10 @@ use bytes::Bytes; use futures::{stream::BoxStream, StreamExt, TryStreamExt}; +use object_store::aws::AmazonS3; +use object_store::azure::MicrosoftAzure; +use object_store::gcp::GoogleCloudStorage; +use object_store::list::{PaginatedListOptions, PaginatedListResult, PaginatedListStore}; use std::borrow::Cow; use std::ops::Range; use std::sync::OnceLock; @@ -58,58 +62,47 @@ 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 -// on the `self` lifetime. Expected to be cleaned up before merge. -// /// Strip the constant prefix from a given path -fn strip_prefix(prefix: &Path, path: Path) -> Path { - // Note cannot use match because of borrow checker - if let Some(suffix) = path.prefix_match(prefix) { - return suffix.collect(); +fn strip_prefix(prefix: Option<&Path>, path: Path) -> Path { + if let Some(prefix) = prefix { + // Note cannot use match because of borrow checker + if let Some(suffix) = path.prefix_match(prefix) { + return suffix.collect(); + } + path + } else { + path } - path } /// Strip the constant prefix from a given ObjectMeta fn strip_meta(prefix: Option<&Path>, meta: ObjectMeta) -> ObjectMeta { - if let Some(prefix) = prefix { - ObjectMeta { - last_modified: meta.last_modified, - size: meta.size, - location: strip_prefix(prefix, meta.location), - e_tag: meta.e_tag, - version: None, - } - } else { - meta + ObjectMeta { + last_modified: meta.last_modified, + size: meta.size, + location: strip_prefix(prefix, meta.location), + e_tag: meta.e_tag, + 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(&self, location: &Path, payload: PutPayload) -> Result { @@ -166,7 +159,7 @@ impl ObjectStore for MaybePrefixedStore { async fn head(&self, location: &Path) -> Result { let full_path = self.full_path(location); let meta = self.inner.head(&full_path).await?; - Ok(self.strip_meta(meta)) + Ok(strip_meta(self.prefix.as_ref(), meta)) } async fn delete(&self, location: &Path) -> Result<()> { @@ -200,18 +193,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(&self, from: &Path, to: &Path) -> Result<()> { @@ -238,3 +220,83 @@ impl ObjectStore for MaybePrefixedStore { self.inner.rename_if_not_exists(&full_from, &full_to).await } } + +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, + }) + } +} + +#[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, + }) + } +} + +#[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, + }) + } +} From 4ea16c040651b2e66de2822eaf9833e08d15cce8 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Wed, 3 Sep 2025 16:53:13 -0400 Subject: [PATCH 02/15] Create MaybePaginatedListStore --- obstore/src/list.rs | 58 ++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 55 insertions(+), 3 deletions(-) diff --git a/obstore/src/list.rs b/obstore/src/list.rs index 1a4e9cd6..6832872f 100644 --- a/obstore/src/list.rs +++ b/obstore/src/list.rs @@ -8,17 +8,69 @@ 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::path::Path; use object_store::{ListResult, ObjectMeta, ObjectStore}; -use pyo3::exceptions::{PyImportError, PyStopAsyncIteration, PyStopIteration}; +use pyo3::exceptions::{PyImportError, PyStopAsyncIteration, PyStopIteration, PyValueError}; use pyo3::prelude::*; +use pyo3::pybacked::PyBackedStr; use pyo3::types::PyDict; -use pyo3::{intern, IntoPyObjectExt}; +use pyo3::{intern, IntoPyObjectExt, PyTypeInfo}; use pyo3_arrow::{PyRecordBatch, PyTable}; use pyo3_async_runtimes::tokio::get_runtime; -use pyo3_object_store::{PyObjectStore, PyObjectStoreError, PyObjectStoreResult}; +use pyo3_object_store::{ + PyAzureStore, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, PyObjectStore, + PyObjectStoreError, PyObjectStoreResult, PyS3Store, +}; use tokio::sync::Mutex; +enum MaybePaginatedListStore { + SupportsPagination(Arc), + NoPagination(Arc), +} + +impl<'py> FromPyObject<'py> for MaybePaginatedListStore { + fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + if let Ok(store) = ob.downcast::() { + Ok(Self::SupportsPagination(store.get().as_ref().clone())) + } else if let Ok(store) = ob.downcast::() { + Ok(Self::SupportsPagination(store.get().as_ref().clone())) + } else if let Ok(store) = ob.downcast::() { + Ok(Self::SupportsPagination(store.get().as_ref().clone())) + } else if let Ok(store) = ob.downcast::() { + Ok(Self::NoPagination(store.get().as_ref().clone())) + } else if let Ok(store) = ob.downcast::() { + Ok(Self::NoPagination(store.get().as_ref().clone())) + } else if let Ok(store) = ob.downcast::() { + Ok(Self::NoPagination(store.get().as_ref().clone())) + } else { + let py = ob.py(); + // Check for object-store instance from other library + let cls_name = ob + .getattr(intern!(py, "__class__"))? + .getattr(intern!(py, "__name__"))? + .extract::()?; + if [ + PyAzureStore::NAME, + PyGCSStore::NAME, + PyHttpStore::NAME, + PyLocalStore::NAME, + PyMemoryStore::NAME, + PyS3Store::NAME, + ] + .contains(&cls_name.as_ref()) + { + return Err(PyValueError::new_err("You must use an object store instance exported from **the same library** as this function. They cannot be used across libraries.\nThis is because object store instances are compiled with a specific version of Rust and Python." )); + } + + Err(PyValueError::new_err(format!( + "Expected an object store instance, got {}", + ob.repr()? + ))) + } + } +} + pub(crate) struct PyObjectMeta(ObjectMeta); impl PyObjectMeta { From 958fd4d68ffd896d05c53828ba92112426f148f8 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Wed, 3 Sep 2025 23:46:28 -0400 Subject: [PATCH 03/15] Claude-created paginated listing --- obstore/src/list.rs | 118 ++++++++++++++++++++++++++++++++++++++++---- tests/test_list.py | 63 ++++++++++++++++++++++- 2 files changed, 171 insertions(+), 10 deletions(-) diff --git a/obstore/src/list.rs b/obstore/src/list.rs index 6832872f..46d72fa6 100644 --- a/obstore/src/list.rs +++ b/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), NoPagination(Arc), } @@ -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, offset: Option, 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, + 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( + (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, + prefix: Option, + offset: Option, +) -> BoxStream<'static, object_store::Result> { + // 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 +} diff --git a/tests/test_list.py b/tests/test_list.py index a151d760..03f12128 100644 --- a/tests/test_list.py +++ b/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() From a86bfc83dd833fd99db2131ee71451a077d5305c Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Tue, 2 Jun 2026 14:29:59 -0400 Subject: [PATCH 04/15] use generic for cleaner code --- pyo3-object_store/src/prefix.rs | 47 +-------------------------------- 1 file changed, 1 insertion(+), 46 deletions(-) diff --git a/pyo3-object_store/src/prefix.rs b/pyo3-object_store/src/prefix.rs index 7fcc870c..497083b6 100644 --- a/pyo3-object_store/src/prefix.rs +++ b/pyo3-object_store/src/prefix.rs @@ -6,9 +6,6 @@ use bytes::Bytes; use futures::{stream::BoxStream, StreamExt, TryStreamExt}; use http::Method; -use object_store::aws::AmazonS3; -use object_store::azure::MicrosoftAzure; -use object_store::gcp::GoogleCloudStorage; use object_store::list::{PaginatedListOptions, PaginatedListResult, PaginatedListStore}; use object_store::signer::Signer; use std::borrow::Cow; @@ -262,49 +259,7 @@ fn create_paginated_list_prefix<'a>( } #[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, - }) - } -} - -#[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, - }) - } -} - -#[async_trait::async_trait] -impl PaginatedListStore for MaybePrefixedStore { +impl PaginatedListStore for MaybePrefixedStore { async fn list_paginated( &self, prefix: Option<&str>, From 3909ea6fc05a123507be58ff144450cd86247ac5 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Tue, 2 Jun 2026 15:10:43 -0400 Subject: [PATCH 05/15] refactor PyObjectStore into enum of typed stores --- obstore/src/get.rs | 4 +-- pyo3-object_store/src/store.rs | 57 ++++++++++++++++++++++++++-------- 2 files changed, 46 insertions(+), 15 deletions(-) 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/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() } } From 06e1f1b4ee69147cbe3de8a9c994b3eec938be19 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Tue, 2 Jun 2026 15:10:52 -0400 Subject: [PATCH 06/15] store config under Arc --- pyo3-object_store/src/aws/store.rs | 6 +++--- pyo3-object_store/src/azure/store.rs | 6 +++--- pyo3-object_store/src/gcp/store.rs | 6 +++--- pyo3-object_store/src/http.rs | 6 +++--- pyo3-object_store/src/local.rs | 6 +++--- 5 files changed, 15 insertions(+), 15 deletions(-) diff --git a/pyo3-object_store/src/aws/store.rs b/pyo3-object_store/src/aws/store.rs index 8ed0cd83..c266192f 100644 --- a/pyo3-object_store/src/aws/store.rs +++ b/pyo3-object_store/src/aws/store.rs @@ -66,7 +66,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 { @@ -129,13 +129,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 1b93cfa6..cef31f2f 100644 --- a/pyo3-object_store/src/azure/store.rs +++ b/pyo3-object_store/src/azure/store.rs @@ -72,7 +72,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 { @@ -146,13 +146,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 5e5178df..3aab326f 100644 --- a/pyo3-object_store/src/gcp/store.rs +++ b/pyo3-object_store/src/gcp/store.rs @@ -64,7 +64,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 { @@ -114,13 +114,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 75801561..ec29914b 100644 --- a/pyo3-object_store/src/local.rs +++ b/pyo3-object_store/src/local.rs @@ -34,7 +34,7 @@ impl LocalConfig { #[pyclass(name = "LocalStore", frozen, subclass, from_py_object)] pub struct PyLocalStore { store: Arc, - config: LocalConfig, + config: Arc, } impl AsRef> for PyLocalStore { @@ -70,11 +70,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, - }, + }), }) } From a7465b7eef8aaf1b56333456f7edc44b5e4da0b0 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Tue, 2 Jun 2026 15:32:37 -0400 Subject: [PATCH 07/15] clean up list --- Cargo.lock | 1 + Cargo.toml | 1 + obstore/Cargo.toml | 1 + obstore/src/list.rs | 169 ++++++++++++++++----------- pyo3-object_store/src/aws/store.rs | 12 -- pyo3-object_store/src/azure/store.rs | 12 -- pyo3-object_store/src/gcp/store.rs | 12 -- pyo3-object_store/src/local.rs | 66 +---------- 8 files changed, 107 insertions(+), 167 deletions(-) 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/list.rs b/obstore/src/list.rs index 3e12c23f..3fb23f3c 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,71 +9,18 @@ use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; use futures::stream::{BoxStream, Fuse}; use futures::StreamExt; use indexmap::IndexMap; -use object_store::list::{PaginatedListOptions, PaginatedListStore}; +use object_store::list::{PaginatedListOptions, PaginatedListResult, PaginatedListStore}; use object_store::{ListResult, ObjectMeta, ObjectStore}; -use pyo3::exceptions::{PyImportError, PyStopAsyncIteration, PyStopIteration, PyValueError}; +use pyo3::exceptions::{PyImportError, PyStopAsyncIteration, PyStopIteration}; use pyo3::prelude::*; -use pyo3::pybacked::PyBackedStr; use pyo3::types::PyDict; -use pyo3::{intern, IntoPyObjectExt, PyTypeInfo}; +use pyo3::{intern, IntoPyObjectExt}; use pyo3_arrow::export::{Arro3RecordBatch, Arro3Table}; use pyo3_arrow::PyTable; use pyo3_async_runtimes::tokio::get_runtime; -use pyo3_object_store::{ - PyAzureStore, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, PyObjectStore, - PyObjectStoreError, PyObjectStoreResult, PyPath, PyS3Store, -}; +use pyo3_object_store::{PyObjectStore, PyObjectStoreError, PyObjectStoreResult, PyPath}; use tokio::sync::Mutex; -pub(crate) enum MaybePaginatedListStore { - SupportsPagination(Arc), - NoPagination(Arc), -} - -impl<'py> FromPyObject<'_, 'py> for MaybePaginatedListStore { - type Error = PyErr; - - fn extract(ob: Borrowed<'_, 'py, PyAny>) -> PyResult { - if let Ok(store) = ob.cast::() { - Ok(Self::SupportsPagination(store.get().as_ref().clone())) - } else if let Ok(store) = ob.cast::() { - Ok(Self::SupportsPagination(store.get().as_ref().clone())) - } else if let Ok(store) = ob.cast::() { - Ok(Self::SupportsPagination(store.get().as_ref().clone())) - } else if let Ok(store) = ob.cast::() { - Ok(Self::NoPagination(store.get().as_ref().clone())) - } else if let Ok(store) = ob.cast::() { - Ok(Self::NoPagination(store.get().as_ref().clone())) - } else if let Ok(store) = ob.cast::() { - Ok(Self::NoPagination(store.get().as_ref().clone())) - } else { - let py = ob.py(); - // Check for object-store instance from other library - let cls_name = ob - .getattr(intern!(py, "__class__"))? - .getattr(intern!(py, "__name__"))? - .extract::()?; - if [ - PyAzureStore::type_object(py).name()?.to_str()?, - PyGCSStore::type_object(py).name()?.to_str()?, - PyHttpStore::type_object(py).name()?.to_str()?, - PyLocalStore::type_object(py).name()?.to_str()?, - PyMemoryStore::type_object(py).name()?.to_str()?, - PyS3Store::type_object(py).name()?.to_str()?, - ] - .contains(&cls_name.as_str()) - { - return Err(PyValueError::new_err("You must use an object store instance exported from **the same library** as this function. They cannot be used across libraries.\nThis is because object store instances are compiled with a specific version of Rust and Python." )); - } - - Err(PyValueError::new_err(format!( - "Expected an object store instance, got {}", - ob.repr()? - ))) - } - } -} - pub(crate) struct PyObjectMeta(ObjectMeta); impl PyObjectMeta { @@ -401,11 +349,105 @@ 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(s3) => Self::Native(s3.into_inner()), + PyObjectStore::Azure(azure) => Self::Native(azure.into_inner()), + PyObjectStore::Gcs(gcs) => Self::Native(gcs.into_inner()), + PyObjectStore::Http(http) => Self::Emulated(http.into_inner()), + PyObjectStore::Local(local) => Self::Emulated(local.into_inner()), + PyObjectStore::Memory(memory) => Self::Emulated(memory.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 { + // Split a path like "some/prefix/abc" into (Some(Path("some/prefix")), Some("abc")) + // This allows us to do a substring prefix match after the / delimiter + let (list_path, list_prefix_match): (Option, Option) = + if let Some(list_prefix) = prefix { + if let Some((list_path, list_prefix_match)) = list_prefix.rsplit_once('/') { + // There's a / in the prefix, so we assume the part before the last / is a + // path, and the end is a substring match + ( + Some(object_store::path::Path::parse(list_path)?), + Some(list_prefix_match.to_string()), + ) + } else { + // No / in prefix, so we assume it's a substring + (None, Some(list_prefix.to_string())) + } + } else { + (None, None) + }; + + let list_result = store.list_with_delimiter(list_path.as_ref()).await?; + + // Filter list result to include only results with the given prefix after the / delimiter + let filtered_list_result = if let Some(list_prefix_match) = list_prefix_match { + let filtered_common_prefixes = list_result + .common_prefixes + .into_iter() + .filter(|p| p.as_ref().starts_with(&list_prefix_match)) + .collect(); + let filtered_objects = list_result + .objects + .into_iter() + .filter(|obj| obj.location.as_ref().starts_with(&list_prefix_match)) + .collect(); + ListResult { + common_prefixes: filtered_common_prefixes, + objects: filtered_objects, + } + } else { + list_result + }; + + Ok(PaginatedListResult { + result: filtered_list_result, + // emulated stores do not support pagination + 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: MaybePaginatedListStore, + store: PyObjectStore, prefix: Option, offset: Option, chunk_size: usize, @@ -423,14 +465,8 @@ pub(crate) fn list( .map_err(|err| PyImportError::new_err(format!("{msg}\n\n{err}")))?; } - 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) - } - }; + 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)) } @@ -524,6 +560,7 @@ fn create_paginated_stream( Box::pin(stream) } +// I think this isn't used anymore, and is subsumed into emulate_paginated_list? fn create_filtered_stream( store: Arc, prefix: Option, diff --git a/pyo3-object_store/src/aws/store.rs b/pyo3-object_store/src/aws/store.rs index c266192f..56826bb7 100644 --- a/pyo3-object_store/src/aws/store.rs +++ b/pyo3-object_store/src/aws/store.rs @@ -3,7 +3,6 @@ use std::sync::Arc; use itertools::Itertools; use object_store::aws::{AmazonS3, AmazonS3Builder, AmazonS3ConfigKey}; -use object_store::list::{PaginatedListOptions, PaginatedListResult, PaginatedListStore}; use object_store::ObjectStoreScheme; use pyo3::prelude::*; use pyo3::pybacked::PyBackedStr; @@ -429,14 +428,3 @@ fn parse_url( Ok(config) } - -#[async_trait::async_trait] -impl PaginatedListStore for PyS3Store { - async fn list_paginated( - &self, - prefix: Option<&str>, - opts: PaginatedListOptions, - ) -> object_store::Result { - self.store.list_paginated(prefix, opts).await - } -} diff --git a/pyo3-object_store/src/azure/store.rs b/pyo3-object_store/src/azure/store.rs index cef31f2f..eb596fd8 100644 --- a/pyo3-object_store/src/azure/store.rs +++ b/pyo3-object_store/src/azure/store.rs @@ -2,7 +2,6 @@ use std::collections::HashMap; use std::sync::Arc; use object_store::azure::{AzureConfigKey, MicrosoftAzure, MicrosoftAzureBuilder}; -use object_store::list::{PaginatedListOptions, PaginatedListResult, PaginatedListStore}; use object_store::ObjectStoreScheme; use pyo3::prelude::*; use pyo3::pybacked::PyBackedStr; @@ -488,14 +487,3 @@ fn parse_url(config: Option, parsed: &Url) -> object_store::Resul Ok(config) } - -#[async_trait::async_trait] -impl PaginatedListStore for PyAzureStore { - async fn list_paginated( - &self, - prefix: Option<&str>, - opts: PaginatedListOptions, - ) -> object_store::Result { - self.store.list_paginated(prefix, opts).await - } -} diff --git a/pyo3-object_store/src/gcp/store.rs b/pyo3-object_store/src/gcp/store.rs index 3aab326f..8620bf53 100644 --- a/pyo3-object_store/src/gcp/store.rs +++ b/pyo3-object_store/src/gcp/store.rs @@ -2,7 +2,6 @@ use std::collections::HashMap; use std::sync::Arc; use object_store::gcp::{GoogleCloudStorage, GoogleCloudStorageBuilder, GoogleConfigKey}; -use object_store::list::{PaginatedListOptions, PaginatedListResult, PaginatedListStore}; use object_store::ObjectStoreScheme; use pyo3::prelude::*; use pyo3::pybacked::PyBackedStr; @@ -379,14 +378,3 @@ fn parse_url(config: Option, parsed: &Url) -> object_store::Resu Ok(config) } - -#[async_trait::async_trait] -impl PaginatedListStore for PyGCSStore { - async fn list_paginated( - &self, - prefix: Option<&str>, - opts: PaginatedListOptions, - ) -> object_store::Result { - self.store.list_paginated(prefix, opts).await - } -} diff --git a/pyo3-object_store/src/local.rs b/pyo3-object_store/src/local.rs index ec29914b..96bdce6f 100644 --- a/pyo3-object_store/src/local.rs +++ b/pyo3-object_store/src/local.rs @@ -1,9 +1,8 @@ use std::fs::create_dir_all; use std::sync::Arc; -use object_store::list::{PaginatedListOptions, PaginatedListResult, PaginatedListStore}; use object_store::local::LocalFileSystem; -use object_store::{ListResult, ObjectStore, ObjectStoreScheme}; +use object_store::ObjectStoreScheme; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::types::{PyDict, PyTuple, PyType}; @@ -142,66 +141,3 @@ impl PyLocalStore { } } } - -/// A custom implementation of PaginatedListStore for LocalFileSystem -/// -/// 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_trait::async_trait] -impl PaginatedListStore for PyLocalStore { - async fn list_paginated( - &self, - prefix: Option<&str>, - _opts: PaginatedListOptions, - ) -> object_store::Result { - // Split a path like "some/prefix/abc" into (Some(Path("some/prefix")), Some("abc")) - // This allows us to do a substring prefix match after the / delimiter - let (list_path, list_prefix_match): (Option, Option) = - if let Some(list_prefix) = prefix { - if let Some((list_path, list_prefix_match)) = list_prefix.rsplit_once('/') { - // There's a / in the prefix, so we assume the part before the last / is a - // path, and the end is a substring match - ( - Some(object_store::path::Path::parse(list_path)?), - Some(list_prefix_match.to_string()), - ) - } else { - // No / in prefix, so we assume it's a substring - (None, Some(list_prefix.to_string())) - } - } else { - (None, None) - }; - - let list_result = self.store.list_with_delimiter(list_path.as_ref()).await?; - - // Filter list result to include only results with the given prefix after the / delimiter - let filtered_list_result = if let Some(list_prefix_match) = list_prefix_match { - let filtered_common_prefixes = list_result - .common_prefixes - .into_iter() - .filter(|p| p.as_ref().starts_with(&list_prefix_match)) - .collect(); - let filtered_objects = list_result - .objects - .into_iter() - .filter(|obj| obj.location.as_ref().starts_with(&list_prefix_match)) - .collect(); - ListResult { - common_prefixes: filtered_common_prefixes, - objects: filtered_objects, - } - } else { - list_result - }; - - Ok(PaginatedListResult { - result: filtered_list_result, - // Local FS does not support pagination - page_token: None, - }) - } -} From 2dfd050119e20ca1f34341d2b39be80f9b006619 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Tue, 2 Jun 2026 15:39:44 -0400 Subject: [PATCH 08/15] reword --- obstore/src/list.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/obstore/src/list.rs b/obstore/src/list.rs index 3fb23f3c..d5e0ae2a 100644 --- a/obstore/src/list.rs +++ b/obstore/src/list.rs @@ -359,12 +359,12 @@ enum MaybePaginatedStore { impl From for MaybePaginatedStore { fn from(store: PyObjectStore) -> Self { match store { - PyObjectStore::S3(s3) => Self::Native(s3.into_inner()), - PyObjectStore::Azure(azure) => Self::Native(azure.into_inner()), - PyObjectStore::Gcs(gcs) => Self::Native(gcs.into_inner()), - PyObjectStore::Http(http) => Self::Emulated(http.into_inner()), - PyObjectStore::Local(local) => Self::Emulated(local.into_inner()), - PyObjectStore::Memory(memory) => Self::Emulated(memory.into_inner()), + 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()), } } } From afaa4147fcd977b2ea68fe2502ccaebc5af95dbd Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Tue, 2 Jun 2026 15:40:14 -0400 Subject: [PATCH 09/15] remove polars from test_list --- tests/test_list.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/test_list.py b/tests/test_list.py index 64563b3f..b0485ee8 100644 --- a/tests/test_list.py +++ b/tests/test_list.py @@ -1,7 +1,6 @@ import tempfile from pathlib import Path -import polars as pl import pyarrow as pa import pytest from arro3.core import RecordBatch, Table @@ -224,12 +223,19 @@ def test_list_substring_filtering_local_store(): assert "data/another.csv" not in paths -def test_list_as_arrow_to_polars(): +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) - _pl_df = pl.DataFrame(next(stream)) - _df = pa.record_batch(next(stream)) + + # 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 From 90710f06b8d9d6d801bda7b99112bfccdf2ab966 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Tue, 2 Jun 2026 16:10:58 -0400 Subject: [PATCH 10/15] remove unused code --- obstore/src/list.rs | 56 --------------------------------------------- 1 file changed, 56 deletions(-) diff --git a/obstore/src/list.rs b/obstore/src/list.rs index d5e0ae2a..0dc3b3c4 100644 --- a/obstore/src/list.rs +++ b/obstore/src/list.rs @@ -559,59 +559,3 @@ fn create_paginated_stream( Box::pin(stream) } - -// I think this isn't used anymore, and is subsumed into emulate_paginated_list? -fn create_filtered_stream( - store: Arc, - prefix: Option, - offset: Option, -) -> BoxStream<'static, object_store::Result> { - // 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('/').next_back() { - if filename.contains(&substring) { - Some(Ok(meta)) - } else { - None - } - } else { - Some(Ok(meta)) - } - } - Err(e) => Some(Err(e)), - } - } - })) - } else { - base_stream - }; - - filtered_stream -} From e80e428c832bcbaceed29977affa61cfd690ed7a Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Tue, 2 Jun 2026 16:14:56 -0400 Subject: [PATCH 11/15] propagate paginated list errors --- obstore/src/list.rs | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/obstore/src/list.rs b/obstore/src/list.rs index 0dc3b3c4..dae977c8 100644 --- a/obstore/src/list.rs +++ b/obstore/src/list.rs @@ -542,20 +542,19 @@ fn create_paginated_stream( Ok(result) => { let next_has_more = result.page_token.is_some(); let next_page_token = result.page_token; - let objects = result.result.objects; + let objects: Vec> = + result.result.objects.into_iter().map(Ok).collect(); 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))) - } + // 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)], (store, prefix, offset, None, false))), } }, ) - .flat_map(|objects| futures::stream::iter(objects.into_iter().map(Ok))); + .flat_map(futures::stream::iter); Box::pin(stream) } From b4bd7b0470018386f75e915b2d48daa18159f2cd Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Tue, 2 Jun 2026 16:25:53 -0400 Subject: [PATCH 12/15] refactor stream state for readability --- obstore/src/list.rs | 100 +++++++++++++++++++++++++++++++------------- 1 file changed, 72 insertions(+), 28 deletions(-) diff --git a/obstore/src/list.rs b/obstore/src/list.rs index dae977c8..567782ee 100644 --- a/obstore/src/list.rs +++ b/obstore/src/list.rs @@ -516,6 +516,71 @@ async fn list_with_delimiter_materialize( Ok(PyListResult::new(list_result, return_arrow)) } +/// Internal stream state +#[derive(Clone)] +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 = PaginatedListOptions { + offset: offset.clone(), + delimiter: None, + max_keys: Some(chunk_size), + page_token: page_token.clone(), + ..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, @@ -524,35 +589,14 @@ fn create_paginated_stream( ) -> BoxStream<'static, object_store::Result> { // 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: Vec> = - result.result.objects.into_iter().map(Ok).collect(); - - let next_state = (store, prefix, offset, next_page_token, 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)], (store, prefix, offset, None, false))), - } + StreamState { + store, + prefix, + offset, + page_token: None, + has_more: true, }, + move |state| stream_step(state, chunk_size), ) .flat_map(futures::stream::iter); From 45b0a652bfe40a806d5146e5c0355a280efa30a1 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Tue, 2 Jun 2026 16:26:19 -0400 Subject: [PATCH 13/15] nits --- obstore/src/list.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/obstore/src/list.rs b/obstore/src/list.rs index 567782ee..757fc6f7 100644 --- a/obstore/src/list.rs +++ b/obstore/src/list.rs @@ -517,7 +517,6 @@ async fn list_with_delimiter_materialize( } /// Internal stream state -#[derive(Clone)] struct StreamState { store: Arc, prefix: Option, @@ -542,11 +541,11 @@ async fn stream_step( return None; } - let opts: PaginatedListOptions = PaginatedListOptions { + let opts = PaginatedListOptions { offset: offset.clone(), delimiter: None, max_keys: Some(chunk_size), - page_token: page_token.clone(), + page_token, ..Default::default() }; From 321dc8d335905850c5048e1ffe06657bf4c27054 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Thu, 4 Jun 2026 17:59:57 -0400 Subject: [PATCH 14/15] add minio integration tests for paginated list Cover the native PaginatedListStore path through obstore.list: multi-page auto-pagination (chunk_size as max_keys forces several continuation tokens), raw-string substring prefixes, and exclusive offset. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/store/test_s3.py | 44 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) 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") From 231ac4d394be7894445a39efaccb10e04d263e19 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Thu, 4 Jun 2026 18:08:03 -0400 Subject: [PATCH 15/15] fix substring prefix listing for emulated stores emulate_paginated_list compared the full object location against only the substring after the last '/', so it matched nothing, and it used list_with_delimiter (one level) which dropped nested keys. List recursively under the last complete path segment and keep keys whose full location starts with the requested prefix string, matching the native PaginatedListStore semantics. Propagate list errors via the stream rather than swallowing them. Tests: emulated (MemoryStore) and local-store substring cases with nested keys, plus a shared helper asserting emulated/native parity. Co-Authored-By: Claude Opus 4.8 (1M context) --- obstore/src/list.rs | 65 +++++++++++---------------- tests/test_list.py | 105 +++++++++++++++++++++++++++----------------- 2 files changed, 90 insertions(+), 80 deletions(-) diff --git a/obstore/src/list.rs b/obstore/src/list.rs index 757fc6f7..0dc96ebe 100644 --- a/obstore/src/list.rs +++ b/obstore/src/list.rs @@ -381,50 +381,35 @@ async fn emulate_paginated_list( store: &Arc, prefix: Option<&str>, ) -> object_store::Result { - // Split a path like "some/prefix/abc" into (Some(Path("some/prefix")), Some("abc")) - // This allows us to do a substring prefix match after the / delimiter - let (list_path, list_prefix_match): (Option, Option) = - if let Some(list_prefix) = prefix { - if let Some((list_path, list_prefix_match)) = list_prefix.rsplit_once('/') { - // There's a / in the prefix, so we assume the part before the last / is a - // path, and the end is a substring match - ( - Some(object_store::path::Path::parse(list_path)?), - Some(list_prefix_match.to_string()), - ) - } else { - // No / in prefix, so we assume it's a substring - (None, Some(list_prefix.to_string())) - } - } else { - (None, None) - }; - - let list_result = store.list_with_delimiter(list_path.as_ref()).await?; + // `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, + }; - // Filter list result to include only results with the given prefix after the / delimiter - let filtered_list_result = if let Some(list_prefix_match) = list_prefix_match { - let filtered_common_prefixes = list_result - .common_prefixes - .into_iter() - .filter(|p| p.as_ref().starts_with(&list_prefix_match)) - .collect(); - let filtered_objects = list_result - .objects - .into_iter() - .filter(|obj| obj.location.as_ref().starts_with(&list_prefix_match)) - .collect(); - ListResult { - common_prefixes: filtered_common_prefixes, - objects: filtered_objects, + 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), } - } else { - list_result - }; + } Ok(PaginatedListResult { - result: filtered_list_result, - // emulated stores do not support pagination + result: ListResult { + common_prefixes: Vec::new(), + objects, + }, + // Emulated stores return everything in a single page. page_token: None, }) } diff --git a/tests/test_list.py b/tests/test_list.py index b0485ee8..cee0365e 100644 --- a/tests/test_list.py +++ b/tests/test_list.py @@ -5,7 +5,7 @@ import pytest from arro3.core import RecordBatch, Table -from obstore.store import LocalStore, MemoryStore +from obstore.store import LocalStore, MemoryStore, ObjectStore, S3Store def test_list(): @@ -165,32 +165,31 @@ async def test_list_with_delimiter_async(): assert objects["path"][1].as_py() == "a/file2.txt" -def test_list_substring_filtering(): +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/another.csv", b"baz") + 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") - # 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() + # 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 == 2 + assert batch.num_rows == 3 def test_list_substring_filtering_local_store(): @@ -198,29 +197,55 @@ def test_list_substring_filtering_local_store(): temp_dir_path = Path(temp_dir) store = LocalStore(temp_dir_path) - # Create directory structure + # Create directory structure, including a nested directory under a matching + # prefix to verify recursive listing. 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] + (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") - # 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 + # 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():