Support non-contiguous put payloads / vectored writes (#5514) #5538

tustvold · 2024-03-21T01:24:21Z

Which issue does this PR close?

Rationale for this change

Allowing non-contiguous payloads avoids needing to bump allocate buffers when writing data where the size isn't known up front. This can be significantly more efficient.

It will also open up interesting possibilities for the parquet writer, where the data is already buffered in chunks

What changes are included in this PR?

Are there any user-facing changes?

FYI @roeap @wjones127 @alamb

tustvold · 2024-03-21T03:14:57Z

object_store/src/aws/client.rs

@@ -438,16 +443,8 @@ impl S3Client {

        let mut builder = self.client.request(Method::POST, url);

-        // Compute checksum - S3 *requires* this for DeleteObjects requests, so we default to


Cleanup given SHA256 is the only checksum we support and are ever likely to

(and to be clear for anyone else reading this PR -- this code is not exposed publicly so it can be changed in the future if needed as well)

tustvold · 2024-03-21T03:16:16Z

object_store/src/client/retry.rs

        async move {
            let mut retries = 0;
            let now = Instant::now();

            loop {
-                let s = req.try_clone().expect("request body must be cloneable");
-                match client.execute(s).await {
+                let mut s = self.try_clone().expect("request body must be cloneable");


We defer setting the body until this point when we have a streaming body, as otherwise this would fail

tustvold · 2024-04-09T14:57:03Z

object_store/src/aws/client.rs

@@ -468,8 +468,8 @@ impl S3Client {
        let response = builder
            .header(CONTENT_TYPE, "application/xml")
            .body(body)
-            .with_aws_sigv4(credential.authorizer(), payload_sha256.as_deref())
-            .send_retry_with_idempotency(&self.config.retry_config, false)


false is the default for a POST request

tustvold · 2024-04-09T14:57:35Z

object_store/src/aws/dynamo.rs

@@ -186,11 +186,7 @@ impl DynamoCommit {
        to: &Path,
    ) -> Result<()> {
        self.conditional_op(client, to, None, || async {
-            client
-                .copy_request(from, to)
-                .set_idempotent(false)


false is the default

tustvold · 2024-04-09T14:59:31Z

object_store/src/azure/client.rs

@@ -237,21 +242,23 @@ impl AzureClient {
            builder = builder.header(CONTENT_TYPE, value);
        }

-        builder = builder


This is moved into send above

tustvold · 2024-04-09T15:00:09Z

object_store/src/client/retry.rs

+                .expect("request body must be cloneable");
+
+            if let Some(x) = &self.payload {
+                *s.body_mut() = Some(x.body());


Setting this here sidesteps issues around clone-ability of stream based request bodies

Maybe a function call like s.set_body(Some(x.body()) instead might be easier to understand the intent

This is an upstream type unfortunately

tustvold · 2024-04-09T15:01:15Z

object_store/src/gcp/client.rs

@@ -327,20 +335,24 @@ impl GoogleCloudStorageClient {
        let builder = self
            .client
            .request(Method::PUT, url)
-            .header(header::CONTENT_TYPE, content_type)
-            .header(header::CONTENT_LENGTH, payload.len())


Moved into send

tustvold · 2024-04-09T15:03:14Z

object_store/src/buffered.rs

@@ -231,7 +231,7 @@ impl std::fmt::Debug for BufWriter {

 enum BufWriterState {
    /// Buffer up to capacity bytes
-    Buffer(Path, Vec<u8>),
+    Buffer(Path, PutPayloadMut),


This change means that we no longer bump allocate 🎉

tustvold · 2024-04-09T15:03:49Z

object_store/src/upload.rs

@@ -121,7 +122,8 @@ impl WriteMultipart {
    pub fn new_with_chunk_size(upload: Box<dyn MultipartUpload>, chunk_size: usize) -> Self {
        Self {
            upload,
-            buffer: Vec::with_capacity(chunk_size),
+            chunk_size,
+            buffer: PutPayloadMut::new(),


This change means we avoid bump allocating but also don't allocate an entire 10MB buffer up-front only to potentially use some fraction of it

object_store/src/payload.rs

Xuanwo · 2024-04-09T15:23:39Z

object_store/src/payload.rs

+
+impl FromIterator<Bytes> for PutPayload {
+    fn from_iter<T: IntoIterator<Item = Bytes>>(iter: T) -> Self {
+        Self(iter.into_iter().collect())


Use Arc::from_iter() directly to reduce one alloaction (after TrustedLen stablized in the future).

I believe this does use Arc::from_iter?

Arc::from_iter has special logic for TrustedLen that can avoid the collect()

If you look at the definition of collect it just invokes from_iter for the collection being collected into, in this case Arc

fn collect<B: FromIterator<Self::Item>>(self) -> B where Self: Sized, { FromIterator::from_iter(self) }

alamb

Thank yoU @tustvold -- I think the structure of this code / API is quite nice and elegant 👌 👏

I do think it is a non trivial API change (not complicated, but will require a bunch of mechanical changes) and so marked the PR as API change. To that end I suggest we add some better documentation to help users map types they are familiar with to the new PutPayload.

Here is a suggestion: tustvold#81

Other than that I had a few suggestions / nits, the only one I feel somewhat strongly about is adding a new put rather than updating the existing write API.

alamb · 2024-04-11T15:43:27Z

object_store/src/aws/client.rs

@@ -266,7 +270,8 @@ pub(crate) struct Request<'a> {
    path: &'a Path,
    config: &'a S3Config,
    builder: RequestBuilder,
-    payload_sha256: Option<Vec<u8>>,
+    payload_sha256: Option<digest::Digest>,


These are unrelated cleanups right? Or is there a reason digest::Digest is preferred over Vec<u8>?

It avoids an allocation, but yes an unrelated cleanup

alamb · 2024-04-11T15:47:09Z

object_store/src/aws/client.rs

@@ -438,16 +443,8 @@ impl S3Client {

        let mut builder = self.client.request(Method::POST, url);

-        // Compute checksum - S3 *requires* this for DeleteObjects requests, so we default to


(and to be clear for anyone else reading this PR -- this code is not exposed publicly so it can be changed in the future if needed as well)

alamb · 2024-04-11T15:51:12Z

object_store/src/lib.rs

    /// should be able to observe a partially written object
-    async fn put(&self, location: &Path, bytes: Bytes) -> Result<PutResult> {
-        self.put_opts(location, bytes, PutOptions::default()).await
+    async fn put(&self, location: &Path, payload: PutPayload) -> Result<PutResult> {


For anyone else reviewing this PR. I think this is an (and the key) API change -- to take a slice of [Bytes] bytes rather than a Bytes directly

alamb · 2024-04-11T16:01:29Z

object_store/src/payload.rs

+    }
+
+    /// Append a [`Bytes`] to this [`PutPayloadMut`]
+    pub fn push(&mut self, bytes: Bytes) {


Can you please document that this is a zero copy API -- specifically that the progress is closed off and the underlying data is not copied.

alamb · 2024-04-11T16:03:25Z

object_store/src/payload.rs

+        if !self.in_progress.is_empty() {
+            let completed = std::mem::take(&mut self.in_progress).into();
+            self.completed.push(completed);
+        }
+        PutPayload(self.completed.into())


nit: I think you can save typing and allocating another Vec with something like

Suggested change

if !self.in_progress.is_empty() {

let completed = std::mem::take(&mut self.in_progress).into();

self.completed.push(completed);

}

PutPayload(self.completed.into())

let Self { in_progress, completed, ..} = self;

if !in_progress.is_empty() {

completed.push(in_progress);

}

PutPayload(completed.into())

Vec::default doesn't allocate

alamb · 2024-04-11T16:04:39Z

object_store/src/payload.rs

+        assert_eq!(chunks[3].len(), 23);
+        assert_eq!(chunks[4].len(), 20);
+    }
+}


It would be good to add a test for push as well (esp to document the non copying behavior)

alamb · 2024-04-11T16:08:18Z

object_store/src/upload.rs

+    /// Put a chunk of data into this [`WriteMultipart`]
+    ///
+    /// See [`Self::write`] for information on backpressure
+    pub fn put(&mut self, mut bytes: Bytes) {


If we are changing the API anyways, is there a reason to keep both write and put? As in why not change write to take a Bytes as one can could then call write via write(my_payload.into()) 🤔

There are circumstances where having the writer perform the buffering might be advantageous, e.g. when integrating with the Write trait

this would be great to include in the docs as I found it non obvious.

Update the docs were updated (thanks 🙏 )

/// Put a chunk of data into this [`WriteMultipart`] without copying /// /// Data is buffered using [`PutPayloadMut::push`]. Implementations looking to /// perform writes from non-owned buffers should prefer [`Self::write`] as this /// will allow multiple calls to share the same underlying allocation.

alamb · 2024-04-11T16:11:52Z

object_store/src/client/retry.rs

+                .expect("request body must be cloneable");
+
+            if let Some(x) = &self.payload {
+                *s.body_mut() = Some(x.body());


Maybe a function call like s.set_body(Some(x.body()) instead might be easier to understand the intent

alamb · 2024-04-11T16:12:38Z

object_store/src/client/retry.rs

+}
+
+impl RetryableRequest {
+    /// Set whether this request is idempotent


Can we document what idempotent means in this context? Like it means the entire request can be retried?

alamb · 2024-04-11T16:16:02Z

object_store/src/client/retry.rs

        loop {
-            let s = req.try_clone().expect("request body must be cloneable");
-            match client.execute(s).await {
+            let mut s = self


I know you like brevity but one character variable names that are different than the first letter of the things they are holding are pretty hard for me to grok. (I realize you can't use r b/c there is another variable below named r).

How about we rename s to request and x to payload?

github-actions bot added the object-store Object Store Interface label Mar 21, 2024

tustvold commented Mar 21, 2024

View reviewed changes

tustvold force-pushed the non-contiguous-write-payloads branch 4 times, most recently from 3640a8c to 4e20f2d Compare April 9, 2024 13:27

Support non-contiguous put payloads (apache#5514)

70de4ee

tustvold force-pushed the non-contiguous-write-payloads branch from 4e20f2d to 70de4ee Compare April 9, 2024 15:02

tustvold commented Apr 9, 2024

View reviewed changes

Docs

4f9c5c7

tustvold commented Apr 9, 2024

View reviewed changes

object_store/src/payload.rs Outdated Show resolved Hide resolved

tustvold marked this pull request as ready for review April 9, 2024 15:07

Add more docs

ff4618a

tustvold changed the title ~~Support non-contiguous put payloads (#5514)~~ Support non-contiguous put payloads / vectored writes (#5514) Apr 9, 2024

Xuanwo reviewed Apr 9, 2024

View reviewed changes

alamb mentioned this pull request Apr 9, 2024

DataFusion weekly project plan (Andrew Lamb) - April 8, 2024 apache/datafusion#10002

Closed

9 tasks

alamb added the api-change Changes to the arrow API label Apr 11, 2024

alamb mentioned this pull request Apr 11, 2024

Additional docs and examples for PutPayload and PutPayloadMut tustvold/arrow-rs#81

Open

alamb approved these changes Apr 11, 2024

View reviewed changes

Review feedback

d238149

tustvold merged commit 8731760 into apache:master Apr 15, 2024
13 checks passed

tustvold mentioned this pull request Apr 15, 2024

Release Object Store 0.10.0 #5647

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support non-contiguous put payloads / vectored writes (#5514) #5538

Support non-contiguous put payloads / vectored writes (#5514) #5538

tustvold commented Mar 21, 2024 •

edited

Loading

tustvold Mar 21, 2024

alamb Apr 11, 2024

tustvold Mar 21, 2024

tustvold Apr 9, 2024

tustvold Apr 9, 2024

tustvold Apr 9, 2024

tustvold Apr 9, 2024

alamb Apr 11, 2024

tustvold Apr 15, 2024

tustvold Apr 9, 2024

tustvold Apr 9, 2024

tustvold Apr 9, 2024

Xuanwo Apr 9, 2024

tustvold Apr 9, 2024

Xuanwo Apr 10, 2024

tustvold Apr 10, 2024 •

edited

Loading

Xuanwo Apr 10, 2024

alamb left a comment

alamb Apr 11, 2024

tustvold Apr 15, 2024

alamb Apr 11, 2024

alamb Apr 11, 2024

alamb Apr 11, 2024

alamb Apr 11, 2024

tustvold Apr 15, 2024

alamb Apr 11, 2024

alamb Apr 11, 2024

tustvold Apr 15, 2024

alamb Apr 16, 2024 •

edited

Loading

alamb Apr 11, 2024

alamb Apr 11, 2024

alamb Apr 11, 2024

		@@ -438,16 +443,8 @@ impl S3Client {

		let mut builder = self.client.request(Method::POST, url);

		// Compute checksum - S3 requires this for DeleteObjects requests, so we default to

Support non-contiguous put payloads / vectored writes (#5514) #5538

Support non-contiguous put payloads / vectored writes (#5514) #5538

Conversation

tustvold commented Mar 21, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tustvold Apr 10, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb Apr 16, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tustvold commented Mar 21, 2024 •

edited

Loading

tustvold Apr 10, 2024 •

edited

Loading

alamb Apr 16, 2024 •

edited

Loading