Copilot commented on code in PR #23656:
URL: https://github.com/apache/datafusion/pull/23656#discussion_r3668352423
##########
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)),
+ )
+ }
Review Comment:
`tmp_dir` is dropped when `build_test_ctx` returns, which can delete the
backing directory while the registered `LocalFileSystem` is still expected to
write/read during the test. This is likely to cause test failures/flakiness.
Keep the `TempDir` alive for the duration of the test (e.g., return it
alongside the `TaskContext`, store it in a test fixture struct, or otherwise
hold it in the calling scope).
##########
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,
Review Comment:
`FileWriteMetadata` is documented as format-agnostic, but the `byte_size`
field description is Parquet-specific (row groups, Parquet footer, page
indexes). Consider rewording `byte_size` docs to be format-neutral (e.g.,
'total bytes written for this file' or 'best-effort encoded data size') and, if
needed, mention Parquet as an example rather than defining semantics in Parquet
terms.
##########
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,
+ };
Review Comment:
These tests hardcode `/tmp` paths and `file:///tmp/...` URLs, which is not
portable (notably on Windows, and some sandboxed CI environments). Prefer
constructing the sink config from the same `TempDir` used by the
`LocalFileSystem` prefix (derive `ListingTableUrl` from the temp path and avoid
absolute `/tmp`), so the tests are hermetic and cross-platform.
##########
datafusion/datasource-parquet/src/sink.rs:
##########
@@ -411,6 +411,31 @@ impl DataSink for ParquetSink {
) -> Result<u64> {
FileSink::write_all(self, data, context).await
}
+
+ fn file_metadata(&self) -> Vec<FileWriteMetadata> {
+ 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();
Review Comment:
This uses `as u64` casts from Parquet metadata values that are typically
signed integers (e.g., `num_rows()` / `compressed_size()` are commonly `i64`).
If a negative value ever occurs (corrupt metadata, unexpected upstream
behavior), `as u64` will wrap to a huge number. Prefer `try_into()` (and either
return an error, or clamp/guard with an internal error) to avoid producing
nonsensical `row_count`/`byte_size`.
--
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]
