Add option to return objects as Arrow in list_with_delimiter (#238)

kylebarron · web-flow · commit cddd07e2df0b · 2025-02-06T23:52:37.000-05:00
* Return ListResult as Arrow

* Add tests for list_with_delimiter
diff --git a/obstore/python/obstore/_list.pyi b/obstore/python/obstore/_list.pyi
@@ -26,7 +26,7 @@ class ObjectMeta(TypedDict):
     version: str | None
     """A version indicator for this object"""
 
-class ListResult(TypedDict):
+class ListResult(TypedDict, Generic[ChunkType]):
     """
     Result of a list call that includes objects, prefixes (directories) and a token for
     the next set of results. Individual result sets may be limited to 1,000 objects
@@ -36,7 +36,7 @@ class ListResult(TypedDict):
     common_prefixes: List[str]
     """Prefixes that are common (like directories)"""
 
-    objects: List[ObjectMeta]
+    objects: ChunkType
     """Object metadata for the listing"""
 
 ChunkType = TypeVar("ChunkType", List[ObjectMeta], RecordBatch)
@@ -188,7 +188,26 @@ def list(
         A ListStream, which you can iterate through to access list results.
     """
 
-def list_with_delimiter(store: ObjectStore, prefix: str | None = None) -> ListResult:
+@overload
+def list_with_delimiter(
+    store: ObjectStore,
+    prefix: str | None = None,
+    *,
+    return_arrow: Literal[True],
+) -> ListResult[RecordBatch]: ...
+@overload
+def list_with_delimiter(
+    store: ObjectStore,
+    prefix: str | None = None,
+    *,
+    return_arrow: Literal[False] = False,
+) -> ListResult[List[ObjectMeta]]: ...
+def list_with_delimiter(
+    store: ObjectStore,
+    prefix: str | None = None,
+    *,
+    return_arrow: bool = False,
+) -> ListResult[RecordBatch] | ListResult[List[ObjectMeta]]:
     """
     List objects with the given prefix and an implementation specific
     delimiter. Returns common prefixes (directories) in addition to object
@@ -209,9 +228,26 @@ def list_with_delimiter(store: ObjectStore, prefix: str | None = None) -> ListRe
         ListResult
     """
 
+@overload
 async def list_with_delimiter_async(
-    store: ObjectStore, prefix: str | None = None
-) -> ListResult:
+    store: ObjectStore,
+    prefix: str | None = None,
+    *,
+    return_arrow: Literal[True],
+) -> ListResult[RecordBatch]: ...
+@overload
+async def list_with_delimiter_async(
+    store: ObjectStore,
+    prefix: str | None = None,
+    *,
+    return_arrow: Literal[False] = False,
+) -> ListResult[List[ObjectMeta]]: ...
+async def list_with_delimiter_async(
+    store: ObjectStore,
+    prefix: str | None = None,
+    *,
+    return_arrow: bool = False,
+) -> ListResult[RecordBatch] | ListResult[List[ObjectMeta]]:
     """Call `list_with_delimiter` asynchronously.
 
     Refer to the documentation for
diff --git a/obstore/src/list.rs b/obstore/src/list.rs
@@ -11,9 +11,9 @@ use indexmap::IndexMap;
 use object_store::path::Path;
 use object_store::{ListResult, ObjectMeta, ObjectStore};
 use pyo3::exceptions::{PyImportError, PyStopAsyncIteration, PyStopIteration};
-use pyo3::intern;
 use pyo3::prelude::*;
 use pyo3::types::PyDict;
+use pyo3::{intern, IntoPyObjectExt};
 use pyo3_arrow::PyRecordBatch;
 use pyo3_object_store::{get_runtime, PyObjectStore, PyObjectStoreError, PyObjectStoreResult};
 use tokio::sync::Mutex;
@@ -47,17 +47,11 @@ impl<'py> IntoPyObject<'py> for PyObjectMeta {
         let mut dict = IndexMap::with_capacity(5);
         // Note, this uses "path" instead of "location" because we standardize the API to accept
         // the keyword "path" everywhere.
-        dict.insert(
-            "path",
-            self.0.location.as_ref().into_pyobject(py)?.into_any(),
-        );
-        dict.insert(
-            "last_modified",
-            self.0.last_modified.into_pyobject(py)?.into_any(),
-        );
-        dict.insert("size", self.0.size.into_pyobject(py)?.into_any());
-        dict.insert("e_tag", self.0.e_tag.into_pyobject(py)?.into_any());
-        dict.insert("version", self.0.version.into_pyobject(py)?);
+        dict.insert("path", self.0.location.as_ref().into_bound_py_any(py)?);
+        dict.insert("last_modified", self.0.last_modified.into_bound_py_any(py)?);
+        dict.insert("size", self.0.size.into_bound_py_any(py)?);
+        dict.insert("e_tag", self.0.e_tag.into_bound_py_any(py)?);
+        dict.insert("version", self.0.version.into_bound_py_any(py)?);
         dict.into_pyobject(py)
     }
 }
@@ -317,7 +311,19 @@ fn object_meta_to_arrow(metas: &[PyObjectMeta]) -> PyRecordBatchWrapper {
     PyRecordBatchWrapper::new(batch)
 }
 
-pub(crate) struct PyListResult(ListResult);
+pub(crate) struct PyListResult {
+    result: ListResult,
+    return_arrow: bool,
+}
+
+impl PyListResult {
+    fn new(result: ListResult, return_arrow: bool) -> Self {
+        Self {
+            result,
+            return_arrow,
+        }
+    }
+}
 
 impl<'py> IntoPyObject<'py> for PyListResult {
     type Target = PyDict;
@@ -328,24 +334,25 @@ impl<'py> IntoPyObject<'py> for PyListResult {
         let mut dict = IndexMap::with_capacity(2);
         dict.insert(
             "common_prefixes",
-            self.0
+            self.result
                 .common_prefixes
                 .into_iter()
                 .map(String::from)
                 .collect::<Vec<_>>()
-                .into_pyobject(py)?
-                .into_any(),
-        );
-        dict.insert(
-            "objects",
-            self.0
-                .objects
-                .into_iter()
-                .map(PyObjectMeta)
-                .collect::<Vec<_>>()
-                .into_pyobject(py)?
-                .into_any(),
+                .into_bound_py_any(py)?,
         );
+        let objects = self
+            .result
+            .objects
+            .into_iter()
+            .map(PyObjectMeta)
+            .collect::<Vec<_>>();
+        let objects = if self.return_arrow {
+            object_meta_to_arrow(&objects).into_bound_py_any(py)
+        } else {
+            objects.into_bound_py_any(py)
+        }?;
+        dict.insert("objects", objects);
         dict.into_pyobject(py)
     }
 }
@@ -383,41 +390,48 @@ pub(crate) fn list(
 }
 
 #[pyfunction]
-#[pyo3(signature = (store, prefix=None))]
+#[pyo3(signature = (store, prefix=None, *, return_arrow=false))]
 pub(crate) fn list_with_delimiter(
     py: Python,
     store: PyObjectStore,
     prefix: Option<String>,
+    return_arrow: bool,
 ) -> PyObjectStoreResult<PyListResult> {
     let runtime = get_runtime(py)?;
     py.allow_threads(|| {
         let out = runtime.block_on(list_with_delimiter_materialize(
             store.into_inner(),
             prefix.map(|s| s.into()).as_ref(),
+            return_arrow,
         ))?;
         Ok::<_, PyObjectStoreError>(out)
     })
 }
 
 #[pyfunction]
-#[pyo3(signature = (store, prefix=None))]
+#[pyo3(signature = (store, prefix=None, *, return_arrow=false))]
 pub(crate) fn list_with_delimiter_async(
     py: Python,
     store: PyObjectStore,
     prefix: Option<String>,
+    return_arrow: bool,
 ) -> PyResult<Bound<PyAny>> {
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
-        let out =
-            list_with_delimiter_materialize(store.into_inner(), prefix.map(|s| s.into()).as_ref())
-                .await?;
+        let out = list_with_delimiter_materialize(
+            store.into_inner(),
+            prefix.map(|s| s.into()).as_ref(),
+            return_arrow,
+        )
+        .await?;
         Ok(out)
     })
 }
 
 async fn list_with_delimiter_materialize(
     store: Arc<dyn ObjectStore>,
     prefix: Option<&Path>,
+    return_arrow: bool,
 ) -> PyObjectStoreResult<PyListResult> {
     let list_result = store.list_with_delimiter(prefix).await?;
-    Ok(PyListResult(list_result))
+    Ok(PyListResult::new(list_result, return_arrow))
 }
diff --git a/obstore/src/put.rs b/obstore/src/put.rs
@@ -12,10 +12,10 @@ use object_store::{
     WriteMultipart,
 };
 use pyo3::exceptions::{PyStopAsyncIteration, PyStopIteration, PyValueError};
-use pyo3::intern;
 use pyo3::prelude::*;
 use pyo3::pybacked::PyBackedStr;
 use pyo3::types::PyDict;
+use pyo3::{intern, IntoPyObjectExt};
 use pyo3_bytes::PyBytes;
 use pyo3_file::PyFileLikeObject;
 use pyo3_object_store::{get_runtime, PyObjectStore, PyObjectStoreResult};
@@ -283,8 +283,8 @@ impl<'py> IntoPyObject<'py> for PyPutResult {
 
     fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
         let mut dict = IndexMap::with_capacity(2);
-        dict.insert("e_tag", self.0.e_tag.into_pyobject(py)?.into_any());
-        dict.insert("version", self.0.version.into_pyobject(py)?.into_any());
+        dict.insert("e_tag", self.0.e_tag.into_bound_py_any(py)?);
+        dict.insert("version", self.0.version.into_bound_py_any(py)?);
         dict.into_pyobject(py)
     }
 }
diff --git a/pyo3-object_store/src/http.rs b/pyo3-object_store/src/http.rs
@@ -17,7 +17,7 @@ struct HTTPConfig {
 
 impl HTTPConfig {
     fn __getnewargs_ex__(&self, py: Python) -> PyResult<PyObject> {
-        let args = PyTuple::new(py, vec![self.url.clone().into_pyobject(py)?])?.into_py_any(py)?;
+        let args = PyTuple::new(py, vec![self.url.clone()])?.into_py_any(py)?;
         let kwargs = PyDict::new(py);
 
         if let Some(client_options) = &self.client_options {
diff --git a/pyo3-object_store/src/local.rs b/pyo3-object_store/src/local.rs
@@ -20,8 +20,7 @@ struct LocalConfig {
 
 impl LocalConfig {
     fn __getnewargs_ex__(&self, py: Python) -> PyResult<PyObject> {
-        let args =
-            PyTuple::new(py, vec![self.prefix.clone().into_pyobject(py)?])?.into_py_any(py)?;
+        let args = PyTuple::new(py, vec![self.prefix.clone()])?.into_py_any(py)?;
         let kwargs = PyDict::new(py);
         kwargs.set_item(intern!(py, "automatic_cleanup"), self.automatic_cleanup)?;
         kwargs.set_item(intern!(py, "mkdir"), self.mkdir)?;
@@ -119,10 +118,8 @@ impl PyLocalStore {
         // So we manually convert to a pathlib.Path
         if let Some(prefix) = &self.config.prefix {
             let pathlib_mod = py.import(intern!(py, "pathlib"))?;
-            let path_object = pathlib_mod.call_method1(
-                intern!(py, "Path"),
-                PyTuple::new(py, vec![prefix.into_pyobject(py)?])?,
-            )?;
+            let path_object =
+                pathlib_mod.call_method1(intern!(py, "Path"), PyTuple::new(py, vec![prefix])?)?;
             path_object.into_py_any(py)
         } else {
             Ok(py.None())
diff --git a/pyo3-object_store/src/simple.rs b/pyo3-object_store/src/simple.rs
@@ -39,7 +39,7 @@ pub fn from_url(
                 retry_config,
                 kwargs.map(|x| x.extract()).transpose()?,
             )?;
-            Ok(store.into_pyobject(py)?.into_py_any(py)?)
+            Ok(store.into_py_any(py)?)
         }
         ObjectStoreScheme::GoogleCloudStorage => {
             let store = PyGCSStore::from_url(
@@ -50,7 +50,7 @@ pub fn from_url(
                 retry_config,
                 kwargs.map(|x| x.extract()).transpose()?,
             )?;
-            Ok(store.into_pyobject(py)?.into_py_any(py)?)
+            Ok(store.into_py_any(py)?)
         }
         ObjectStoreScheme::MicrosoftAzure => {
             let store = PyAzureStore::from_url(
@@ -61,7 +61,7 @@ pub fn from_url(
                 retry_config,
                 kwargs.map(|x| x.extract()).transpose()?,
             )?;
-            Ok(store.into_pyobject(py)?.into_py_any(py)?)
+            Ok(store.into_py_any(py)?)
         }
         ObjectStoreScheme::Http => {
             raise_if_config_passed(config, kwargs, "http")?;
@@ -71,7 +71,7 @@ pub fn from_url(
                 client_options,
                 retry_config,
             )?;
-            Ok(store.into_pyobject(py)?.into_py_any(py)?)
+            Ok(store.into_py_any(py)?)
         }
         ObjectStoreScheme::Local => {
             let mut automatic_cleanup = false;
@@ -92,12 +92,12 @@ pub fn from_url(
                 automatic_cleanup,
                 mkdir,
             )?;
-            Ok(store.into_pyobject(py)?.into_py_any(py)?)
+            Ok(store.into_py_any(py)?)
         }
         ObjectStoreScheme::Memory => {
             raise_if_config_passed(config, kwargs, "memory")?;
             let store: PyMemoryStore = Arc::new(InMemory::new()).into();
-            Ok(store.into_pyobject(py)?.into_py_any(py)?)
+            Ok(store.into_py_any(py)?)
         }
         scheme => Err(GenericError::new_err(format!("Unknown URL scheme {:?}", scheme,)).into()),
     }
diff --git a/tests/test_list.py b/tests/test_list.py
@@ -58,3 +58,68 @@ async def test_list_stream_async():
     batch = await stream.collect_async()
     assert isinstance(batch, RecordBatch)
     assert batch.num_rows == 100
+
+
+def test_list_with_delimiter():
+    store = MemoryStore()
+
+    obs.put(store, "a/file1.txt", b"foo")
+    obs.put(store, "a/file2.txt", b"bar")
+    obs.put(store, "b/file3.txt", b"baz")
+
+    list_result1 = obs.list_with_delimiter(store)
+    assert list_result1["common_prefixes"] == ["a", "b"]
+    assert list_result1["objects"] == []
+
+    list_result2 = obs.list_with_delimiter(store, "a")
+    assert list_result2["common_prefixes"] == []
+    assert list_result2["objects"][0]["path"] == "a/file1.txt"
+    assert list_result2["objects"][1]["path"] == "a/file2.txt"
+
+    list_result3 = obs.list_with_delimiter(store, "b")
+    assert list_result3["common_prefixes"] == []
+    assert list_result3["objects"][0]["path"] == "b/file3.txt"
+
+    # Test returning arrow
+    list_result1 = obs.list_with_delimiter(store, return_arrow=True)
+    assert list_result1["common_prefixes"] == ["a", "b"]
+    assert list_result1["objects"].num_rows == 0
+
+    list_result2 = obs.list_with_delimiter(store, "a", return_arrow=True)
+    assert list_result2["common_prefixes"] == []
+    assert list_result2["objects"].num_rows == 2
+    assert list_result2["objects"]["path"][0].as_py() == "a/file1.txt"
+    assert list_result2["objects"]["path"][1].as_py() == "a/file2.txt"
+
+
+@pytest.mark.asyncio
+async def test_list_with_delimiter_async():
+    store = MemoryStore()
+
+    await obs.put_async(store, "a/file1.txt", b"foo")
+    await obs.put_async(store, "a/file2.txt", b"bar")
+    await obs.put_async(store, "b/file3.txt", b"baz")
+
+    list_result1 = await obs.list_with_delimiter_async(store)
+    assert list_result1["common_prefixes"] == ["a", "b"]
+    assert list_result1["objects"] == []
+
+    list_result2 = await obs.list_with_delimiter_async(store, "a")
+    assert list_result2["common_prefixes"] == []
+    assert list_result2["objects"][0]["path"] == "a/file1.txt"
+    assert list_result2["objects"][1]["path"] == "a/file2.txt"
+
+    list_result3 = await obs.list_with_delimiter_async(store, "b")
+    assert list_result3["common_prefixes"] == []
+    assert list_result3["objects"][0]["path"] == "b/file3.txt"
+
+    # Test returning arrow
+    list_result1 = await obs.list_with_delimiter_async(store, return_arrow=True)
+    assert list_result1["common_prefixes"] == ["a", "b"]
+    assert list_result1["objects"].num_rows == 0
+
+    list_result2 = await obs.list_with_delimiter_async(store, "a", return_arrow=True)
+    assert list_result2["common_prefixes"] == []
+    assert list_result2["objects"].num_rows == 2
+    assert list_result2["objects"]["path"][0].as_py() == "a/file1.txt"
+    assert list_result2["objects"]["path"][1].as_py() == "a/file2.txt"
