timsaucer commented on code in PR #23656:
URL: https://github.com/apache/datafusion/pull/23656#discussion_r3668541611
##########
datafusion/datasource/src/sink.rs:
##########
@@ -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.
Review Comment:
`FileWriteMetada` should be format agnostic, right? Here we're talking about
row group sizes, which isn't a concept in all of the datasources. I *think* we
just need a comment update here, but if this is specific to parquet then it
probably needs to go into the `format_metadata`.
##########
datafusion/datasource-parquet/src/sink.rs:
##########
@@ -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<TaskContext> {
+ 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<ParquetSink>, 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();
Review Comment:
Same comment as above about reusing the `record_batch!` macro
##########
datafusion/datasource/src/sink.rs:
##########
@@ -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<Bytes>,
Review Comment:
I suppose the other option, which matches some metadata approaches is that
these be key/value pairs. That has the advantage of being human readable but
the downside of potential data loss on conversion, such as `f64 -> String`.
Here we have to pay a serialization/deserialization cost so parsing strings is
likely not too much less performant. It's probably not a hot path anyway.
##########
datafusion/datasource-parquet/src/sink.rs:
##########
@@ -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<TaskContext> {
+ 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<ParquetSink>, 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()
Review Comment:
There is a `record_batch!` macro that makes this a little more pleasant.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
