feat: open_reader accepts optional size to skip HEAD request

Author: fvaleye

obstore/python/obstore/_buffered.pyi

@@ -21,6 +21,7 @@ def open_reader(
     path: str,
     *,
     buffer_size: int = 1024 * 1024,
+    size: int | None = None,
 ) -> ReadableFile:
     """Open a readable file object from the specified location.
 
@@ -30,6 +31,11 @@ def open_reader(
 
     Keyword Args:
         buffer_size: The minimum number of bytes to read in a single request. Up to `buffer_size` bytes will be buffered in memory.
+        size: Optional byte size of the object. When provided, skips the HEAD request used to fetch the file size. Useful for callers that already know the size from external metadata.
+
+            The caller is responsible for accuracy: a value larger than the actual file surfaces as a read-time range error, a value smaller causes silent truncation.
+
+            When `size` is provided, the resulting reader's `meta` attribute omits `last_modified` (since it was not fetched). Callers that need that field should call `open_reader` without `size`. Defaults to `None`.
 
     Returns:
         ReadableFile
@@ -41,6 +47,7 @@ async def open_reader_async(
     path: str,
     *,
     buffer_size: int = 1024 * 1024,
+    size: int | None = None,
 ) -> AsyncReadableFile:
     """Call `open_reader` asynchronously, returning a readable file object with asynchronous operations.
 

obstore/src/buffered.rs

@@ -2,11 +2,13 @@ use std::io::SeekFrom;
 use std::sync::Arc;
 
 use bytes::Bytes;
+use chrono::{DateTime, Utc};
+use indexmap::IndexMap;
 use object_store::buffered::{BufReader, BufWriter};
 use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt};
 use pyo3::exceptions::{PyIOError, PyStopAsyncIteration, PyStopIteration};
 use pyo3::prelude::*;
-use pyo3::types::PyString;
+use pyo3::types::{PyDict, PyString};
 use pyo3::{intern, IntoPyObjectExt};
 use pyo3_async_runtimes::tokio::{future_into_py, get_runtime};
 use pyo3_bytes::PyBytes;
@@ -15,62 +17,79 @@ use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncSeekExt, AsyncWriteExt, Line
 use tokio::sync::Mutex;
 
 use crate::attributes::PyAttributes;
-use crate::list::PyObjectMeta;
 use crate::tags::PyTagSet;
 
 #[pyfunction]
-#[pyo3(signature = (store, path, *, buffer_size=1024 * 1024))]
+#[pyo3(signature = (store, path, *, buffer_size=1024 * 1024, size=None))]
 pub(crate) fn open_reader(
     py: Python,
     store: PyObjectStore,
     path: PyPath,
     buffer_size: usize,
+    size: Option<u64>,
 ) -> PyObjectStoreResult<PyReadableFile> {
     let store = store.into_inner();
     let runtime = get_runtime();
-    let (reader, meta) = py.detach(|| runtime.block_on(create_reader(store, path, buffer_size)))?;
-    Ok(PyReadableFile::new(reader, meta, false))
+    let was_hinted = size.is_some();
+    let (reader, meta) =
+        py.detach(|| runtime.block_on(create_reader(store, path, buffer_size, size)))?;
+    Ok(PyReadableFile::new(reader, meta, was_hinted, false))
 }
 
 #[pyfunction]
-#[pyo3(signature = (store, path, *, buffer_size=1024 * 1024))]
+#[pyo3(signature = (store, path, *, buffer_size=1024 * 1024, size=None))]
 pub(crate) fn open_reader_async(
     py: Python,
     store: PyObjectStore,
     path: PyPath,
     buffer_size: usize,
+    size: Option<u64>,
 ) -> PyResult<Bound<PyAny>> {
     let store = store.into_inner();
+    let was_hinted = size.is_some();
     future_into_py(py, async move {
-        let (reader, meta) = create_reader(store, path, buffer_size).await?;
-        Ok(PyReadableFile::new(reader, meta, true))
+        let (reader, meta) = create_reader(store, path, buffer_size, size).await?;
+        Ok(PyReadableFile::new(reader, meta, was_hinted, true))
     })
 }
 
 async fn create_reader(
     store: Arc<dyn ObjectStore>,
     path: PyPath,
     capacity: usize,
+    size: Option<u64>,
 ) -> PyObjectStoreResult<(BufReader, ObjectMeta)> {
-    let meta = store
-        .head(path.as_ref())
-        .await
-        .map_err(PyObjectStoreError::ObjectStoreError)?;
+    let meta = match size {
+        Some(size) => ObjectMeta {
+            location: path.as_ref().clone(),
+            last_modified: DateTime::<Utc>::from_timestamp(0, 0)
+                .expect("unix epoch is a valid DateTime"),
+            size,
+            e_tag: None,
+            version: None,
+        },
+        None => store
+            .head(path.as_ref())
+            .await
+            .map_err(PyObjectStoreError::ObjectStoreError)?,
+    };
     Ok((BufReader::with_capacity(store, &meta, capacity), meta))
 }
 
 #[pyclass(name = "ReadableFile", frozen)]
 pub(crate) struct PyReadableFile {
     reader: Arc<Mutex<BufReader>>,
     meta: ObjectMeta,
+    was_hinted: bool,
     r#async: bool,
 }
 
 impl PyReadableFile {
-    fn new(reader: BufReader, meta: ObjectMeta, r#async: bool) -> Self {
+    fn new(reader: BufReader, meta: ObjectMeta, was_hinted: bool, r#async: bool) -> Self {
         Self {
             reader: Arc::new(Mutex::new(reader)),
             meta,
+            was_hinted,
             r#async,
         }
     }
@@ -92,8 +111,19 @@ impl PyReadableFile {
     fn close(&self) {}
 
     #[getter]
-    fn meta(&self) -> PyObjectMeta {
-        self.meta.clone().into()
+    fn meta<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
+        let mut dict = IndexMap::with_capacity(5);
+        dict.insert("path", self.meta.location.as_ref().into_bound_py_any(py)?);
+        if !self.was_hinted {
+            dict.insert(
+                "last_modified",
+                self.meta.last_modified.into_bound_py_any(py)?,
+            );
+        }
+        dict.insert("size", self.meta.size.into_bound_py_any(py)?);
+        dict.insert("e_tag", self.meta.e_tag.clone().into_bound_py_any(py)?);
+        dict.insert("version", self.meta.version.clone().into_bound_py_any(py)?);
+        dict.into_pyobject(py)
     }
 
     #[pyo3(signature = (size = None, /))]

tests/test_buffered.py

@@ -112,3 +112,77 @@ async def test_read_past_eof_async():
     buf = BytesIO(data)
     expected = buf.read(20)
     assert memoryview(expected) == memoryview(buffer)
+
+
+def test_open_reader_size_hint_sync():
+    store = MemoryStore()
+    data = b"x" * 1000
+    path = "sized.bin"
+    obs.put(store, path, data)
+
+    file = obs.open_reader(store, path, size=len(data))
+    assert file.size == len(data)
+    assert memoryview(data) == memoryview(file.read())
+
+
+@pytest.mark.asyncio
+async def test_open_reader_size_hint_async():
+    store = MemoryStore()
+    data = b"x" * 1000
+    path = "sized.bin"
+    await obs.put_async(store, path, data)
+
+    file = await obs.open_reader_async(store, path, size=len(data))
+    assert file.size == len(data)
+    assert memoryview(data) == memoryview(await file.read())
+
+
+def test_open_reader_size_hint_larger_than_actual_errors_on_read():
+    store = MemoryStore()
+    data = b"x" * 1000
+    path = "sized.bin"
+    obs.put(store, path, data)
+
+    file = obs.open_reader(store, path, size=5000)
+    assert file.size == 5000
+    with pytest.raises(OSError, match="range"):
+        file.read()
+
+
+def test_open_reader_size_hint_smaller_than_actual_truncates():
+    store = MemoryStore()
+    data = b"x" * 1000
+    path = "sized.bin"
+    obs.put(store, path, data)
+
+    file = obs.open_reader(store, path, size=500)
+    assert file.size == 500
+    buffer = file.read()
+    assert memoryview(data[:500]) == memoryview(buffer)
+
+
+def test_open_reader_size_hint_zero_byte_file():
+    store = MemoryStore()
+    path = "empty.bin"
+    obs.put(store, path, b"")
+
+    file = obs.open_reader(store, path, size=0)
+    assert file.size == 0
+    assert memoryview(b"") == memoryview(file.read())
+
+
+def test_open_reader_meta_last_modified_depends_on_size_hint():
+    store = MemoryStore()
+    data = b"x" * 1000
+    path = "sized.bin"
+    obs.put(store, path, data)
+
+    hinted = obs.open_reader(store, path, size=len(data))
+    unhinted = obs.open_reader(store, path)
+
+    assert "last_modified" not in hinted.meta
+    assert "last_modified" in unhinted.meta
+    assert hinted.meta["size"] == len(data)
+    assert unhinted.meta["size"] == len(data)
+    assert hinted.meta["e_tag"] is None
+    assert hinted.meta["path"] == path