Add write support to Avro format

datafusion/datasource-avro/Cargo.toml

@@ -45,13 +45,18 @@ arrow-avro = { workspace = true }
 async-trait = { workspace = true }
 bytes = { workspace = true }
 datafusion-common = { workspace = true, features = ["object_store"] }
+datafusion-common-runtime = { workspace = true }
 datafusion-datasource = { workspace = true }
+datafusion-execution = { workspace = true }
+datafusion-expr = { workspace = true }
 datafusion-physical-expr-adapter = { workspace = true }
+datafusion-physical-expr-common = { workspace = true }
 datafusion-physical-plan = { workspace = true }
 datafusion-proto-models = { workspace = true, optional = true }
 datafusion-session = { workspace = true }
 futures = { workspace = true }
 object_store = { workspace = true }
+tokio = { workspace = true }
 
 [dev-dependencies]
 

datafusion/datasource-avro/src/file_format.rs

@@ -17,29 +17,53 @@
 
 //! Apache Avro [`FileFormat`] abstractions
 use std::collections::HashMap;
-use std::fmt;
+use std::fmt::Debug;
 use std::sync::Arc;
+use std::{fmt, io};
 
 use crate::read_avro_schema_from_reader;
 use crate::source::AvroSource;
 
 use arrow::datatypes::Schema;
 use arrow::datatypes::SchemaRef;
-use datafusion_common::DEFAULT_AVRO_EXTENSION;
+use arrow_avro::errors::AvroError;
+use arrow_avro::writer::format::AvroOcfFormat;
+use arrow_avro::writer::{AvroWriter, WriterBuilder};
 use datafusion_common::GetExt;
 use datafusion_common::internal_err;
+use datafusion_common::not_impl_err;
 use datafusion_common::parsers::CompressionTypeVariant;
-use datafusion_common::{Result, Statistics};
+use datafusion_common::{DEFAULT_AVRO_EXTENSION, Diagnostic};
+use datafusion_common::{DataFusionError, Result, Statistics, internal_datafusion_err};
+use datafusion_common_runtime::{JoinSet, SpawnedTask};
+use datafusion_datasource::display::FileGroupDisplay;
 use datafusion_datasource::file::FileSource;
 use datafusion_datasource::file_compression_type::FileCompressionType;
 use datafusion_datasource::file_format::{FileFormat, FileFormatFactory};
 use datafusion_datasource::file_scan_config::FileScanConfig;
+use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig};
+use datafusion_datasource::sink::{DataSink, DataSinkExec};
 use datafusion_datasource::source::DataSourceExec;
-use datafusion_physical_plan::ExecutionPlan;
+use datafusion_datasource::write::demux::DemuxedStreamReceiver;
+use datafusion_datasource::write::{
+    ObjectWriterBuilder, SharedBuffer, get_writer_schema,
+};
+use datafusion_execution::{SendableRecordBatchStream, TaskContext};
+use datafusion_expr::dml::InsertOp;
+use datafusion_physical_expr_common::sort_expr::LexRequirement;
+use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan};
 use datafusion_session::Session;
 
 use async_trait::async_trait;
 use object_store::{GetResultPayload, ObjectMeta, ObjectStore, ObjectStoreExt};
+use tokio::io::AsyncWriteExt;
+
+/// Initial writing buffer size. Note this is just a size hint for efficiency. It
+/// will grow beyond the set value if needed.
+const INITIAL_BUFFER_BYTES: usize = 1048576;
+
+/// If the buffered Avro data exceeds this size, it is flushed to object store
+const BUFFER_FLUSH_BYTES: usize = 1024000;
 
 #[derive(Default)]
 /// Factory struct used to create [`AvroFormat`]
@@ -66,7 +90,7 @@ impl FileFormatFactory for AvroFormatFactory {
     }
 }
 
-impl fmt::Debug for AvroFormatFactory {
+impl Debug for AvroFormatFactory {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("AvroFormatFactory").finish()
     }
@@ -147,10 +171,173 @@ impl FileFormat for AvroFormat {
         Ok(DataSourceExec::from_data_source(conf))
     }
 
+    async fn create_writer_physical_plan(
+        &self,
+        input: Arc<dyn ExecutionPlan>,
+        _state: &dyn Session,
+        conf: FileSinkConfig,
+        order_requirements: Option<LexRequirement>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        if conf.insert_op != InsertOp::Append {
+            return not_impl_err!("Overwrites are not implemented yet for Avro format");
+        }
+
+        let sink = Arc::new(AvroFileSink::new(conf));
+
+        Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _)
+    }
+
     fn file_source(
         &self,
         table_schema: datafusion_datasource::TableSchema,
     ) -> Arc<dyn FileSource> {
         Arc::new(AvroSource::new(table_schema))
     }
 }
+
+/// Implements [`FileSink`] for Avro Object Container Files
+struct AvroFileSink {
+    config: FileSinkConfig,
+}
+
+impl AvroFileSink {
+    fn new(config: FileSinkConfig) -> Self {
+        Self { config }
+    }
+}
+
+#[async_trait]
+impl FileSink for AvroFileSink {
+    fn config(&self) -> &FileSinkConfig {
+        &self.config
+    }
+
+    async fn spawn_writer_tasks_and_join(
+        &self,
+        context: &Arc<TaskContext>,
+        demux_task: SpawnedTask<Result<()>>,
+        mut file_stream_rx: DemuxedStreamReceiver,
+        object_store: Arc<dyn ObjectStore>,
+    ) -> Result<u64> {
+        let mut file_write_tasks: JoinSet<std::result::Result<usize, DataFusionError>> =
+            JoinSet::new();
+
+        let writer_schema = get_writer_schema(&self.config);
+        while let Some((path, mut rx)) = file_stream_rx.recv().await {
+            let shared_buffer = SharedBuffer::new(INITIAL_BUFFER_BYTES);
+            let mut avro_writer: AvroWriter<SharedBuffer> =
+                WriterBuilder::new(writer_schema.as_ref().clone())
+                    .build::<_, AvroOcfFormat>(shared_buffer.clone())
+                    .map_err(|e| {
+                        internal_datafusion_err!("Failed to create Avro writer: {e}")
+                    })?;
+            let mut object_store_writer = ObjectWriterBuilder::new(
+                FileCompressionType::UNCOMPRESSED,
+                &path,
+                Arc::clone(&object_store),
+            )
+            .with_buffer_size(Some(
+                context
+                    .session_config()
+                    .options()
+                    .execution
+                    .objectstore_writer_buffer_size,
+            ))
+            .build()?;
+            file_write_tasks.spawn(async move {
+                let mut row_count = 0;
+                while let Some(batch) = rx.recv().await {
+                    row_count += batch.num_rows();
+                    avro_writer
+                        .write(&batch)
+                        .map_err(|e| internal_datafusion_err!("{e}"))?;
+                    let mut buff_to_flush = shared_buffer.buffer.try_lock().unwrap();
+                    if buff_to_flush.len() > BUFFER_FLUSH_BYTES {
+                        object_store_writer
+                            .write_all(buff_to_flush.as_slice())
+                            .await?;
+                        buff_to_flush.clear();
+                    }
+                }
+                if let Err(e) = avro_writer.finish() {
+                    return Err(match e {
+                        AvroError::NYI(e) => DataFusionError::NotImplemented(e),
+                        AvroError::EOF(e) => DataFusionError::IoError(io::Error::new(
+                            io::ErrorKind::UnexpectedEof,
+                            e,
+                        )),
+                        AvroError::ArrowError(e) => DataFusionError::ArrowError(e, None),
+                        AvroError::External(e) => DataFusionError::External(e),
+                        AvroError::IoError(msg, e) => DataFusionError::IoError(e)
+                            .with_diagnostic(Diagnostic::new_error(msg, None)),
+                        _ => internal_datafusion_err!("{e}"),
+                    });
+                }
+                let final_buff = shared_buffer.buffer.try_lock().unwrap();
+
+                object_store_writer.write_all(final_buff.as_slice()).await?;
+                object_store_writer.shutdown().await?;
+                Ok(row_count)
+            });
+        }
+
+        let mut row_count = 0;
+        while let Some(result) = file_write_tasks.join_next().await {
+            match result {
+                Ok(r) => {
+                    row_count += r?;
+                }
+                Err(e) => {
+                    if e.is_panic() {
+                        std::panic::resume_unwind(e.into_panic());
+                    } else {
+                        unreachable!();
+                    }
+                }
+            }
+        }
+
+        demux_task
+            .join_unwind()
+            .await
+            .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??;
+        Ok(row_count as u64)
+    }
+}
+
+impl Debug for AvroFileSink {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("AvroFileSink").finish()
+    }
+}
+
+impl DisplayAs for AvroFileSink {
+    fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match t {
+            DisplayFormatType::Default | DisplayFormatType::Verbose => {
+                write!(f, "AvroFileSink(file_groups=")?;
+                FileGroupDisplay(&self.config.file_group).fmt_as(t, f)?;
+                write!(f, ")")
+            }
+            DisplayFormatType::TreeRender => {
+                writeln!(f, "format: avro")?;
+                write!(f, "file={}", self.config.original_url)
+            }
+        }
+    }
+}
+
+#[async_trait]
+impl DataSink for AvroFileSink {
+    fn schema(&self) -> &SchemaRef {
+        self.config.output_schema()
+    }
+
+    async fn write_all(
+        &self,
+        data: SendableRecordBatchStream,
+        context: &Arc<TaskContext>,
+    ) -> Result<u64> {
+        FileSink::write_all(self, data, context).await
+    }
+}