diff --git a/datafusion/datasource-csv/src/file_format.rs b/datafusion/datasource-csv/src/file_format.rs index 6b131f2beed10..7fd9cffa7081f 100644 --- a/datafusion/datasource-csv/src/file_format.rs +++ b/datafusion/datasource-csv/src/file_format.rs @@ -43,14 +43,19 @@ use datafusion_datasource::file_format::{ DEFAULT_SCHEMA_INFER_MAX_RECORD, FileFormat, FileFormatFactory, }; use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; -use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig}; +use datafusion_datasource::file_sink_config::{ + FileSink, FileSinkConfig, FileSinkMetrics, +}; use datafusion_datasource::sink::{DataSink, DataSinkExec}; use datafusion_datasource::write::BatchSerializer; use datafusion_datasource::write::demux::DemuxedStreamReceiver; -use datafusion_datasource::write::orchestration::spawn_writer_tasks_and_join; +use datafusion_datasource::write::orchestration::{ + StatelessWriterOptions, spawn_writer_tasks_and_join_with_metrics, +}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::dml::InsertOp; use datafusion_physical_expr_common::sort_expr::LexRequirement; +use datafusion_physical_plan::metrics::MetricsSet; use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; use datafusion_session::Session; @@ -741,6 +746,8 @@ pub struct CsvSink { /// Config options for writing data config: FileSinkConfig, writer_options: CsvWriterOptions, + /// Metrics for tracking write operations + metrics: FileSinkMetrics, } impl Debug for CsvSink { @@ -771,6 +778,7 @@ impl CsvSink { Self { config, writer_options, + metrics: FileSinkMetrics::new(), } } @@ -800,11 +808,12 @@ impl FileSink for CsvSink { .with_builder(builder) .with_header(header), ) as _; - spawn_writer_tasks_and_join( + spawn_writer_tasks_and_join_with_metrics( context, serializer, - self.writer_options.compression.into(), - self.writer_options.compression_level, + StatelessWriterOptions::new(self.writer_options.compression.into()) + .with_compression_level(self.writer_options.compression_level) + .with_metrics(&self.metrics), object_store, demux_task, file_stream_rx, @@ -815,6 +824,10 @@ impl FileSink for CsvSink { #[async_trait] impl DataSink for CsvSink { + fn metrics(&self) -> Option { + Some(self.metrics.snapshot()) + } + fn schema(&self) -> &SchemaRef { self.config.output_schema() } diff --git a/datafusion/datasource-json/src/file_format.rs b/datafusion/datasource-json/src/file_format.rs index 43bde2a039059..eba11c8e1568f 100644 --- a/datafusion/datasource-json/src/file_format.rs +++ b/datafusion/datasource-json/src/file_format.rs @@ -46,15 +46,20 @@ use datafusion_datasource::file_format::{ DEFAULT_SCHEMA_INFER_MAX_RECORD, FileFormat, FileFormatFactory, }; use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; -use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig}; +use datafusion_datasource::file_sink_config::{ + FileSink, FileSinkConfig, FileSinkMetrics, +}; use datafusion_datasource::sink::{DataSink, DataSinkExec}; use datafusion_datasource::source::DataSourceExec; use datafusion_datasource::write::BatchSerializer; use datafusion_datasource::write::demux::DemuxedStreamReceiver; -use datafusion_datasource::write::orchestration::spawn_writer_tasks_and_join; +use datafusion_datasource::write::orchestration::{ + StatelessWriterOptions, spawn_writer_tasks_and_join_with_metrics, +}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::dml::InsertOp; use datafusion_physical_expr_common::sort_expr::LexRequirement; +use datafusion_physical_plan::metrics::MetricsSet; use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; use datafusion_session::Session; @@ -411,6 +416,8 @@ pub struct JsonSink { config: FileSinkConfig, /// Writer options for underlying Json writer writer_options: JsonWriterOptions, + /// Metrics for tracking write operations + metrics: FileSinkMetrics, } impl Debug for JsonSink { @@ -441,6 +448,7 @@ impl JsonSink { Self { config, writer_options, + metrics: FileSinkMetrics::new(), } } @@ -464,11 +472,12 @@ impl FileSink for JsonSink { object_store: Arc, ) -> Result { let serializer = Arc::new(JsonSerializer::new()) as _; - spawn_writer_tasks_and_join( + spawn_writer_tasks_and_join_with_metrics( context, serializer, - self.writer_options.compression.into(), - self.writer_options.compression_level, + StatelessWriterOptions::new(self.writer_options.compression.into()) + .with_compression_level(self.writer_options.compression_level) + .with_metrics(&self.metrics), object_store, demux_task, file_stream_rx, @@ -479,6 +488,10 @@ impl FileSink for JsonSink { #[async_trait] impl DataSink for JsonSink { + fn metrics(&self) -> Option { + Some(self.metrics.snapshot()) + } + fn schema(&self) -> &SchemaRef { self.config.output_schema() } diff --git a/datafusion/datasource-parquet/src/sink.rs b/datafusion/datasource-parquet/src/sink.rs index f15f67aab0a87..ba3c5294f74ac 100644 --- a/datafusion/datasource-parquet/src/sink.rs +++ b/datafusion/datasource-parquet/src/sink.rs @@ -31,7 +31,9 @@ use datafusion_common::{DataFusionError, HashMap, Result, internal_datafusion_er 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::file_sink_config::{ + FileSink, FileSinkConfig, FileSinkMetrics, +}; use datafusion_datasource::sink::DataSink; use datafusion_datasource::write::demux::DemuxedStreamReceiver; use datafusion_datasource::write::{ @@ -40,10 +42,7 @@ use datafusion_datasource::write::{ use datafusion_execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; -use datafusion_physical_plan::metrics::{ - ElapsedComputeFutureExt, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, - MetricsSet, Time, -}; +use datafusion_physical_plan::metrics::{ElapsedComputeFutureExt, MetricsSet, Time}; use datafusion_physical_plan::{DisplayAs, DisplayFormatType}; use object_store::ObjectStore; use object_store::buffered::BufWriter; @@ -83,7 +82,7 @@ pub struct ParquetSink { /// Optional sorting columns to write to Parquet metadata sorting_columns: Option>, /// Metrics for tracking write operations - metrics: ExecutionPlanMetricsSet, + metrics: FileSinkMetrics, } impl Debug for ParquetSink { @@ -116,7 +115,7 @@ impl ParquetSink { parquet_options, written: Default::default(), sorting_columns: None, - metrics: ExecutionPlanMetricsSet::new(), + metrics: FileSinkMetrics::new(), } } @@ -262,16 +261,12 @@ impl FileSink for ParquetSink { mut file_stream_rx: DemuxedStreamReceiver, object_store: Arc, ) -> Result { - let rows_written_counter = MetricBuilder::new(&self.metrics) - .with_category(MetricCategory::Rows) - .global_counter("rows_written"); + let rows_written_counter = self.metrics.rows_written().clone(); // Note: bytes_written is the sum of compressed row group sizes, which // may differ slightly from the actual on-disk file size (excludes footer, // page indexes, and other Parquet metadata overhead). - let bytes_written_counter = MetricBuilder::new(&self.metrics) - .with_category(MetricCategory::Bytes) - .global_counter("bytes_written"); - let elapsed_compute = MetricBuilder::new(&self.metrics).elapsed_compute(0); + let bytes_written_counter = self.metrics.bytes_written().clone(); + let elapsed_compute = self.metrics.elapsed_compute().clone(); let parquet_opts = &self.parquet_options; @@ -397,7 +392,7 @@ impl FileSink for ParquetSink { #[async_trait] impl DataSink for ParquetSink { fn metrics(&self) -> Option { - Some(self.metrics.clone_inner()) + Some(self.metrics.snapshot()) } fn schema(&self) -> &SchemaRef { diff --git a/datafusion/datasource/src/file_sink_config.rs b/datafusion/datasource/src/file_sink_config.rs index 1abce86a3565f..6261b3e211ca6 100644 --- a/datafusion/datasource/src/file_sink_config.rs +++ b/datafusion/datasource/src/file_sink_config.rs @@ -28,10 +28,72 @@ use datafusion_common_runtime::SpawnedTask; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::dml::InsertOp; +use datafusion_physical_plan::metrics::{ + Count, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, MetricsSet, Time, +}; use async_trait::async_trait; use object_store::ObjectStore; +/// Common metrics exposed by file sinks. +#[derive(Debug)] +pub struct FileSinkMetrics { + metrics: ExecutionPlanMetricsSet, + rows_written: Count, + bytes_written: Count, + elapsed_compute: Time, +} + +impl FileSinkMetrics { + /// Create a new set of file sink metrics. + pub fn new() -> Self { + let metrics = ExecutionPlanMetricsSet::new(); + let rows_written = MetricBuilder::new(&metrics) + .with_category(MetricCategory::Rows) + .global_counter("rows_written"); + let bytes_written = MetricBuilder::new(&metrics) + .with_category(MetricCategory::Bytes) + .global_counter("bytes_written"); + let elapsed_compute = MetricBuilder::new(&metrics).elapsed_compute(0); + Self { + metrics, + rows_written, + bytes_written, + elapsed_compute, + } + } + + /// Return the number of rows written. + pub fn rows_written(&self) -> &Count { + &self.rows_written + } + + /// Return the number of bytes written. + /// + /// The exact meaning is format-specific. For stateless sinks, this is the + /// number of serialized bytes passed to the writer before any stream-level + /// compression is applied. + pub fn bytes_written(&self) -> &Count { + &self.bytes_written + } + + /// Return the elapsed compute time. + pub fn elapsed_compute(&self) -> &Time { + &self.elapsed_compute + } + + /// Return a snapshot of the metrics. + pub fn snapshot(&self) -> MetricsSet { + self.metrics.clone_inner() + } +} + +impl Default for FileSinkMetrics { + fn default() -> Self { + Self::new() + } +} + /// Determines how `FileSink` output paths are interpreted. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum FileOutputMode { diff --git a/datafusion/datasource/src/write/orchestration.rs b/datafusion/datasource/src/write/orchestration.rs index 39c91a1c0d676..40eb1c10249be 100644 --- a/datafusion/datasource/src/write/orchestration.rs +++ b/datafusion/datasource/src/write/orchestration.rs @@ -24,6 +24,7 @@ use std::sync::Arc; use super::demux::DemuxedStreamReceiver; use super::{BatchSerializer, ObjectWriterBuilder}; use crate::file_compression_type::FileCompressionType; +use crate::file_sink_config::FileSinkMetrics; use datafusion_common::error::Result; use arrow::array::RecordBatch; @@ -50,6 +51,9 @@ pub(crate) enum SerializedRecordBatchResult { /// the number of rows successfully written row_count: usize, + + /// the number of serialized bytes successfully written + bytes_written: usize, }, Failure { /// As explained in [`serialize_rb_stream_to_object_store`]: @@ -64,8 +68,12 @@ pub(crate) enum SerializedRecordBatchResult { impl SerializedRecordBatchResult { /// Create the success variant - pub fn success(writer: WriterType, row_count: usize) -> Self { - Self::Success { writer, row_count } + pub fn success(writer: WriterType, row_count: usize, bytes_written: usize) -> Self { + Self::Success { + writer, + row_count, + bytes_written, + } } pub fn failure(writer: Option, err: DataFusionError) -> Self { @@ -111,9 +119,11 @@ pub(crate) async fn serialize_rb_stream_to_object_store( }); let mut row_count = 0; + let mut bytes_written = 0; while let Some(task) = rx.recv().await { match task.join().await { Ok(Ok((cnt, bytes))) => { + let byte_len = bytes.len(); match writer.write_all(&bytes).await { Ok(_) => (), Err(e) => { @@ -124,6 +134,7 @@ pub(crate) async fn serialize_rb_stream_to_object_store( } }; row_count += cnt; + bytes_written += byte_len; } Ok(Err(e)) => { // Return the writer along with the error @@ -151,10 +162,41 @@ pub(crate) async fn serialize_rb_stream_to_object_store( ); } } - SerializedRecordBatchResult::success(writer, row_count) + SerializedRecordBatchResult::success(writer, row_count, bytes_written) } type FileWriteBundle = (Receiver, SerializerType, WriterType); + +/// Options for stateless file writer tasks. +pub struct StatelessWriterOptions<'a> { + compression: FileCompressionType, + compression_level: Option, + metrics: Option<&'a FileSinkMetrics>, +} + +impl<'a> StatelessWriterOptions<'a> { + /// Create options for the specified compression type. + pub fn new(compression: FileCompressionType) -> Self { + Self { + compression, + compression_level: None, + metrics: None, + } + } + + /// Set the compression level. + pub fn with_compression_level(mut self, compression_level: Option) -> Self { + self.compression_level = compression_level; + self + } + + /// Record common file sink metrics while writing. + pub fn with_metrics(mut self, metrics: &'a FileSinkMetrics) -> Self { + self.metrics = Some(metrics); + self + } +} + /// Contains the common logic for serializing RecordBatches and /// writing the resulting bytes to an ObjectStore. /// Serialization is assumed to be stateless, i.e. @@ -162,9 +204,10 @@ type FileWriteBundle = (Receiver, SerializerType, WriterType); /// dependency on the RecordBatches before or after. pub(crate) async fn stateless_serialize_and_write_files( mut rx: Receiver, - tx: tokio::sync::oneshot::Sender, + tx: tokio::sync::oneshot::Sender<(u64, usize)>, ) -> Result<()> { let mut row_count = 0; + let mut bytes_written = 0; // tracks if any writers encountered an error triggering the need to abort let mut any_errors = false; // tracks the specific error triggering abort @@ -185,9 +228,11 @@ pub(crate) async fn stateless_serialize_and_write_files( SerializedRecordBatchResult::Success { writer, row_count: cnt, + bytes_written: bytes, } => { finished_writers.push(writer); row_count += cnt; + bytes_written += bytes; } SerializedRecordBatchResult::Failure { writer, err } => { finished_writers.extend(writer); @@ -233,7 +278,7 @@ pub(crate) async fn stateless_serialize_and_write_files( } } - tx.send(row_count as u64).map_err(|_| { + tx.send((row_count as u64, bytes_written)).map_err(|_| { internal_datafusion_err!( "Error encountered while sending row count back to file sink!" ) @@ -251,8 +296,32 @@ pub async fn spawn_writer_tasks_and_join( compression_level: Option, object_store: Arc, demux_task: SpawnedTask>, + file_stream_rx: DemuxedStreamReceiver, +) -> Result { + spawn_writer_tasks_and_join_with_metrics( + context, + serializer, + StatelessWriterOptions::new(compression) + .with_compression_level(compression_level), + object_store, + demux_task, + file_stream_rx, + ) + .await +} + +/// Like [`spawn_writer_tasks_and_join`], and also records common file sink metrics. +pub async fn spawn_writer_tasks_and_join_with_metrics( + context: &Arc, + serializer: Arc, + options: StatelessWriterOptions<'_>, + object_store: Arc, + demux_task: SpawnedTask>, mut file_stream_rx: DemuxedStreamReceiver, ) -> Result { + let _timer = options + .metrics + .map(|metrics| metrics.elapsed_compute().timer()); let rb_buffer_size = &context .session_config() .options() @@ -265,17 +334,20 @@ pub async fn spawn_writer_tasks_and_join( stateless_serialize_and_write_files(rx_file_bundle, tx_row_cnt).await }); while let Some((location, rb_stream)) = file_stream_rx.recv().await { - let writer = - ObjectWriterBuilder::new(compression, &location, Arc::clone(&object_store)) - .with_buffer_size(Some( - context - .session_config() - .options() - .execution - .objectstore_writer_buffer_size, - )) - .with_compression_level(compression_level) - .build()?; + let writer = ObjectWriterBuilder::new( + options.compression, + &location, + Arc::clone(&object_store), + ) + .with_buffer_size(Some( + context + .session_config() + .options() + .execution + .objectstore_writer_buffer_size, + )) + .with_compression_level(options.compression_level) + .build()?; if tx_file_bundle .send((rb_stream, Arc::clone(&serializer), writer)) @@ -298,8 +370,13 @@ pub async fn spawn_writer_tasks_and_join( r1.map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; r2.map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; - // Return total row count: - rx_row_cnt.await.map_err(|_| { + // Return total row count and update optional metrics: + let (rows_written, bytes_written) = rx_row_cnt.await.map_err(|_| { internal_datafusion_err!("Did not receive row count from write coordinator") - }) + })?; + if let Some(metrics) = options.metrics { + metrics.rows_written().add(rows_written as usize); + metrics.bytes_written().add(bytes_written); + } + Ok(rows_written) } diff --git a/datafusion/sqllogictest/test_files/copy.slt b/datafusion/sqllogictest/test_files/copy.slt index 77977a6afcb11..68b269646a781 100644 --- a/datafusion/sqllogictest/test_files/copy.slt +++ b/datafusion/sqllogictest/test_files/copy.slt @@ -211,6 +211,26 @@ Plan with Metrics 03)----ProjectionExec: expr=[col1@0 as col1, upper(col2@1) as col2_upper], metrics=[output_rows=2, elapsed_compute=, output_bytes=, output_batches=1, expr_0_eval_time=, expr_1_eval_time=] 04)------DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] +# Verify CsvSink exposes rows_written, bytes_written, and 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/table_csv_metrics.csv' STORED AS CSV; +---- +Plan with Metrics +01)DataSinkExec: sink=CsvSink(file_groups=[]), metrics=[elapsed_compute=, bytes_written=, rows_written=2] +02)--SortExec: expr=[col1@0 ASC NULLS LAST], preserve_partitioning=[false], metrics=[output_rows=2, elapsed_compute=, output_bytes=, output_batches=, 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=, output_bytes=, output_batches=1, expr_0_eval_time=, expr_1_eval_time=] +04)------DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] + +# Verify JsonSink exposes rows_written, bytes_written, and 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/table_json_metrics.json' STORED AS JSON; +---- +Plan with Metrics +01)DataSinkExec: sink=JsonSink(file_groups=[]), metrics=[elapsed_compute=, bytes_written=, rows_written=2] +02)--SortExec: expr=[col1@0 ASC NULLS LAST], preserve_partitioning=[false], metrics=[output_rows=2, elapsed_compute=, output_bytes=, output_batches=, 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=, output_bytes=, output_batches=1, expr_0_eval_time=, expr_1_eval_time=] +04)------DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] + # Copy to directory as partitioned files with keep_partition_by_columns enabled query I COPY (values ('1', 'a'), ('2', 'b'), ('3', 'c')) TO 'test_files/scratch/copy/partitioned_table4/' STORED AS parquet PARTITIONED BY (column1)