fix: fix elapsed_compute metric in ParquetSink to report encoding time only (#21825)

## Which issue does this PR close?

- Closes #21797.

## Rationale for this change

`ParquetSink` registered an `elapsed_compute` metric using a single
wall-clock timer that spanned the entire write operation — upstream
batch wait, CPU Arrow→Parquet encoding, and object-store I/O all rolled
into one number. This
made the metric misleading: it inflated `elapsed_compute` with I/O
latency, which is inconsistent with how every other operator in
DataFusion reports this metric (CPU time only).

## What changes are included in this PR?

Two write paths are fixed independently:

**Sequential path** (`allow_single_file_parallelism = false` or CDC
enabled):
- A new `TimingWriter<W>` wrapper implements `AsyncFileWriter` and
records wall-clock time spent in I/O calls (`write` / `complete`).
- The total time inside `writer.write()` and `writer.close()` is
accumulated in `total_write_time`. After all tasks join,
`elapsed_compute` is set to `total_write_time − io_time`, isolating pure
Arrow→Parquet encoding time.

**Parallel path** (`allow_single_file_parallelism = true`, default):
- `encoding_time: Time` (a clone of the registered `elapsed_compute`
metric) is threaded through the five-function call chain down to the two
leaf sites: `writer.write()` in `column_serializer_task` and
`writer.close()` in `spawn_rg_join_and_finalize_task`. Since `Time` is
`Arc<AtomicUsize>`, all concurrent column tasks accumulate directly into
the registered metric.
- Note: on the parallel path, `append_to_row_group()` in
`concatenate_parallel_row_groups` is interleaved with I/O and cannot be
cleanly isolated. It is excluded from `elapsed_compute`. This is
acceptable since it operates on already-encoded data and represents a
small fraction of total encoding CPU time.

## Are these changes tested?

Yes. Two tests are added/extended in
`datafusion/core/src/dataframe/parquet.rs`:
- `test_parquet_sink_metrics_sequential` (new): verifies
`elapsed_compute > 0` with `allow_single_file_parallelism = false`.
- `test_parquet_sink_metrics_parallel`: extended with an
`elapsed_compute > 0` assertion (was previously missing for the parallel
path).

The existing `test_parquet_sink_metrics` test (parallel path, default
config) already asserted `elapsed_compute > 0` and continues to pass.

## Are there any user-facing changes?

No API changes. The `elapsed_compute` metric was already surfaced — this
PR makes its value accurate rather than introducing a new metric.

Author: fred1268

Cargo.lock

@@ -183,6 +183,7 @@ parquet = { version = "58.1.0", default-features = false, features = [
 ] }
 pbjson = { version = "0.9.0" }
 pbjson-types = "0.9"
+pin-project = "1"
 # Should match arrow-flight's version of prost.
 prost = "0.14.1"
 rand = "0.9"

Cargo.toml

@@ -558,6 +558,66 @@ mod tests {
 
         assert_eq!(metric_usize(&aggregated, "rows_written"), 50);
         assert!(metric_usize(&aggregated, "bytes_written") > 0);
+        assert!(
+            metric_usize(&aggregated, "elapsed_compute") > 0,
+            "expected elapsed_compute > 0 on parallel path"
+        );
+
+        Ok(())
+    }
+
+    /// Test that ParquetSink reports a non-zero elapsed_compute on the sequential
+    /// write path (allow_single_file_parallelism = false), where elapsed_compute
+    /// is computed as total_write_time - io_time via TimingWriter.
+    #[tokio::test]
+    async fn test_parquet_sink_metrics_sequential() -> Result<()> {
+        use arrow::array::Int32Array;
+        use arrow::datatypes::{DataType, Field, Schema};
+        use arrow::record_batch::RecordBatch;
+        use datafusion_execution::TaskContext;
+
+        use futures::TryStreamExt;
+
+        let ctx = SessionContext::new();
+        ctx.sql("SET datafusion.execution.parquet.allow_single_file_parallelism = false")
+            .await?
+            .collect()
+            .await?;
+
+        let tmp_dir = TempDir::new()?;
+        let output_path = tmp_dir.path().join("metrics_sequential.parquet");
+        let output_path_str = output_path.to_str().unwrap();
+
+        let schema =
+            Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
+        let ids: Vec<i32> = (0..50).collect();
+        let batch = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![Arc::new(Int32Array::from(ids))],
+        )?;
+        ctx.register_batch("source_seq", batch)?;
+
+        let df = ctx
+            .sql(&format!(
+                "COPY source_seq TO '{output_path_str}' STORED AS PARQUET"
+            ))
+            .await?;
+        let plan = df.create_physical_plan().await?;
+        let task_ctx = Arc::new(TaskContext::from(&ctx.state()));
+        let stream = plan.execute(0, task_ctx)?;
+        let _batches: Vec<_> = stream.try_collect().await?;
+
+        let metrics = plan
+            .metrics()
+            .expect("DataSinkExec should return metrics from ParquetSink");
+        let aggregated = metrics.aggregate_by_name();
+
+        assert_eq!(metric_usize(&aggregated, "rows_written"), 50);
+        assert!(metric_usize(&aggregated, "bytes_written") > 0);
+        assert!(
+            metric_usize(&aggregated, "elapsed_compute") > 0,
+            "expected elapsed_compute > 0 on sequential path"
+        );
 
         Ok(())
     }

datafusion/core/src/dataframe/parquet.rs

@@ -55,7 +55,8 @@ use datafusion_execution::{SendableRecordBatchStream, TaskContext};
 use datafusion_expr::dml::InsertOp;
 use datafusion_physical_expr_common::sort_expr::{LexOrdering, LexRequirement};
 use datafusion_physical_plan::metrics::{
-    ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, MetricsSet,
+    ElapsedComputeFutureExt, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory,
+    MetricsSet, Time,
 };
 use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan};
 use datafusion_session::Session;
@@ -1340,8 +1341,6 @@ impl FileSink for ParquetSink {
             .global_counter("bytes_written");
         let elapsed_compute = MetricBuilder::new(&self.metrics).elapsed_compute(0);
 
-        let write_start = datafusion_common::instant::Instant::now();
-
         let parquet_opts = &self.parquet_options;
 
         let mut file_write_tasks: JoinSet<
@@ -1375,17 +1374,20 @@ impl FileSink for ParquetSink {
                     .await?;
                 let reservation = MemoryConsumer::new(format!("ParquetSink[{path}]"))
                     .register(context.memory_pool());
-                file_write_tasks.spawn(async move {
-                    while let Some(batch) = rx.recv().await {
-                        writer.write(&batch).await?;
-                        reservation.try_resize(writer.memory_size())?;
+                file_write_tasks.spawn(
+                    async move {
+                        while let Some(batch) = rx.recv().await {
+                            writer.write(&batch).await?;
+                            reservation.try_resize(writer.memory_size())?;
+                        }
+                        let parquet_meta_data = writer
+                            .close()
+                            .await
+                            .map_err(|e| DataFusionError::ParquetError(Box::new(e)))?;
+                        Ok((path, parquet_meta_data))
                     }
-                    let parquet_meta_data = writer
-                        .close()
-                        .await
-                        .map_err(|e| DataFusionError::ParquetError(Box::new(e)))?;
-                    Ok((path, parquet_meta_data))
-                });
+                    .with_elapsed_compute(elapsed_compute.clone()),
+                );
             } else {
                 let writer = ObjectWriterBuilder::new(
                     // Parquet files as a whole are never compressed, since they
@@ -1402,20 +1404,20 @@ impl FileSink for ParquetSink {
                         .objectstore_writer_buffer_size,
                 ))
                 .build()?;
-                let schema = get_writer_schema(&self.config);
-                let props = parquet_props.clone();
-                let skip_arrow_metadata = self.parquet_options.global.skip_arrow_metadata;
-                let parallel_options_clone = parallel_options.clone();
-                let pool = Arc::clone(context.memory_pool());
+                let ctx = ParquetFileWriteContext {
+                    schema: get_writer_schema(&self.config),
+                    props: Arc::new(parquet_props),
+                    skip_arrow_metadata: self.parquet_options.global.skip_arrow_metadata,
+                    parallel_options: Arc::new(parallel_options.clone()),
+                    pool: Arc::clone(context.memory_pool()),
+                };
+                let encoding_time = elapsed_compute.clone();
                 file_write_tasks.spawn(async move {
                     let parquet_meta_data = output_single_parquet_file_parallelized(
                         writer,
                         rx,
-                        schema,
-                        &props,
-                        skip_arrow_metadata,
-                        parallel_options_clone,
-                        pool,
+                        ctx,
+                        encoding_time,
                     )
                     .await?;
                     Ok((path, parquet_meta_data))
@@ -1456,8 +1458,6 @@ impl FileSink for ParquetSink {
             .await
             .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??;
 
-        elapsed_compute.add_elapsed(write_start);
-
         Ok(rows_written_counter.value() as u64)
     }
 }
@@ -1487,8 +1487,10 @@ async fn column_serializer_task(
     mut rx: Receiver<ArrowLeafColumn>,
     mut writer: ArrowColumnWriter,
     reservation: MemoryReservation,
+    encoding_time: Time,
 ) -> Result<(ArrowColumnWriter, MemoryReservation)> {
     while let Some(col) = rx.recv().await {
+        let _timer = encoding_time.timer();
         writer.write(&col)?;
         reservation.try_resize(writer.memory_size())?;
     }
@@ -1505,6 +1507,7 @@ fn spawn_column_parallel_row_group_writer(
     col_writers: Vec<ArrowColumnWriter>,
     max_buffer_size: usize,
     pool: &Arc<dyn MemoryPool>,
+    encoding_time: &Time,
 ) -> Result<(Vec<ColumnWriterTask>, Vec<ColSender>)> {
     let num_columns = col_writers.len();
 
@@ -1522,6 +1525,7 @@ fn spawn_column_parallel_row_group_writer(
             receive_array,
             writer,
             reservation,
+            encoding_time.clone(),
         ));
         col_writer_tasks.push(task);
     }
@@ -1536,6 +1540,22 @@ struct ParallelParquetWriterOptions {
     max_buffered_record_batches_per_stream: usize,
 }
 
+/// Write configuration inputs shared across all parallel tasks that encode a
+/// single Parquet file. These values are invariant for the duration of one file
+/// write and do not change per row-group or per column.
+///
+/// Separating these from per-call parameters (`object_store_writer`, `data`,
+/// `encoding_time`) keeps the deep parallel call chain below the argument-count
+/// limit without mixing configuration with runtime state.
+#[derive(Clone)]
+struct ParquetFileWriteContext {
+    schema: Arc<Schema>,
+    props: Arc<WriterProperties>,
+    skip_arrow_metadata: bool,
+    parallel_options: Arc<ParallelParquetWriterOptions>,
+    pool: Arc<dyn MemoryPool>,
+}
+
 /// This is the return type of calling [ArrowColumnWriter].close() on each column
 /// i.e. the Vec of encoded columns which can be appended to a row group
 type RBStreamSerializeResult = Result<(Vec<ArrowColumnChunk>, MemoryReservation, usize)>;
@@ -1570,6 +1590,7 @@ fn spawn_rg_join_and_finalize_task(
     column_writer_tasks: Vec<ColumnWriterTask>,
     rg_rows: usize,
     pool: &Arc<dyn MemoryPool>,
+    encoding_time: Time,
 ) -> SpawnedTask<RBStreamSerializeResult> {
     let rg_reservation =
         MemoryConsumer::new("ParquetSink(SerializedRowGroupWriter)").register(pool);
@@ -1584,6 +1605,7 @@ fn spawn_rg_join_and_finalize_task(
                 .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??;
             let encoded_size = writer.get_estimated_total_bytes();
             rg_reservation.grow(encoded_size);
+            let _timer = encoding_time.timer();
             finalized_rg.push(writer.close()?);
         }
 
@@ -1603,21 +1625,25 @@ fn spawn_parquet_parallel_serialization_task(
     row_group_writer_factory: ArrowRowGroupWriterFactory,
     mut data: Receiver<RecordBatch>,
     serialize_tx: Sender<SpawnedTask<RBStreamSerializeResult>>,
-    schema: Arc<Schema>,
-    writer_props: Arc<WriterProperties>,
-    parallel_options: Arc<ParallelParquetWriterOptions>,
-    pool: Arc<dyn MemoryPool>,
+    ctx: ParquetFileWriteContext,
+    encoding_time: Time,
 ) -> SpawnedTask<Result<(), DataFusionError>> {
     SpawnedTask::spawn(async move {
-        let max_buffer_rb = parallel_options.max_buffered_record_batches_per_stream;
-        let max_row_group_rows = writer_props
+        let max_buffer_rb = ctx.parallel_options.max_buffered_record_batches_per_stream;
+        let max_row_group_rows = ctx
+            .props
             .max_row_group_row_count()
             .unwrap_or(DEFAULT_MAX_ROW_GROUP_ROW_COUNT);
         let mut row_group_index = 0;
         let col_writers =
             row_group_writer_factory.create_column_writers(row_group_index)?;
         let (mut column_writer_handles, mut col_array_channels) =
-            spawn_column_parallel_row_group_writer(col_writers, max_buffer_rb, &pool)?;
+            spawn_column_parallel_row_group_writer(
+                col_writers,
+                max_buffer_rb,
+                &ctx.pool,
+                &encoding_time,
+            )?;
         let mut current_rg_rows = 0;
 
         while let Some(mut rb) = data.recv().await {
@@ -1629,7 +1655,7 @@ fn spawn_parquet_parallel_serialization_task(
                     send_arrays_to_col_writers(
                         &col_array_channels,
                         &rb,
-                        Arc::clone(&schema),
+                        Arc::clone(&ctx.schema),
                     )
                     .await?;
                     current_rg_rows += rb.num_rows();
@@ -1640,7 +1666,7 @@ fn spawn_parquet_parallel_serialization_task(
                     send_arrays_to_col_writers(
                         &col_array_channels,
                         &a,
-                        Arc::clone(&schema),
+                        Arc::clone(&ctx.schema),
                     )
                     .await?;
 
@@ -1651,7 +1677,8 @@ fn spawn_parquet_parallel_serialization_task(
                     let finalize_rg_task = spawn_rg_join_and_finalize_task(
                         column_writer_handles,
                         max_row_group_rows,
-                        &pool,
+                        &ctx.pool,
+                        encoding_time.clone(),
                     );
 
                     // Do not surface error from closed channel (means something
@@ -1670,7 +1697,8 @@ fn spawn_parquet_parallel_serialization_task(
                         spawn_column_parallel_row_group_writer(
                             col_writers,
                             max_buffer_rb,
-                            &pool,
+                            &ctx.pool,
+                            &encoding_time,
                         )?;
                 }
             }
@@ -1682,7 +1710,8 @@ fn spawn_parquet_parallel_serialization_task(
             let finalize_rg_task = spawn_rg_join_and_finalize_task(
                 column_writer_handles,
                 current_rg_rows,
-                &pool,
+                &ctx.pool,
+                encoding_time.clone(),
             );
 
             // Do not surface error from closed channel (means something
@@ -1749,37 +1778,32 @@ async fn concatenate_parallel_row_groups(
 async fn output_single_parquet_file_parallelized(
     object_store_writer: Box<dyn AsyncWrite + Send + Unpin>,
     data: Receiver<RecordBatch>,
-    output_schema: Arc<Schema>,
-    parquet_props: &WriterProperties,
-    skip_arrow_metadata: bool,
-    parallel_options: ParallelParquetWriterOptions,
-    pool: Arc<dyn MemoryPool>,
+    ctx: ParquetFileWriteContext,
+    encoding_time: Time,
 ) -> Result<ParquetMetaData> {
-    let max_rowgroups = parallel_options.max_parallel_row_groups;
+    let max_rowgroups = ctx.parallel_options.max_parallel_row_groups;
     // Buffer size of this channel limits maximum number of RowGroups being worked on in parallel
     let (serialize_tx, serialize_rx) =
         mpsc::channel::<SpawnedTask<RBStreamSerializeResult>>(max_rowgroups);
 
-    let arc_props = Arc::new(parquet_props.clone());
     let merged_buff = SharedBuffer::new(INITIAL_BUFFER_BYTES);
     let options = ArrowWriterOptions::new()
-        .with_properties(parquet_props.clone())
-        .with_skip_arrow_metadata(skip_arrow_metadata);
+        .with_properties((*ctx.props).clone())
+        .with_skip_arrow_metadata(ctx.skip_arrow_metadata);
     let writer = ArrowWriter::try_new_with_options(
         merged_buff.clone(),
-        Arc::clone(&output_schema),
+        Arc::clone(&ctx.schema),
         options,
     )?;
     let (writer, row_group_writer_factory) = writer.into_serialized_writer()?;
 
+    let pool = Arc::clone(&ctx.pool);
     let launch_serialization_task = spawn_parquet_parallel_serialization_task(
         row_group_writer_factory,
         data,
         serialize_tx,
-        Arc::clone(&output_schema),
-        Arc::clone(&arc_props),
-        parallel_options.into(),
-        Arc::clone(&pool),
+        ctx,
+        encoding_time,
     );
     let parquet_meta_data = concatenate_parallel_row_groups(
         writer,

datafusion/datasource-parquet/src/file_format.rs

@@ -49,6 +49,7 @@ hashbrown = { workspace = true }
 indexmap = { workspace = true }
 itertools = { workspace = true }
 parking_lot = { workspace = true }
+pin-project = { workspace = true }
 
 [dev-dependencies]
 criterion = { workspace = true }