perf: Reduce memory overhead for multipart uploads

obstore/python/obstore/_put.pyi

@@ -109,7 +109,7 @@ def put(
     tags: dict[str, str] | None = None,
     mode: PutMode | None = None,
     use_multipart: bool | None = None,
-    chunk_size: int = ...,
+    chunk_size: int = 5 * 1024 * 1024,
     max_concurrency: int = 12,
 ) -> PutResult:
     """Save the provided bytes to the specified location.
@@ -179,7 +179,7 @@ async def put_async(
     tags: dict[str, str] | None = None,
     mode: PutMode | None = None,
     use_multipart: bool | None = None,
-    chunk_size: int = ...,
+    chunk_size: int = 5 * 1024 * 1024,
     max_concurrency: int = 12,
 ) -> PutResult:
     """Call `put` asynchronously.

obstore/src/put.rs

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 use std::fs::File;
-use std::io::{BufReader, Cursor, Read, Seek, SeekFrom};
+use std::io::{BufReader, Read, Seek, SeekFrom};
+use std::num::NonZeroUsize;
 use std::path::PathBuf;
 use std::sync::Arc;
 
@@ -65,7 +66,6 @@ impl<'py> FromPyObject<'_, 'py> for PyUpdateVersion {
 pub(crate) enum PullSource {
     File(BufReader<File>),
     FileLike(PyFileLikeObject),
-    Buffer(Cursor<Bytes>),
 }
 
 impl PullSource {
@@ -78,8 +78,8 @@ impl PullSource {
     }
 
     /// Whether to use multipart uploads.
-    fn use_multipart(&mut self, chunk_size: usize) -> PyObjectStoreResult<bool> {
-        Ok(self.nbytes()? > chunk_size)
+    fn use_multipart(&mut self, chunk_size: NonZeroUsize) -> PyObjectStoreResult<bool> {
+        Ok(self.nbytes()? > chunk_size.get())
     }
 }
 
@@ -88,7 +88,6 @@ impl Read for PullSource {
         match self {
             Self::File(f) => f.read(buf),
             Self::FileLike(f) => f.read(buf),
-            Self::Buffer(f) => f.read(buf),
         }
     }
 }
@@ -98,7 +97,6 @@ impl Seek for PullSource {
         match self {
             Self::File(f) => f.seek(pos),
             Self::FileLike(f) => f.seek(pos),
-            Self::Buffer(f) => f.seek(pos),
         }
     }
 }
@@ -199,6 +197,9 @@ impl AsyncPushSource {
 
 // #[derive(Debug)]
 pub(crate) enum PutInput {
+    /// A buffer protocol object
+    Buffer(Bytes),
+
     /// Input that we can pull from
     Pull(PullSource),
 
@@ -211,25 +212,24 @@ pub(crate) enum PutInput {
 
 impl PutInput {
     /// Whether to use multipart uploads.
-    fn use_multipart(&mut self, chunk_size: usize) -> PyObjectStoreResult<bool> {
+    fn use_multipart(&mut self, chunk_size: NonZeroUsize) -> PyObjectStoreResult<bool> {
         match self {
+            Self::Buffer(buffer) => Ok(buffer.len() > chunk_size.get()),
             Self::Pull(pull_source) => pull_source.use_multipart(chunk_size),
             // We always use multipart uploads for push-based sources because we have no way of
             // knowing how large they'll be and we don't want to buffer them into memory.
-            _ => Ok(true),
+            Self::SyncPush(_) | Self::AsyncPush(_) => Ok(true),
         }
     }
 
     async fn read_all(&mut self) -> PyObjectStoreResult<PutPayload> {
         match self {
-            Self::Pull(pull_source) => match pull_source {
-                PullSource::Buffer(buffer) => Ok(buffer.get_ref().clone().into()),
-                source => {
-                    let mut buf = Vec::new();
-                    source.read_to_end(&mut buf)?;
-                    Ok(Bytes::from(buf).into())
-                }
-            },
+            Self::Buffer(buffer) => Ok(buffer.clone().into()),
+            Self::Pull(pull_source) => {
+                let mut buf = Vec::new();
+                pull_source.read_to_end(&mut buf)?;
+                Ok(Bytes::from(buf).into())
+            }
             Self::SyncPush(push_source) => push_source.read_all(),
             Self::AsyncPush(push_source) => push_source.read_all().await,
         }
@@ -246,9 +246,7 @@ impl<'py> FromPyObject<'_, 'py> for PutInput {
                 path,
             )?))))
         } else if let Ok(buffer) = obj.extract::<PyBytes>() {
-            Ok(Self::Pull(PullSource::Buffer(Cursor::new(
-                buffer.into_inner(),
-            ))))
+            Ok(Self::Buffer(buffer.into_inner()))
         }
         // Check for file-like object
         else if obj.hasattr(intern!(py, "read"))? && obj.hasattr(intern!(py, "seek"))? {
@@ -301,7 +299,7 @@ impl<'py> IntoPyObject<'py> for PyPutResult {
 }
 
 #[pyfunction]
-#[pyo3(signature = (store, path, file, *, attributes=None, tags=None, mode=None, use_multipart=None, chunk_size=5242880, max_concurrency=12))]
+#[pyo3(signature = (store, path, file, *, attributes=None, tags=None, mode=None, use_multipart=None, chunk_size=NonZeroUsize::new(5 * 1024 * 1024).unwrap(), max_concurrency=12))]
 #[allow(clippy::too_many_arguments)]
 pub(crate) fn put(
     py: Python,
@@ -312,7 +310,7 @@ pub(crate) fn put(
     tags: Option<PyTagSet>,
     mode: Option<PyPutMode>,
     use_multipart: Option<bool>,
-    chunk_size: usize,
+    chunk_size: NonZeroUsize,
     max_concurrency: usize,
 ) -> PyObjectStoreResult<PyPutResult> {
     if matches!(file, PutInput::AsyncPush(_)) {
@@ -360,7 +358,7 @@ pub(crate) fn put(
 }
 
 #[pyfunction]
-#[pyo3(signature = (store, path, file, *, attributes=None, tags=None, mode=None, use_multipart=None, chunk_size=5242880, max_concurrency=12))]
+#[pyo3(signature = (store, path, file, *, attributes=None, tags=None, mode=None, use_multipart=None, chunk_size=NonZeroUsize::new(5 * 1024 * 1024).unwrap(), max_concurrency=12))]
 #[allow(clippy::too_many_arguments)]
 pub(crate) fn put_async(
     py: Python,
@@ -371,7 +369,7 @@ pub(crate) fn put_async(
     tags: Option<PyTagSet>,
     mode: Option<PyPutMode>,
     use_multipart: Option<bool>,
-    chunk_size: usize,
+    chunk_size: NonZeroUsize,
     max_concurrency: usize,
 ) -> PyResult<Bound<PyAny>> {
     let mut use_multipart = if let Some(use_multipart) = use_multipart {
@@ -442,7 +440,7 @@ async fn put_multipart_inner(
     store: Arc<dyn ObjectStore>,
     path: &Path,
     reader: PutInput,
-    chunk_size: usize,
+    chunk_size: NonZeroUsize,
     max_concurrency: usize,
     attributes: Option<PyAttributes>,
     tags: Option<PyTagSet>,
@@ -457,7 +455,7 @@ async fn put_multipart_inner(
     }
 
     let upload = store.put_multipart_opts(path, opts).await?;
-    let mut writer = WriteMultipart::new_with_chunk_size(upload, chunk_size);
+    let mut writer = WriteMultipart::new_with_chunk_size(upload, chunk_size.get());
 
     // Make sure to call abort if the multipart upload failed for any reason
     match write_multipart(&mut writer, reader, chunk_size, max_concurrency).await {
@@ -472,21 +470,30 @@ async fn put_multipart_inner(
 async fn write_multipart(
     writer: &mut WriteMultipart,
     reader: PutInput,
-    chunk_size: usize,
+    chunk_size: NonZeroUsize,
     max_concurrency: usize,
 ) -> PyObjectStoreResult<()> {
-    // Match across pull, push, async push
     match reader {
-        PutInput::Pull(mut pull_reader) => loop {
-            let mut scratch_buffer = vec![0; chunk_size];
-            let read_size = pull_reader.read(&mut scratch_buffer)?;
-            if read_size == 0 {
-                break;
-            } else {
+        PutInput::Buffer(buffer) => {
+            let mut offset = 0;
+            while offset < buffer.len() {
+                let end = (offset + chunk_size.get()).min(buffer.len());
+                writer.wait_for_capacity(max_concurrency).await?;
+                writer.put(buffer.slice(offset..end));
+                offset = end;
+            }
+        }
+        PutInput::Pull(mut pull_reader) => {
+            let mut scratch_buffer = vec![0; chunk_size.get()];
+            loop {
+                let read_size = pull_reader.read(&mut scratch_buffer)?;
+                if read_size == 0 {
+                    break;
+                }
                 writer.wait_for_capacity(max_concurrency).await?;
                 writer.write(&scratch_buffer[0..read_size]);
             }
-        },
+        }
         PutInput::SyncPush(push_reader) => {
             for buf in push_reader {
                 writer.wait_for_capacity(max_concurrency).await?;

tests/test_put.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import itertools
 from tempfile import TemporaryDirectory
 
@@ -109,3 +111,53 @@ def test_put_sync_iterable_local_store():
         store.put(path, iterator)
 
         assert store.get(path).bytes() == data
+
+
+CHUNK_SIZE = 1024
+
+
+@pytest.mark.parametrize("nbytes", [0, 1, CHUNK_SIZE - 1, CHUNK_SIZE, CHUNK_SIZE + 1])
+def test_put_multipart_buffer_chunk_boundaries(nbytes: int):
+    """Buffers are chunked correctly regardless of alignment to chunk_size."""
+    store = MemoryStore()
+    data = bytes(range(256)) * (nbytes // 256 + 1)
+    data = data[:nbytes]
+
+    store.put("file1.txt", data, use_multipart=True, chunk_size=CHUNK_SIZE)
+
+    assert store.get("file1.txt").bytes() == data
+
+
+@pytest.mark.parametrize("wrapper", [bytes, bytearray, memoryview])
+def test_put_multipart_buffer_types(
+    wrapper: type[bytes | bytearray | memoryview],
+):
+    """Any buffer-protocol input is uploaded verbatim."""
+    store = MemoryStore()
+    data = b"the quick brown fox jumps over the lazy dog," * 500
+
+    store.put("file1.txt", wrapper(data), use_multipart=True, chunk_size=CHUNK_SIZE)
+
+    assert store.get("file1.txt").bytes() == data
+
+
+def test_put_multipart_buffer_slice():
+    """A memoryview into the middle of a larger buffer uploads only its own bytes."""
+    store = MemoryStore()
+    backing = bytes(range(256)) * 40
+    view = memoryview(backing)[1000:8000]
+
+    store.put("file1.txt", view, use_multipart=True, chunk_size=CHUNK_SIZE)
+
+    assert store.get("file1.txt").bytes() == backing[1000:8000]
+
+
+def test_put_multipart_buffer_local_store():
+    """The buffer multipart path round-trips through a store with real parts."""
+    with TemporaryDirectory() as tmpdir:
+        store = LocalStore(tmpdir)
+        data = b"the quick brown fox jumps over the lazy dog," * 500
+
+        store.put("big-data.txt", data, use_multipart=True, chunk_size=CHUNK_SIZE)
+
+        assert store.get("big-data.txt").bytes() == data