Use object_store:BufWriter to replace put_multipart

Cargo.toml

@@ -93,7 +93,7 @@ indexmap = "2.0.0"
 itertools = "0.12"
 log = "^0.4"
 num_cpus = "1.13.0"
-object_store = { version = "0.9.0", default-features = false }
+object_store = { version = "0.9.1", default-features = false }
 parking_lot = "0.12"
 parquet = { version = "50.0.0", default-features = false, features = ["arrow", "async", "object_store"] }
 rand = "0.8"

datafusion/core/src/datasource/file_format/file_compression_type.rs

@@ -43,6 +43,7 @@ use futures::stream::BoxStream;
 use futures::StreamExt;
 #[cfg(feature = "compression")]
 use futures::TryStreamExt;
+use object_store::buffered::BufWriter;
 use tokio::io::AsyncWrite;
 #[cfg(feature = "compression")]
 use tokio_util::io::{ReaderStream, StreamReader};
@@ -152,7 +153,7 @@ impl FileCompressionType {
     /// according to this `FileCompressionType`.
     pub fn convert_async_writer(
         &self,
-        w: Box<dyn AsyncWrite + Send + Unpin>,
+        w: BufWriter,
     ) -> Result<Box<dyn AsyncWrite + Send + Unpin>> {
         Ok(match self.variant {
             #[cfg(feature = "compression")]
@@ -169,7 +170,7 @@ impl FileCompressionType {
                     "Compression feature is not enabled".to_owned(),
                 ))
             }
-            UNCOMPRESSED => w,
+            UNCOMPRESSED => Box::new(w),
         })
     }
 

datafusion/core/src/datasource/file_format/parquet.rs

@@ -23,7 +23,7 @@ use std::fmt::Debug;
 use std::sync::Arc;
 
 use super::write::demux::start_demuxer_task;
-use super::write::{create_writer, AbortableWrite, SharedBuffer};
+use super::write::{create_writer, SharedBuffer};
 use super::{FileFormat, FileScanConfig};
 use crate::arrow::array::{
     BooleanArray, Float32Array, Float64Array, Int32Array, Int64Array, RecordBatch,
@@ -56,6 +56,7 @@ use datafusion_physical_plan::metrics::MetricsSet;
 
 use async_trait::async_trait;
 use bytes::{BufMut, BytesMut};
+use object_store::buffered::BufWriter;
 use parquet::arrow::arrow_writer::{
     compute_leaves, get_column_writers, ArrowColumnChunk, ArrowColumnWriter,
     ArrowLeafColumn,
@@ -616,20 +617,14 @@ impl ParquetSink {
         location: &Path,
         object_store: Arc<dyn ObjectStore>,
         parquet_props: WriterProperties,
-    ) -> Result<
-        AsyncArrowWriter<Box<dyn tokio::io::AsyncWrite + std::marker::Send + Unpin>>,
-    > {
-        let (_, multipart_writer) = object_store
-            .put_multipart(location)
-            .await
-            .map_err(DataFusionError::ObjectStore)?;
+    ) -> Result<AsyncArrowWriter<BufWriter>> {
+        let buf_writer = BufWriter::new(object_store, location.clone());
         let writer = AsyncArrowWriter::try_new(
-            multipart_writer,
+            buf_writer,
             self.get_writer_schema(),
             PARQUET_WRITER_BUFFER_SIZE,
             Some(parquet_props),
         )?;
-
         Ok(writer)
     }
 
@@ -947,7 +942,7 @@ async fn concatenate_parallel_row_groups(
     mut serialize_rx: Receiver<SpawnedTask<RBStreamSerializeResult>>,
     schema: Arc<Schema>,
     writer_props: Arc<WriterProperties>,
-    mut object_store_writer: AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+    mut object_store_writer: Box<dyn AsyncWrite + Send + Unpin>,
 ) -> Result<FileMetaData> {
     let merged_buff = SharedBuffer::new(INITIAL_BUFFER_BYTES);
 
@@ -989,7 +984,7 @@ async fn concatenate_parallel_row_groups(
 /// task then stitches these independent RowGroups together and streams this large
 /// single parquet file to an ObjectStore in multiple parts.
 async fn output_single_parquet_file_parallelized(
-    object_store_writer: AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+    object_store_writer: Box<dyn AsyncWrite + Send + Unpin>,
     data: Receiver<RecordBatch>,
     output_schema: Arc<Schema>,
     parquet_props: &WriterProperties,

datafusion/core/src/datasource/file_format/write/mod.rs

@@ -18,21 +18,18 @@
 //! Module containing helper methods/traits related to enabling
 //! write support for the various file formats
 
-use std::io::{Error, Write};
-use std::pin::Pin;
+use std::io::Write;
 use std::sync::Arc;
-use std::task::{Context, Poll};
 
 use crate::datasource::file_format::file_compression_type::FileCompressionType;
 use crate::error::Result;
 
 use arrow_array::RecordBatch;
-use datafusion_common::DataFusionError;
 
 use bytes::Bytes;
-use futures::future::BoxFuture;
+use object_store::buffered::BufWriter;
 use object_store::path::Path;
-use object_store::{MultipartId, ObjectStore};
+use object_store::ObjectStore;
 use tokio::io::AsyncWrite;
 
 pub(crate) mod demux;
@@ -69,79 +66,6 @@ impl Write for SharedBuffer {
     }
 }
 
-/// Stores data needed during abortion of MultiPart writers
-#[derive(Clone)]
-pub(crate) struct MultiPart {
-    /// A shared reference to the object store
-    store: Arc<dyn ObjectStore>,
-    multipart_id: MultipartId,
-    location: Path,
-}
-
-impl MultiPart {
-    /// Create a new `MultiPart`
-    pub fn new(
-        store: Arc<dyn ObjectStore>,
-        multipart_id: MultipartId,
-        location: Path,
-    ) -> Self {
-        Self {
-            store,
-            multipart_id,
-            location,
-        }
-    }
-}
-
-/// A wrapper struct with abort method and writer
-pub(crate) struct AbortableWrite<W: AsyncWrite + Unpin + Send> {
-    writer: W,
-    multipart: MultiPart,
-}
-
-impl<W: AsyncWrite + Unpin + Send> AbortableWrite<W> {
-    /// Create a new `AbortableWrite` instance with the given writer, and write mode.
-    pub(crate) fn new(writer: W, multipart: MultiPart) -> Self {
-        Self { writer, multipart }
-    }
-
-    /// handling of abort for different write modes
-    pub(crate) fn abort_writer(&self) -> Result<BoxFuture<'static, Result<()>>> {
-        let multi = self.multipart.clone();
-        Ok(Box::pin(async move {
-            multi
-                .store
-                .abort_multipart(&multi.location, &multi.multipart_id)
-                .await
-                .map_err(DataFusionError::ObjectStore)
-        }))
-    }
-}
-
-impl<W: AsyncWrite + Unpin + Send> AsyncWrite for AbortableWrite<W> {
-    fn poll_write(
-        self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-        buf: &[u8],
-    ) -> Poll<std::result::Result<usize, Error>> {
-        Pin::new(&mut self.get_mut().writer).poll_write(cx, buf)
-    }
-
-    fn poll_flush(
-        self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-    ) -> Poll<std::result::Result<(), Error>> {
-        Pin::new(&mut self.get_mut().writer).poll_flush(cx)
-    }
-
-    fn poll_shutdown(
-        self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-    ) -> Poll<std::result::Result<(), Error>> {
-        Pin::new(&mut self.get_mut().writer).poll_shutdown(cx)
-    }
-}
-
 /// A trait that defines the methods required for a RecordBatch serializer.
 pub trait BatchSerializer: Sync + Send {
     /// Asynchronously serializes a `RecordBatch` and returns the serialized bytes.
@@ -150,19 +74,13 @@ pub trait BatchSerializer: Sync + Send {
     fn serialize(&self, batch: RecordBatch, initial: bool) -> Result<Bytes>;
 }
 
-/// Returns an [`AbortableWrite`] which writes to the given object store location
+/// Returns an [`AsyncWrite`] which writes to the given object store location
 /// with the specified compression
 pub(crate) async fn create_writer(
     file_compression_type: FileCompressionType,
     location: &Path,
     object_store: Arc<dyn ObjectStore>,
-) -> Result<AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>> {
-    let (multipart_id, writer) = object_store
-        .put_multipart(location)
-        .await
-        .map_err(DataFusionError::ObjectStore)?;
-    Ok(AbortableWrite::new(
-        file_compression_type.convert_async_writer(writer)?,
-        MultiPart::new(object_store, multipart_id, location.clone()),
-    ))
+) -> Result<Box<dyn AsyncWrite + Send + Unpin>> {
+    let buf_writer = BufWriter::new(object_store, location.clone());
+    file_compression_type.convert_async_writer(buf_writer)
 }

datafusion/core/src/datasource/file_format/write/orchestration.rs

@@ -22,7 +22,7 @@
 use std::sync::Arc;
 
 use super::demux::start_demuxer_task;
-use super::{create_writer, AbortableWrite, BatchSerializer};
+use super::{create_writer, BatchSerializer};
 use crate::datasource::file_format::file_compression_type::FileCompressionType;
 use crate::datasource::physical_plan::FileSinkConfig;
 use crate::error::Result;
@@ -39,7 +39,7 @@ use tokio::io::{AsyncWrite, AsyncWriteExt};
 use tokio::sync::mpsc::{self, Receiver};
 use tokio::task::JoinSet;
 
-type WriterType = AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>;
+type WriterType = Box<dyn AsyncWrite + Send + Unpin>;
 type SerializerType = Arc<dyn BatchSerializer>;
 
 /// Serializes a single data stream in parallel and writes to an ObjectStore
@@ -49,7 +49,7 @@ type SerializerType = Arc<dyn BatchSerializer>;
 pub(crate) async fn serialize_rb_stream_to_object_store(
     mut data_rx: Receiver<RecordBatch>,
     serializer: Arc<dyn BatchSerializer>,
-    mut writer: AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+    mut writer: WriterType,
 ) -> std::result::Result<(WriterType, u64), (WriterType, DataFusionError)> {
     let (tx, mut rx) =
         mpsc::channel::<SpawnedTask<Result<(usize, Bytes), DataFusionError>>>(100);
@@ -173,19 +173,9 @@ pub(crate) async fn stateless_serialize_and_write_files(
 
     // Finalize or abort writers as appropriate
     for mut writer in finished_writers.into_iter() {
-        match any_errors {
-            true => {
-                let abort_result = writer.abort_writer();
-                if abort_result.is_err() {
-                    any_abort_errors = true;
-                }
-            }
-            false => {
-                writer.shutdown()
+        writer.shutdown()
                     .await
                     .map_err(|_| internal_datafusion_err!("Error encountered while finalizing writes! Partial results may have been written to ObjectStore!"))?;
-            }
-        }
     }
 
     if any_errors {

datafusion/core/src/datasource/physical_plan/csv.rs

@@ -44,6 +44,7 @@ use datafusion_physical_expr::{EquivalenceProperties, LexOrdering};
 
 use bytes::{Buf, Bytes};
 use futures::{ready, StreamExt, TryStreamExt};
+use object_store::buffered::BufWriter;
 use object_store::{GetOptions, GetResultPayload, ObjectStore};
 use tokio::io::AsyncWriteExt;
 use tokio::task::JoinSet;
@@ -471,7 +472,7 @@ pub async fn plan_to_csv(
 
         let mut stream = plan.execute(i, task_ctx.clone())?;
         join_set.spawn(async move {
-            let (_, mut multipart_writer) = storeref.put_multipart(&file).await?;
+            let mut buf_writer = BufWriter::new(storeref, file.clone());
             let mut buffer = Vec::with_capacity(1024);
             //only write headers on first iteration
             let mut write_headers = true;
@@ -481,15 +482,12 @@ pub async fn plan_to_csv(
                     .build(buffer);
                 writer.write(&batch)?;
                 buffer = writer.into_inner();
-                multipart_writer.write_all(&buffer).await?;
+                buf_writer.write_all(&buffer).await?;
                 buffer.clear();
                 //prevent writing headers more than once
                 write_headers = false;
             }
-            multipart_writer
-                .shutdown()
-                .await
-                .map_err(DataFusionError::from)
+            buf_writer.shutdown().await.map_err(DataFusionError::from)
         });
     }
 

datafusion/core/src/datasource/physical_plan/json.rs

@@ -43,6 +43,7 @@ use datafusion_physical_expr::{EquivalenceProperties, LexOrdering};
 
 use bytes::{Buf, Bytes};
 use futures::{ready, StreamExt, TryStreamExt};
+use object_store::buffered::BufWriter;
 use object_store::{self, GetOptions, GetResultPayload, ObjectStore};
 use tokio::io::AsyncWriteExt;
 use tokio::task::JoinSet;
@@ -338,21 +339,18 @@ pub async fn plan_to_json(
 
         let mut stream = plan.execute(i, task_ctx.clone())?;
         join_set.spawn(async move {
-            let (_, mut multipart_writer) = storeref.put_multipart(&file).await?;
+            let mut buf_writer = BufWriter::new(storeref, file.clone());
 
             let mut buffer = Vec::with_capacity(1024);
             while let Some(batch) = stream.next().await.transpose()? {
                 let mut writer = json::LineDelimitedWriter::new(buffer);
                 writer.write(&batch)?;
                 buffer = writer.into_inner();
-                multipart_writer.write_all(&buffer).await?;
+                buf_writer.write_all(&buffer).await?;
                 buffer.clear();
             }
 
-            multipart_writer
-                .shutdown()
-                .await
-                .map_err(DataFusionError::from)
+            buf_writer.shutdown().await.map_err(DataFusionError::from)
         });
     }
 

datafusion/core/src/datasource/physical_plan/parquet/mod.rs

@@ -52,6 +52,7 @@ use futures::future::BoxFuture;
 use futures::{StreamExt, TryStreamExt};
 use itertools::Itertools;
 use log::debug;
+use object_store::buffered::BufWriter;
 use object_store::path::Path;
 use object_store::ObjectStore;
 use parquet::arrow::arrow_reader::ArrowReaderOptions;
@@ -698,11 +699,11 @@ pub async fn plan_to_parquet(
         let propclone = writer_properties.clone();
 
         let storeref = store.clone();
-        let (_, multipart_writer) = storeref.put_multipart(&file).await?;
+        let buf_writer = BufWriter::new(storeref, file.clone());
         let mut stream = plan.execute(i, task_ctx.clone())?;
         join_set.spawn(async move {
             let mut writer = AsyncArrowWriter::try_new(
-                multipart_writer,
+                buf_writer,
                 plan.schema(),
                 10485760,
                 propclone,