tustvold commented on code in PR #5205:
URL: https://github.com/apache/arrow-rs/pull/5205#discussion_r1426967218
##########
object_store/src/multipart.rs:
##########
@@ -316,3 +317,136 @@ pub trait MultiPartStore: Send + Sync + 'static {
/// Aborts a multipart upload
async fn abort_multipart(&self, path: &Path, id: &MultipartId) ->
Result<()>;
}
+
+/// Create a lazy multipart writer for a given [`ObjectStore`] and [`Path`].
+pub fn put_multipart_lazy(
+ store: Arc<dyn ObjectStore>,
+ path: Path,
+) -> Box<dyn AsyncWrite + Send + Unpin> {
+ Box::new(LazyWriteMultiPart::new(store, path))
+}
+
+/// Wrapper around a [`ObjectStore`] and [`Path`] that implements
[`AsyncWrite`]
+///
+/// A multipart upload using `ObjectStore::put_multipart` will only be created
if the size exceeds 10 MB,
+/// otherwise a direct PUT will be performed on shutdown.
+pub struct LazyWriteMultiPart {
+ store: Arc<dyn ObjectStore>,
+ path: Path,
+ part_size: usize,
+ multipart_writer: Option<Box<dyn AsyncWrite + Send + Unpin>>,
+ buffer: Vec<u8>,
+ create_task: Option<BoxedTryFuture<Box<dyn AsyncWrite + Send + Unpin>>>,
+ put_task: Option<BoxedTryFuture<()>>,
+}
+
+impl LazyWriteMultiPart {
+ /// Create a new lazy multipart upload.
+ pub fn new(store: Arc<dyn ObjectStore>, path: Path) -> Self {
+ Self {
+ store,
+ path,
+ part_size: 10 * 1024 * 1024,
+ multipart_writer: None,
+ buffer: Vec::new(),
+ create_task: None,
+ put_task: None,
+ }
+ }
+
+ fn do_flush(
+ mut self: Pin<&mut Self>,
+ cx: &mut std::task::Context<'_>,
+ ) -> Poll<Result<(), io::Error>> {
+ let buffer = std::mem::take(&mut self.buffer);
Review Comment:
Perhaps this could be moved into the conditional block below?
##########
object_store/src/multipart.rs:
##########
@@ -316,3 +317,136 @@ pub trait MultiPartStore: Send + Sync + 'static {
/// Aborts a multipart upload
async fn abort_multipart(&self, path: &Path, id: &MultipartId) ->
Result<()>;
}
+
+/// Create a lazy multipart writer for a given [`ObjectStore`] and [`Path`].
+pub fn put_multipart_lazy(
+ store: Arc<dyn ObjectStore>,
+ path: Path,
+) -> Box<dyn AsyncWrite + Send + Unpin> {
+ Box::new(LazyWriteMultiPart::new(store, path))
+}
+
+/// Wrapper around a [`ObjectStore`] and [`Path`] that implements
[`AsyncWrite`]
+///
+/// A multipart upload using `ObjectStore::put_multipart` will only be created
if the size exceeds 10 MB,
+/// otherwise a direct PUT will be performed on shutdown.
+pub struct LazyWriteMultiPart {
+ store: Arc<dyn ObjectStore>,
+ path: Path,
+ part_size: usize,
+ multipart_writer: Option<Box<dyn AsyncWrite + Send + Unpin>>,
+ buffer: Vec<u8>,
+ create_task: Option<BoxedTryFuture<Box<dyn AsyncWrite + Send + Unpin>>>,
+ put_task: Option<BoxedTryFuture<()>>,
Review Comment:
This might be easier to follow as an enumeration, e.g.
```
enum LazyWriteState {
Buffer(Arc<dyn ObjectStore>, Path, Vec<u8>),
Put(BoxFuture<'static, Result<()>>),
PutMultiPart(BoxFuture<'static, Result<Box<dyn AsyncWrite + Send +
Unpin>>),
AsyncWrite(Box<dyn AsyncWrite + Send + Unpin>)
Error,
}
```
##########
object_store/src/multipart.rs:
##########
@@ -316,3 +317,136 @@ pub trait MultiPartStore: Send + Sync + 'static {
/// Aborts a multipart upload
async fn abort_multipart(&self, path: &Path, id: &MultipartId) ->
Result<()>;
}
+
+/// Create a lazy multipart writer for a given [`ObjectStore`] and [`Path`].
+pub fn put_multipart_lazy(
+ store: Arc<dyn ObjectStore>,
+ path: Path,
+) -> Box<dyn AsyncWrite + Send + Unpin> {
+ Box::new(LazyWriteMultiPart::new(store, path))
+}
+
+/// Wrapper around a [`ObjectStore`] and [`Path`] that implements
[`AsyncWrite`]
+///
+/// A multipart upload using `ObjectStore::put_multipart` will only be created
if the size exceeds 10 MB,
+/// otherwise a direct PUT will be performed on shutdown.
+pub struct LazyWriteMultiPart {
+ store: Arc<dyn ObjectStore>,
+ path: Path,
+ part_size: usize,
+ multipart_writer: Option<Box<dyn AsyncWrite + Send + Unpin>>,
+ buffer: Vec<u8>,
+ create_task: Option<BoxedTryFuture<Box<dyn AsyncWrite + Send + Unpin>>>,
+ put_task: Option<BoxedTryFuture<()>>,
+}
+
+impl LazyWriteMultiPart {
+ /// Create a new lazy multipart upload.
+ pub fn new(store: Arc<dyn ObjectStore>, path: Path) -> Self {
+ Self {
+ store,
+ path,
+ part_size: 10 * 1024 * 1024,
+ multipart_writer: None,
+ buffer: Vec::new(),
+ create_task: None,
+ put_task: None,
+ }
+ }
+
+ fn do_flush(
+ mut self: Pin<&mut Self>,
+ cx: &mut std::task::Context<'_>,
+ ) -> Poll<Result<(), io::Error>> {
+ let buffer = std::mem::take(&mut self.buffer);
+ if let Some(multipart_writer) = self.multipart_writer.as_mut() {
+ if !buffer.is_empty() {
+ Pin::new(multipart_writer)
+ .poll_write(cx, &buffer)
+ .map(|_| Ok(()))
+ } else {
+ Pin::new(multipart_writer).poll_flush(cx)
+ }
+ } else {
+ self.buffer = buffer;
+ Poll::Ready(Ok(()))
+ }
+ }
+}
+
+impl AsyncWrite for LazyWriteMultiPart {
+ fn poll_write(
+ mut self: Pin<&mut Self>,
+ cx: &mut std::task::Context<'_>,
+ buf: &[u8],
+ ) -> Poll<Result<usize, io::Error>> {
+ if let Some(multipart_writer) = self.multipart_writer.as_mut() {
+ Pin::new(multipart_writer).poll_write(cx, buf)
+ } else {
+ let buf_len = buf.len();
+ let new_len = self.buffer.len() + buf_len;
+ if new_len > self.part_size {
+ let inner = Arc::clone(&self.store);
+ let path = self.path.clone();
+ let create_task = self.create_task.get_or_insert_with(|| {
+ Box::pin(async move {
+ let (_, multipart_writer) =
inner.put_multipart(&path).await?;
+ Ok(multipart_writer)
+ })
+ });
+ let multipart_writer = ready!(Pin::new(create_task).poll(cx))?;
+ self.multipart_writer = Some(multipart_writer);
+
+ let to_copy = std::cmp::min(buf_len, self.buffer.len());
+ let mut write_buf = std::mem::take(&mut self.buffer);
+ self.buffer = write_buf.split_off(to_copy);
+ self.buffer.extend_from_slice(buf);
+
+
Pin::new(self.multipart_writer.as_mut().unwrap()).poll_write(cx, &write_buf)
Review Comment:
I think this is incorrect if the inner poll_write doesn't accept the entire
write_buf? I think we need to keep track of where in self.buffer we have
flushed from, and loop with this. This would also avoid the copy above
##########
object_store/src/multipart.rs:
##########
@@ -316,3 +317,136 @@ pub trait MultiPartStore: Send + Sync + 'static {
/// Aborts a multipart upload
async fn abort_multipart(&self, path: &Path, id: &MultipartId) ->
Result<()>;
}
+
+/// Create a lazy multipart writer for a given [`ObjectStore`] and [`Path`].
+pub fn put_multipart_lazy(
Review Comment:
Should the point at which it switches to multipart be configurable?
##########
object_store/src/multipart.rs:
##########
@@ -316,3 +317,136 @@ pub trait MultiPartStore: Send + Sync + 'static {
/// Aborts a multipart upload
async fn abort_multipart(&self, path: &Path, id: &MultipartId) ->
Result<()>;
}
+
+/// Create a lazy multipart writer for a given [`ObjectStore`] and [`Path`].
+pub fn put_multipart_lazy(
+ store: Arc<dyn ObjectStore>,
+ path: Path,
+) -> Box<dyn AsyncWrite + Send + Unpin> {
+ Box::new(LazyWriteMultiPart::new(store, path))
+}
+
+/// Wrapper around a [`ObjectStore`] and [`Path`] that implements
[`AsyncWrite`]
+///
+/// A multipart upload using `ObjectStore::put_multipart` will only be created
if the size exceeds 10 MB,
+/// otherwise a direct PUT will be performed on shutdown.
+pub struct LazyWriteMultiPart {
Review Comment:
```suggestion
struct LazyWriteMultiPart {
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
