diff --git a/obstore/python/obstore/_put.pyi b/obstore/python/obstore/_put.pyi index 30ea7f91..47005fc1 100644 --- a/obstore/python/obstore/_put.pyi +++ b/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. diff --git a/obstore/src/put.rs b/obstore/src/put.rs index 68433107..c46ea608 100644 --- a/obstore/src/put.rs +++ b/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), FileLike(PyFileLikeObject), - Buffer(Cursor), } impl PullSource { @@ -78,8 +78,8 @@ impl PullSource { } /// Whether to use multipart uploads. - fn use_multipart(&mut self, chunk_size: usize) -> PyObjectStoreResult { - Ok(self.nbytes()? > chunk_size) + fn use_multipart(&mut self, chunk_size: NonZeroUsize) -> PyObjectStoreResult { + 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 { + fn use_multipart(&mut self, chunk_size: NonZeroUsize) -> PyObjectStoreResult { 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 { 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::() { - 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, mode: Option, use_multipart: Option, - chunk_size: usize, + chunk_size: NonZeroUsize, max_concurrency: usize, ) -> PyObjectStoreResult { 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, mode: Option, use_multipart: Option, - chunk_size: usize, + chunk_size: NonZeroUsize, max_concurrency: usize, ) -> PyResult> { let mut use_multipart = if let Some(use_multipart) = use_multipart { @@ -442,7 +440,7 @@ async fn put_multipart_inner( store: Arc, path: &Path, reader: PutInput, - chunk_size: usize, + chunk_size: NonZeroUsize, max_concurrency: usize, attributes: Option, tags: Option, @@ -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?; diff --git a/tests/test_put.py b/tests/test_put.py index 4f60d5e8..154bd2d1 100644 --- a/tests/test_put.py +++ b/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