This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24925-e9da08d89fed6380047109cce60ea3ecc4cc30f7 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit a7ea1201bccb5cd9807039f31cf5a883d815e842 Author: Pepijn Van Eeckhoudt <[email protected]> AuthorDate: Thu Sep 10 09:14:00 2026 +0000 Add write support to Avro format (#24925) ## Which issue does this PR close? - Closes #7679. ## Rationale for this change The avro integration was still missing the necessary glue code to support writing. This MR adds the missing bits. ## What changes are included in this PR? - Adds a write physical plan in the avro file format This MR was authored using Claude Code and subsequently reviewed by myself. ## What is the testing strategy for this PR? - Added SLTs to verify Avro roundtripping - Manual testing ## Are there any user-facing changes? Yes, `COPY ... TO ... STORED AS AVRO` now works. --- Cargo.lock | 5 + datafusion/datasource-avro/Cargo.toml | 5 + datafusion/datasource-avro/src/file_format.rs | 197 ++++++++++++++++++++++- datafusion/sqllogictest/test_files/copy_avro.slt | 166 +++++++++++++++++++ 4 files changed, 368 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 18fde39a95..6d4c38d8d9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1994,13 +1994,18 @@ dependencies = [ "async-trait", "bytes", "datafusion-common", + "datafusion-common-runtime", "datafusion-datasource", + "datafusion-execution", + "datafusion-expr", "datafusion-physical-expr-adapter", + "datafusion-physical-expr-common", "datafusion-physical-plan", "datafusion-proto-models", "datafusion-session", "futures", "object_store", + "tokio", ] [[package]] diff --git a/datafusion/datasource-avro/Cargo.toml b/datafusion/datasource-avro/Cargo.toml index 70b675d63f..63e949e21d 100644 --- a/datafusion/datasource-avro/Cargo.toml +++ b/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] diff --git a/datafusion/datasource-avro/src/file_format.rs b/datafusion/datasource-avro/src/file_format.rs index 9d9d3279c0..93dad800c6 100644 --- a/datafusion/datasource-avro/src/file_format.rs +++ b/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,6 +171,22 @@ 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, @@ -154,3 +194,150 @@ impl FileFormat for AvroFormat { 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 + } +} diff --git a/datafusion/sqllogictest/test_files/copy_avro.slt b/datafusion/sqllogictest/test_files/copy_avro.slt new file mode 100644 index 0000000000..6fa23fd3e2 --- /dev/null +++ b/datafusion/sqllogictest/test_files/copy_avro.slt @@ -0,0 +1,166 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# tests for `COPY ... TO ... STORED AS AVRO` + +statement ok +create table source_table(col1 integer, col2 varchar) as values (1, 'Foo'), (2, 'Bar'); + +# Copy from table to a single Avro file +query I +COPY source_table TO 'test_files/scratch/copy_avro/table.avro' STORED AS AVRO; +---- +2 + +# Validate single Avro file output +statement ok +CREATE EXTERNAL TABLE validate_avro_file +STORED AS AVRO +LOCATION 'test_files/scratch/copy_avro/table.avro'; + +query IT +select * from validate_avro_file; +---- +1 Foo +2 Bar + +# Copy from table to a directory of (multiple) Avro files +query I +COPY (select * from source_table UNION ALL select * from source_table) to 'test_files/scratch/copy_avro/table_dir' STORED AS AVRO; +---- +4 + +# Validate multiple Avro file output +statement ok +CREATE EXTERNAL TABLE validate_avro_dir STORED AS AVRO LOCATION 'test_files/scratch/copy_avro/table_dir'; + +query IT +select * from validate_avro_dir order by col1, col2; +---- +1 Foo +1 Foo +2 Bar +2 Bar + +# Copy 0 rows to a single Avro file output +query I +COPY (SELECT 1 AS id WHERE FALSE) TO 'test_files/scratch/copy_avro/table_no_rows.avro' STORED AS AVRO; +---- +0 + +statement ok +CREATE EXTERNAL TABLE validate_avro_no_rows STORED AS AVRO LOCATION 'test_files/scratch/copy_avro/table_no_rows.avro'; + +# validate the Avro file contains 0 rows. +query I +SELECT count(id) FROM validate_avro_no_rows; +---- +0 + +# Copy to directory as partitioned files +query I +COPY (values (1, 'a', 'x'), (2, 'b', 'y'), (3, 'c', 'z')) TO 'test_files/scratch/copy_avro/partitioned_table/' STORED AS AVRO PARTITIONED BY (column2, column3); +---- +3 + +# validate multiple partitioned Avro file output +statement ok +CREATE EXTERNAL TABLE validate_partitioned_avro STORED AS AVRO +LOCATION 'test_files/scratch/copy_avro/partitioned_table/' PARTITIONED BY (column2, column3); + +query ITT +select * from validate_partitioned_avro order by column1, column2, column3; +---- +1 a x +2 b y +3 c z + +statement ok +CREATE EXTERNAL TABLE validate_partitioned_avro_a_x STORED AS AVRO +LOCATION 'test_files/scratch/copy_avro/partitioned_table/column2=a/column3=x'; + +query I +select * from validate_partitioned_avro_a_x order by column1; +---- +1 + +# Copy a variety of data types, including nested/complex types (list, struct, +# map), to a single Avro file and validate the round trip. The schema is +# inferred from the self-describing Avro file rather than declared explicitly. +query I +COPY (values + (arrow_cast(1, 'Int32'), arrow_cast(2, 'Int64'), arrow_cast(3.1, 'Float64'), 19968::date, true, 'x', + make_array(1, 2, 3), named_struct('a', 1, 'b', 'foo'), MAP(['k1', 'k2'], [10, 20])), + (arrow_cast(11, 'Int32'), arrow_cast(22, 'Int64'), arrow_cast(3.2, 'Float64'), 19969::date, false, 'y', + make_array(4, 5), named_struct('a', 2, 'b', 'bar'), MAP(['k3'], [30])) +) TO 'test_files/scratch/copy_avro/table_types.avro' STORED AS AVRO; +---- +2 + +statement ok +CREATE EXTERNAL TABLE validate_avro_types STORED AS AVRO LOCATION 'test_files/scratch/copy_avro/table_types.avro'; + +query IIRDBT??? +select * from validate_avro_types order by column1; +---- +1 2 3.1 2024-09-02 true x [1, 2, 3] {a: 1, b: foo} {k1: 10, k2: 20} +11 22 3.2 2024-09-03 false y [4, 5] {a: 2, b: bar} {k3: 30} + +# EXPLAIN output for COPY ... STORED AS AVRO +query TT +EXPLAIN COPY source_table TO 'test_files/scratch/copy_avro/table_explain.avro' STORED AS AVRO; +---- +logical_plan +01)CopyTo: format=avro output_url=test_files/scratch/copy_avro/table_explain.avro options: () +02)--TableScan: source_table projection=[col1, col2] +physical_plan +01)DataSinkExec: sink=AvroFileSink(file_groups=[]) +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +# EXPLAIN ANALYZE COPY ... STORED AS AVRO. Note the AvroFileSink does not +# currently report rows_written/bytes_written/elapsed_compute metrics. +query TT +EXPLAIN ANALYZE COPY (SELECT col1, upper(col2) AS col2_upper FROM source_table ORDER BY col1) TO 'test_files/scratch/copy_avro/table_metrics.avro' STORED AS AVRO; +---- +Plan with Metrics +01)DataSinkExec: sink=AvroFileSink(file_groups=[]), metrics=[] +02)--SortExec: expr=[col1@0 ASC NULLS LAST], preserve_partitioning=[false], metrics=[output_rows=2, elapsed_compute=<slt:ignore>, output_bytes=<slt:ignore>, output_batches=<slt:ignore>, spill_count=0, spilled_bytes=0.0 B, spilled_rows=0] +03)----ProjectionExec: expr=[col1@0 as col1, upper(col2@1) as col2_upper], metrics=[output_rows=2, elapsed_compute=<slt:ignore>, output_bytes=<slt:ignore>, output_batches=1, expr_0_eval_time=<slt:ignore>, expr_1_eval_time=<slt:ignore>] +04)------DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] + +# Error case: Format not explicitly set and unable to infer extension from output path +query error DataFusion error: Invalid or Unsupported Configuration: Format not explicitly set and unable to get file extension! Use STORED AS to define file format. +EXPLAIN COPY source_table to 'test_files/scratch/copy_avro/no_format_specified'; + +# The following test writes sufficient data to exceed `datasource::file_format::avro::BUFFER_FLUSH_BYTES` +# so that we excercise the intermediate flush code path. +statement ok +CREATE EXTERNAL TABLE repeat_much STORED AS PARQUET LOCATION 'data/repeat_much.snappy.parquet'; + +query I +COPY repeat_much TO 'test_files/scratch/copy_avro/repeat_much.avro' STORED AS AVRO; +---- +1957500 + +statement ok +CREATE EXTERNAL TABLE repeat_much_copy STORED AS AVRO LOCATION 'test_files/scratch/copy_avro/repeat_much.avro'; + +query I +SELECT sum(a) FROM repeat_much UNION ALL SELECT sum(a) FROM repeat_much_copy; +---- +4798655628750 +4798655628750 \ No newline at end of file --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
