feat: add metrics tracking for file sinks and update orchestration fo…

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<MetricsSet> {
+        Some(self.metrics.snapshot())
+    }
+
     fn schema(&self) -> &SchemaRef {
         self.config.output_schema()
     }

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<dyn ObjectStore>,
     ) -> Result<u64> {
         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<MetricsSet> {
+        Some(self.metrics.snapshot())
+    }
+
     fn schema(&self) -> &SchemaRef {
         self.config.output_schema()
     }

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<Vec<SortingColumn>>,
     /// 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<dyn ObjectStore>,
     ) -> Result<u64> {
-        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<MetricsSet> {
-        Some(self.metrics.clone_inner())
+        Some(self.metrics.snapshot())
     }
 
     fn schema(&self) -> &SchemaRef {

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 {