From f9c560a5cafb255a56ced8f0ace19c7ccbb12d98 Mon Sep 17 00:00:00 2001 From: Jared Yu Date: Fri, 10 Jul 2026 21:10:21 -0700 Subject: [PATCH 1/3] 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()); + } +} From 7ef1553193a4f070d7c0ed7f7f0c5a7db13380a7 Mon Sep 17 00:00:00 2001 From: Jared Yu Date: Wed, 29 Jul 2026 19:01:17 -0700 Subject: [PATCH 2/3] Address review feedback: fix TempDir lifetime, portable paths, safe casts, format-agnostic docs, record_batch! macro --- datafusion/datasource-parquet/Cargo.toml | 1 + datafusion/datasource-parquet/src/sink.rs | 202 ++++++++++++++-------- datafusion/datasource/src/sink.rs | 8 +- 3 files changed, 134 insertions(+), 77 deletions(-) diff --git a/datafusion/datasource-parquet/Cargo.toml b/datafusion/datasource-parquet/Cargo.toml index 32424069c17a0..3752fdcfa39c3 100644 --- a/datafusion/datasource-parquet/Cargo.toml +++ b/datafusion/datasource-parquet/Cargo.toml @@ -62,6 +62,7 @@ criterion = { workspace = true } datafusion-functions = { workspace = true } datafusion-functions-nested = { workspace = true } tempfile = { workspace = true } +url = { workspace = true } # Note: add additional linter rules in lib.rs. # Rust does not support workspace + new linter rules in subcrates yet diff --git a/datafusion/datasource-parquet/src/sink.rs b/datafusion/datasource-parquet/src/sink.rs index d27b66c4c6f42..0d5fcbb46473d 100644 --- a/datafusion/datasource-parquet/src/sink.rs +++ b/datafusion/datasource-parquet/src/sink.rs @@ -417,11 +417,17 @@ impl DataSink for ParquetSink { written .iter() .map(|(path, parquet_meta)| { - let row_count = parquet_meta.file_metadata().num_rows() as u64; + let row_count: u64 = parquet_meta + .file_metadata() + .num_rows() + .try_into() + .unwrap_or(0); let byte_size: u64 = parquet_meta .row_groups() .iter() - .map(|rg| rg.compressed_size() as u64) + .map(|rg| { + u64::try_from(rg.compressed_size()).unwrap_or(0) + }) .sum(); FileWriteMetadata { @@ -781,7 +787,7 @@ async fn output_single_parquet_file_parallelized( #[cfg(test)] mod tests { use super::*; - use arrow::array::{ArrayRef, StringArray}; + use arrow::array::record_batch; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::config::TableParquetOptions; use datafusion_datasource::PartitionedFile; @@ -795,38 +801,42 @@ mod tests { use datafusion_expr::dml::InsertOp; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use object_store::local::LocalFileSystem; + use tempfile::TempDir; + + /// Test fixture that keeps the TempDir alive for the duration of the test. + struct TestFixture { + _tmp_dir: TempDir, + ctx: Arc, + sink: Arc, + schema: SchemaRef, + } - fn build_test_ctx(store_url: &ObjectStoreUrl) -> Arc { - let tmp_dir = tempfile::TempDir::new().unwrap(); + fn setup_test_fixture() -> TestFixture { + let tmp_dir = TempDir::new().unwrap(); + let tmp_path = tmp_dir.path().to_owned(); let local = Arc::new( - LocalFileSystem::new_with_prefix(&tmp_dir) + LocalFileSystem::new_with_prefix(&tmp_path) .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)), + let object_store_url = ObjectStoreUrl::local_filesystem(); + let table_url = ListingTableUrl::parse( + url::Url::from_directory_path(&tmp_path).unwrap().as_str(), ) - } + .unwrap(); - 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()], + file_group: FileGroup::new(vec![PartitionedFile::new( + tmp_path.to_string_lossy().to_string(), + 1, + )]), + table_paths: vec![table_url], output_schema: Arc::clone(&schema), table_partition_cols: vec![], insert_op: InsertOp::Overwrite, @@ -835,46 +845,64 @@ mod tests { file_output_mode: FileOutputMode::Automatic, }; - let parquet_sink = Arc::new(ParquetSink::new( + let sink = Arc::new(ParquetSink::new( file_sink_config, TableParquetOptions::default(), )); - (parquet_sink, schema, object_store_url) + let session = SessionConfig::default(); + let runtime = RuntimeEnv::default(); + runtime + .object_store_registry + .register_store(object_store_url.as_ref(), local); + + let ctx = Arc::new( + TaskContext::default() + .with_session_config(session) + .with_runtime(Arc::new(runtime)), + ); + + TestFixture { + _tmp_dir: tmp_dir, + ctx, + sink, + schema, + } } - 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() + fn make_test_batch(_schema: &SchemaRef) -> RecordBatch { + record_batch!( + ("a", Utf8, ["foo", "bar", "baz"]), + ("b", Utf8, ["one", "two", "three"]) + ) + .unwrap() } #[test] fn file_metadata_empty_before_write() { - let (sink, _schema, _url) = create_test_sink(); + let fixture = setup_test_fixture(); assert!( - sink.file_metadata().is_empty(), + fixture.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 fixture = setup_test_fixture(); + let batch = make_test_batch(&fixture.schema); let data: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( - Arc::clone(&schema), + Arc::clone(&fixture.schema), futures::stream::iter(vec![Ok(batch)]), )); - let count = DataSink::write_all(sink.as_ref(), data, &ctx) + let count = DataSink::write_all(fixture.sink.as_ref(), data, &fixture.ctx) .await .unwrap(); assert_eq!(count, 3, "should have written 3 rows"); - let metadata = sink.file_metadata(); + let metadata = fixture.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"); @@ -884,20 +912,19 @@ mod tests { #[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 fixture = setup_test_fixture(); + let batch = make_test_batch(&fixture.schema); let data: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( - Arc::clone(&schema), + Arc::clone(&fixture.schema), futures::stream::iter(vec![Ok(batch)]), )); - let count = DataSink::write_all(sink.as_ref(), data, &ctx) + let count = DataSink::write_all(fixture.sink.as_ref(), data, &fixture.ctx) .await .unwrap(); - let sum_rows: u64 = sink.file_metadata().iter().map(|f| f.row_count).sum(); + let sum_rows: u64 = fixture.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" @@ -906,21 +933,20 @@ mod tests { #[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 fixture = setup_test_fixture(); + let batch = make_test_batch(&fixture.schema); let data: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( - Arc::clone(&schema), + Arc::clone(&fixture.schema), futures::stream::iter(vec![Ok(batch)]), )); - DataSink::write_all(sink.as_ref(), data, &ctx) + DataSink::write_all(fixture.sink.as_ref(), data, &fixture.ctx) .await .unwrap(); - let file_meta = sink.file_metadata(); - let written = sink.written(); + let file_meta = fixture.sink.file_metadata(); + let written = fixture.sink.written(); assert_eq!( file_meta.len(), @@ -933,12 +959,17 @@ mod tests { 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_rows: u64 = parquet_meta + .file_metadata() + .num_rows() + .try_into() + .unwrap_or(0); + assert_eq!(fm.row_count, expected_rows); let expected_bytes: u64 = parquet_meta .row_groups() .iter() - .map(|rg| rg.compressed_size() as u64) + .map(|rg| u64::try_from(rg.compressed_size()).unwrap_or(0)) .sum(); assert_eq!(fm.byte_size, expected_bytes); } @@ -946,38 +977,50 @@ mod tests { #[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 fixture = setup_test_fixture(); + let batch = make_test_batch(&fixture.schema); let data: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( - Arc::clone(&schema), + Arc::clone(&fixture.schema), futures::stream::iter(vec![Ok(batch)]), )); - DataSink::write_all(sink.as_ref(), data, &ctx) + DataSink::write_all(fixture.sink.as_ref(), data, &fixture.ctx) .await .unwrap(); - let first = sink.file_metadata(); - let second = sink.file_metadata(); + let first = fixture.sink.file_metadata(); + let second = fixture.sink.file_metadata(); assert_eq!(first, second); } #[tokio::test] async fn partitioned_write_produces_multiple_metadata_entries() { + let tmp_dir = TempDir::new().unwrap(); + let tmp_path = tmp_dir.path().to_owned(); + let local = Arc::new( + LocalFileSystem::new_with_prefix(&tmp_path) + .expect("should create object store"), + ); + + let object_store_url = ObjectStoreUrl::local_filesystem(); + let table_url = ListingTableUrl::parse( + url::Url::from_directory_path(&tmp_path).unwrap().as_str(), + ) + .unwrap(); + 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(), - ], + file_group: FileGroup::new(vec![PartitionedFile::new( + tmp_path.to_string_lossy().to_string(), + 1, + )]), + table_paths: vec![table_url], output_schema: Arc::clone(&schema), table_partition_cols: vec![("a".to_string(), DataType::Utf8)], insert_op: InsertOp::Overwrite, @@ -991,12 +1034,24 @@ mod tests { 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 batch = record_batch!( + ("a", Utf8, ["x", "y", "x"]), + ("b", Utf8, ["one", "two", "three"]) + ) + .unwrap(); + + let session = SessionConfig::default(); + let runtime = RuntimeEnv::default(); + runtime + .object_store_registry + .register_store(object_store_url.as_ref(), local); + + let ctx = Arc::new( + TaskContext::default() + .with_session_config(session) + .with_runtime(Arc::new(runtime)), + ); - let ctx = build_test_ctx(&object_store_url); let data: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( Arc::clone(&schema), futures::stream::iter(vec![Ok(batch)]), @@ -1023,17 +1078,16 @@ mod tests { /// 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 fixture = setup_test_fixture(); + let batch = make_test_batch(&fixture.schema); let data: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( - Arc::clone(&schema), + Arc::clone(&fixture.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) + let count = DataSink::write_all(fixture.sink.as_ref(), data, &fixture.ctx) .await .unwrap(); assert_eq!(count, 3); @@ -1041,9 +1095,9 @@ mod tests { // 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)), + datafusion_physical_plan::empty::EmptyExec::new(Arc::clone(&fixture.schema)), ); - let exec = DataSinkExec::new(input, sink, None); + let exec = DataSinkExec::new(input, fixture.sink, None); let metadata = exec.file_metadata(); assert_eq!(metadata.len(), 1, "should have one file after write"); diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index 6029800062f72..38bad1b0f4423 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -52,10 +52,12 @@ pub struct FileWriteMetadata { pub path: String, /// Number of rows written to this specific file. pub row_count: u64, - /// Sum of compressed row group sizes in bytes. + /// Best-effort total encoded data size 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. + /// How this value is computed depends on the format. For Parquet it is the + /// sum of compressed row-group sizes (which excludes the footer and page + /// indexes); for other formats it may be the full file size or unavailable + /// (zero). Consumers should treat this as an approximation. pub byte_size: u64, /// Format-specific metadata serialized as bytes. /// From a66620e3961fce72884a891a4c9cd616b02f883f Mon Sep 17 00:00:00 2001 From: Jared Yu Date: Thu, 30 Jul 2026 11:15:43 -0700 Subject: [PATCH 3/3] Fix cargo fmt, rename sink_tests to tests, update Cargo.lock, link tracking issue #24010 --- Cargo.lock | 1 + datafusion/datasource-parquet/src/sink.rs | 12 ++++++++---- datafusion/datasource/src/sink.rs | 2 +- 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f00c931f15032..4d46878b648a7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2069,6 +2069,7 @@ dependencies = [ "parquet", "tempfile", "tokio", + "url", ] [[package]] diff --git a/datafusion/datasource-parquet/src/sink.rs b/datafusion/datasource-parquet/src/sink.rs index 0d5fcbb46473d..50564f06af621 100644 --- a/datafusion/datasource-parquet/src/sink.rs +++ b/datafusion/datasource-parquet/src/sink.rs @@ -425,9 +425,7 @@ impl DataSink for ParquetSink { let byte_size: u64 = parquet_meta .row_groups() .iter() - .map(|rg| { - u64::try_from(rg.compressed_size()).unwrap_or(0) - }) + .map(|rg| u64::try_from(rg.compressed_size()).unwrap_or(0)) .sum(); FileWriteMetadata { @@ -437,6 +435,7 @@ impl DataSink for ParquetSink { // 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. + // Tracked by https://github.com/apache/datafusion/issues/24010 format_metadata: None, } }) @@ -924,7 +923,12 @@ mod tests { .await .unwrap(); - let sum_rows: u64 = fixture.sink.file_metadata().iter().map(|f| f.row_count).sum(); + let sum_rows: u64 = fixture + .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" diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index 38bad1b0f4423..21f6fc465218a 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -348,7 +348,7 @@ fn make_count_schema() -> SchemaRef { } #[cfg(test)] -mod sink_tests { +mod tests { use super::*; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_physical_plan::stream::RecordBatchStreamAdapter;