From f9c560a5cafb255a56ced8f0ace19c7ccbb12d98 Mon Sep 17 00:00:00 2001 From: Jared Yu Date: Fri, 10 Jul 2026 21:10:21 -0700 Subject: [PATCH] Datasource: Add file_metadata() to DataSink for per-file write metadata Add FileWriteMetadata struct to datafusion-datasource along with a new default method on the DataSink trait: - file_metadata(): Post-hoc accessor returning per-file path, row count, byte size, and optional format-specific metadata bytes after write_all completes. The default implementation returns an empty Vec, so existing DataSink implementations are unaffected. ParquetSink overrides file_metadata() to expose path, row_count, and byte_size for each written file using the existing self.written HashMap that already collects ParquetMetaData during writes. The typed ParquetMetaData remains accessible via ParquetSink::written() for Rust consumers who need column-level statistics directly. DataSinkExec gains a file_metadata() convenience accessor that delegates to the underlying sink. This enables external table formats (Iceberg, Delta Lake, Hudi) to obtain per-file metadata after a DataFusion write without re-reading file footers or bypassing the write pipeline entirely. Closes apache#23472 --- datafusion/datasource-parquet/src/sink.rs | 302 +++++++++++++++++++++- datafusion/datasource/src/sink.rs | 248 ++++++++++++++++++ 2 files changed, 549 insertions(+), 1 deletion(-) diff --git a/datafusion/datasource-parquet/src/sink.rs b/datafusion/datasource-parquet/src/sink.rs index f15f67aab0a87..d27b66c4c6f42 100644 --- a/datafusion/datasource-parquet/src/sink.rs +++ b/datafusion/datasource-parquet/src/sink.rs @@ -32,7 +32,7 @@ use datafusion_common_runtime::{JoinSet, SpawnedTask}; use datafusion_datasource::display::FileGroupDisplay; use datafusion_datasource::file_compression_type::FileCompressionType; use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig}; -use datafusion_datasource::sink::DataSink; +use datafusion_datasource::sink::{DataSink, FileWriteMetadata}; use datafusion_datasource::write::demux::DemuxedStreamReceiver; use datafusion_datasource::write::{ ObjectWriterBuilder, SharedBuffer, get_writer_schema, @@ -411,6 +411,31 @@ impl DataSink for ParquetSink { ) -> Result { FileSink::write_all(self, data, context).await } + + fn file_metadata(&self) -> Vec { + let written = self.written.lock(); + written + .iter() + .map(|(path, parquet_meta)| { + let row_count = parquet_meta.file_metadata().num_rows() as u64; + let byte_size: u64 = parquet_meta + .row_groups() + .iter() + .map(|rg| rg.compressed_size() as u64) + .sum(); + + FileWriteMetadata { + path: path.to_string(), + row_count, + byte_size, + // Full Parquet metadata is accessible via ParquetSink::written() + // for Rust consumers. Thrift serialization for FFI consumers can + // be added when a concrete use case (e.g. Python via PyO3) needs it. + format_metadata: None, + } + }) + .collect() + } } /// Consumes a stream of [ArrowLeafColumn] via a channel and serializes them using an [ArrowColumnWriter] @@ -752,3 +777,278 @@ async fn output_single_parquet_file_parallelized( .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; Ok(parquet_meta_data) } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ArrayRef, StringArray}; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::config::TableParquetOptions; + use datafusion_datasource::PartitionedFile; + use datafusion_datasource::file_groups::FileGroup; + use datafusion_datasource::file_sink_config::{FileOutputMode, FileSinkConfig}; + use datafusion_datasource::sink::{DataSink, DataSinkExec}; + use datafusion_datasource::url::ListingTableUrl; + use datafusion_execution::config::SessionConfig; + use datafusion_execution::object_store::ObjectStoreUrl; + use datafusion_execution::runtime_env::RuntimeEnv; + use datafusion_expr::dml::InsertOp; + use datafusion_physical_plan::stream::RecordBatchStreamAdapter; + use object_store::local::LocalFileSystem; + + fn build_test_ctx(store_url: &ObjectStoreUrl) -> Arc { + let tmp_dir = tempfile::TempDir::new().unwrap(); + let local = Arc::new( + LocalFileSystem::new_with_prefix(&tmp_dir) + .expect("should create object store"), + ); + + let session = SessionConfig::default(); + let runtime = RuntimeEnv::default(); + runtime + .object_store_registry + .register_store(store_url.as_ref(), local); + + Arc::new( + TaskContext::default() + .with_session_config(session) + .with_runtime(Arc::new(runtime)), + ) + } + + fn create_test_sink() -> (Arc, SchemaRef, ObjectStoreUrl) { + let field_a = Field::new("a", DataType::Utf8, false); + let field_b = Field::new("b", DataType::Utf8, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + let object_store_url = ObjectStoreUrl::local_filesystem(); + + let file_sink_config = FileSinkConfig { + original_url: String::default(), + object_store_url: object_store_url.clone(), + file_group: FileGroup::new(vec![PartitionedFile::new("/tmp".to_string(), 1)]), + table_paths: vec![ListingTableUrl::parse("file:///tmp/test/").unwrap()], + output_schema: Arc::clone(&schema), + table_partition_cols: vec![], + insert_op: InsertOp::Overwrite, + keep_partition_by_columns: false, + file_extension: "parquet".into(), + file_output_mode: FileOutputMode::Automatic, + }; + + let parquet_sink = Arc::new(ParquetSink::new( + file_sink_config, + TableParquetOptions::default(), + )); + + (parquet_sink, schema, object_store_url) + } + + fn make_test_batch(schema: &SchemaRef) -> RecordBatch { + let col_a: ArrayRef = Arc::new(StringArray::from(vec!["foo", "bar", "baz"])); + let col_b: ArrayRef = Arc::new(StringArray::from(vec!["one", "two", "three"])); + RecordBatch::try_new(Arc::clone(schema), vec![col_a, col_b]).unwrap() + } + + #[test] + fn file_metadata_empty_before_write() { + let (sink, _schema, _url) = create_test_sink(); + assert!( + sink.file_metadata().is_empty(), + "file_metadata should be empty before any write" + ); + } + + #[tokio::test] + async fn file_metadata_populated_after_write() { + let (sink, schema, object_store_url) = create_test_sink(); + let ctx = build_test_ctx(&object_store_url); + let batch = make_test_batch(&schema); + + let data: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(vec![Ok(batch)]), + )); + + let count = DataSink::write_all(sink.as_ref(), data, &ctx) + .await + .unwrap(); + assert_eq!(count, 3, "should have written 3 rows"); + + let metadata = sink.file_metadata(); + assert_eq!(metadata.len(), 1, "should have one file metadata entry"); + assert_eq!(metadata[0].row_count, 3); + assert!(metadata[0].byte_size > 0, "byte_size should be non-zero"); + assert!(!metadata[0].path.is_empty(), "path should be non-empty"); + assert!(metadata[0].format_metadata.is_none()); + } + + #[tokio::test] + async fn file_metadata_count_equals_write_all_count() { + let (sink, schema, object_store_url) = create_test_sink(); + let ctx = build_test_ctx(&object_store_url); + let batch = make_test_batch(&schema); + + let data: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(vec![Ok(batch)]), + )); + + let count = DataSink::write_all(sink.as_ref(), data, &ctx) + .await + .unwrap(); + + let sum_rows: u64 = sink.file_metadata().iter().map(|f| f.row_count).sum(); + assert_eq!( + count, sum_rows, + "write_all count must equal sum of per-file row_counts" + ); + } + + #[tokio::test] + async fn file_metadata_consistent_with_written() { + let (sink, schema, object_store_url) = create_test_sink(); + let ctx = build_test_ctx(&object_store_url); + let batch = make_test_batch(&schema); + + let data: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(vec![Ok(batch)]), + )); + + DataSink::write_all(sink.as_ref(), data, &ctx) + .await + .unwrap(); + + let file_meta = sink.file_metadata(); + let written = sink.written(); + + assert_eq!( + file_meta.len(), + written.len(), + "file_metadata and written() should have same number of entries" + ); + + for fm in &file_meta { + let path = Path::from(fm.path.as_str()); + let parquet_meta = written + .get(&path) + .expect("file_metadata path should exist in written()"); + assert_eq!(fm.row_count, parquet_meta.file_metadata().num_rows() as u64); + + let expected_bytes: u64 = parquet_meta + .row_groups() + .iter() + .map(|rg| rg.compressed_size() as u64) + .sum(); + assert_eq!(fm.byte_size, expected_bytes); + } + } + + #[tokio::test] + async fn file_metadata_is_idempotent() { + let (sink, schema, object_store_url) = create_test_sink(); + let ctx = build_test_ctx(&object_store_url); + let batch = make_test_batch(&schema); + + let data: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(vec![Ok(batch)]), + )); + + DataSink::write_all(sink.as_ref(), data, &ctx) + .await + .unwrap(); + + let first = sink.file_metadata(); + let second = sink.file_metadata(); + assert_eq!(first, second); + } + + #[tokio::test] + async fn partitioned_write_produces_multiple_metadata_entries() { + let field_a = Field::new("a", DataType::Utf8, false); + let field_b = Field::new("b", DataType::Utf8, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + let object_store_url = ObjectStoreUrl::local_filesystem(); + + let file_sink_config = FileSinkConfig { + original_url: String::default(), + object_store_url: object_store_url.clone(), + file_group: FileGroup::new(vec![PartitionedFile::new("/tmp".to_string(), 1)]), + table_paths: vec![ + ListingTableUrl::parse("file:///tmp/partitioned/").unwrap(), + ], + output_schema: Arc::clone(&schema), + table_partition_cols: vec![("a".to_string(), DataType::Utf8)], + insert_op: InsertOp::Overwrite, + keep_partition_by_columns: false, + file_extension: "parquet".into(), + file_output_mode: FileOutputMode::Automatic, + }; + + let sink = Arc::new(ParquetSink::new( + file_sink_config, + TableParquetOptions::default(), + )); + + let col_a: ArrayRef = Arc::new(StringArray::from(vec!["x", "y", "x"])); + let col_b: ArrayRef = Arc::new(StringArray::from(vec!["one", "two", "three"])); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![col_a, col_b]).unwrap(); + + let ctx = build_test_ctx(&object_store_url); + let data: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(vec![Ok(batch)]), + )); + + let count = DataSink::write_all(sink.as_ref(), data, &ctx) + .await + .unwrap(); + + let metadata = sink.file_metadata(); + assert_eq!( + metadata.len(), + 2, + "partitioned write with 2 partition values should produce 2 files" + ); + + assert_eq!(count, 3); + let sum_rows: u64 = metadata.iter().map(|f| f.row_count).sum(); + assert_eq!(sum_rows, 3); + } + + /// Verifies that `DataSinkExec::file_metadata()` returns the correct + /// metadata when the underlying `ParquetSink` has written files. + /// This tests the `Arc` dispatch path that real consumers use. + #[tokio::test] + async fn data_sink_exec_e2e_file_metadata_after_execute() { + let (sink, schema, object_store_url) = create_test_sink(); + let ctx = build_test_ctx(&object_store_url); + let batch = make_test_batch(&schema); + + let data: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(vec![Ok(batch)]), + )); + + // Write through the DataSink trait (same call DataSinkExec::execute makes) + let count = DataSink::write_all(sink.as_ref(), data, &ctx) + .await + .unwrap(); + assert_eq!(count, 3); + + // Wrap in DataSinkExec and verify file_metadata is accessible + // through the exec's convenience accessor (the Arc path) + let input: Arc = Arc::new( + datafusion_physical_plan::empty::EmptyExec::new(Arc::clone(&schema)), + ); + let exec = DataSinkExec::new(input, sink, None); + + let metadata = exec.file_metadata(); + assert_eq!(metadata.len(), 1, "should have one file after write"); + assert_eq!(metadata[0].row_count, 3); + assert!(metadata[0].byte_size > 0); + assert!(!metadata[0].path.is_empty()); + } +} diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index 18ebe80773e8a..6029800062f72 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -24,6 +24,7 @@ use std::sync::Arc; use arrow::array::{ArrayRef, RecordBatch, UInt64Array}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use bytes::Bytes; use datafusion_common::{Result, assert_eq_or_internal_err}; use datafusion_execution::TaskContext; use datafusion_physical_expr::{Distribution, EquivalenceProperties}; @@ -40,6 +41,32 @@ use async_trait::async_trait; use datafusion_physical_plan::execution_plan::{EvaluationType, SchedulingType}; use futures::StreamExt; +/// Metadata about a single file produced by a [`DataSink`] write operation. +/// +/// This struct is format-agnostic. The [`Self::format_metadata`] field carries +/// serialized format-specific metadata (e.g., a Parquet file footer serialized +/// via Thrift Compact Protocol). +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FileWriteMetadata { + /// Object-store path where the file was written. + pub path: String, + /// Number of rows written to this specific file. + pub row_count: u64, + /// Sum of compressed row group sizes in bytes. + /// + /// Note: this may differ slightly from the actual on-disk file size as it + /// excludes the Parquet footer, page indexes, and other metadata overhead. + pub byte_size: u64, + /// Format-specific metadata serialized as bytes. + /// + /// For Parquet files this contains the Thrift-serialized `FileMetaData` + /// (the same bytes found in the Parquet footer), enabling consumers to + /// reconstruct column statistics without re-reading the file. + /// + /// For formats that do not produce file-level metadata this is `None`. + pub format_metadata: Option, +} + /// `DataSink` implements writing streams of [`RecordBatch`]es to /// user defined destinations. /// @@ -71,6 +98,21 @@ pub trait DataSink: Any + DisplayAs + Debug + Send + Sync { data: SendableRecordBatchStream, context: &Arc, ) -> Result; + + /// Returns metadata for files written during the most recent + /// [`Self::write_all`] call. + /// + /// This accessor is intended for consumers that need per-file + /// statistics (e.g., column sizes, null counts, value bounds) after + /// a write completes. It should be called after `write_all` returns + /// successfully. + /// + /// The default implementation returns an empty vector. Implementations + /// that collect file metadata during writes (e.g., Parquet sinks) + /// should override this method. + fn file_metadata(&self) -> Vec { + Vec::new() + } } impl dyn DataSink { @@ -145,6 +187,15 @@ impl DataSinkExec { &self.sort_order } + /// Returns per-file metadata from the underlying sink, if available. + /// + /// This is a convenience accessor that delegates to + /// [`DataSink::file_metadata`]. It should be called after the write + /// operation has completed. + pub fn file_metadata(&self) -> Vec { + self.sink.file_metadata() + } + fn create_schema( input: &Arc, schema: SchemaRef, @@ -293,3 +344,200 @@ fn make_count_schema() -> SchemaRef { false, )])) } + +#[cfg(test)] +mod sink_tests { + use super::*; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_physical_plan::stream::RecordBatchStreamAdapter; + use futures::stream; + + /// A minimal DataSink that does NOT override `file_metadata`. + /// Used to verify that the default implementation returns empty. + #[derive(Debug)] + struct MinimalSink { + schema: SchemaRef, + } + + impl MinimalSink { + fn new(schema: SchemaRef) -> Self { + Self { schema } + } + } + + impl DisplayAs for MinimalSink { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "MinimalSink") + } + } + + #[async_trait] + impl DataSink for MinimalSink { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + async fn write_all( + &self, + _data: SendableRecordBatchStream, + _context: &Arc, + ) -> Result { + Ok(42) + } + } + + /// A DataSink that overrides `file_metadata`, simulating a + /// format-specific sink (like ParquetSink) that provides metadata. + #[derive(Debug)] + struct MetadataProvidingSink { + schema: SchemaRef, + } + + impl MetadataProvidingSink { + fn new(schema: SchemaRef) -> Self { + Self { schema } + } + } + + impl DisplayAs for MetadataProvidingSink { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "MetadataProvidingSink") + } + } + + #[async_trait] + impl DataSink for MetadataProvidingSink { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + async fn write_all( + &self, + _data: SendableRecordBatchStream, + _context: &Arc, + ) -> Result { + Ok(100) + } + + fn file_metadata(&self) -> Vec { + vec![ + FileWriteMetadata { + path: "part-0.parquet".to_string(), + row_count: 60, + byte_size: 4096, + format_metadata: Some(Bytes::from_static(b"fake-metadata-0")), + }, + FileWriteMetadata { + path: "part-1.parquet".to_string(), + row_count: 40, + byte_size: 3072, + format_metadata: Some(Bytes::from_static(b"fake-metadata-1")), + }, + ] + } + } + + fn test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])) + } + + fn test_context() -> Arc { + Arc::new(TaskContext::default()) + } + + fn empty_stream(schema: SchemaRef) -> SendableRecordBatchStream { + Box::pin(RecordBatchStreamAdapter::new(schema, stream::empty())) + } + + #[test] + fn file_write_metadata_equality() { + let a = FileWriteMetadata { + path: "a.parquet".to_string(), + row_count: 10, + byte_size: 1024, + format_metadata: None, + }; + let b = a.clone(); + assert_eq!(a, b); + } + + #[test] + fn file_write_metadata_without_format_metadata() { + let meta = FileWriteMetadata { + path: "data.csv".to_string(), + row_count: 50, + byte_size: 1024, + format_metadata: None, + }; + assert_eq!(meta.path, "data.csv"); + assert!(meta.format_metadata.is_none()); + } + + #[tokio::test] + async fn default_file_metadata_returns_empty() { + let schema = test_schema(); + let sink = MinimalSink::new(Arc::clone(&schema)); + let ctx = test_context(); + let data = empty_stream(schema); + + // write_all succeeds + let count = sink.write_all(data, &ctx).await.unwrap(); + assert_eq!(count, 42); + + // default file_metadata returns empty + let metadata = sink.file_metadata(); + assert!(metadata.is_empty()); + } + + #[test] + fn overridden_file_metadata_returns_entries() { + let schema = test_schema(); + let sink = MetadataProvidingSink::new(schema); + + let metadata = sink.file_metadata(); + assert_eq!(metadata.len(), 2); + assert_eq!(metadata[0].path, "part-0.parquet"); + assert_eq!(metadata[1].path, "part-1.parquet"); + } + + #[test] + fn file_metadata_is_idempotent() { + let schema = test_schema(); + let sink = MetadataProvidingSink::new(schema); + + let first = sink.file_metadata(); + let second = sink.file_metadata(); + assert_eq!(first, second); + } + + #[test] + fn data_sink_exec_file_metadata_delegates_to_sink() { + let schema = test_schema(); + let sink: Arc = + Arc::new(MetadataProvidingSink::new(Arc::clone(&schema))); + + let input: Arc = + Arc::new(datafusion_physical_plan::empty::EmptyExec::new(schema)); + + let exec = DataSinkExec::new(input, sink, None); + let metadata = exec.file_metadata(); + + assert_eq!(metadata.len(), 2); + assert_eq!(metadata[0].path, "part-0.parquet"); + assert_eq!(metadata[1].path, "part-1.parquet"); + } + + #[test] + fn data_sink_exec_file_metadata_empty_for_minimal_sink() { + let schema = test_schema(); + let sink: Arc = Arc::new(MinimalSink::new(Arc::clone(&schema))); + + let input: Arc = + Arc::new(datafusion_physical_plan::empty::EmptyExec::new(schema)); + + let exec = DataSinkExec::new(input, sink, None); + let metadata = exec.file_metadata(); + + assert!(metadata.is_empty()); + } +}