feat(parquet): add bytes_processed scan-completion metric

datafusion/core/tests/sql/explain_analyze.rs

@@ -889,6 +889,7 @@ async fn parquet_explain_analyze() {
     );
     assert_contains!(&formatted, "output_rows_skew=0%");
     assert_contains!(&formatted, "scan_efficiency_ratio=13.99%");
+    assert_contains!(&formatted, "bytes_processed=");
 
     // The order of metrics is expected to be the same as the actual pruning order
     // (file-> row-group -> page)

datafusion/datasource-parquet/src/metrics.rs

@@ -103,6 +103,67 @@ pub struct ParquetFileMetrics {
     pub predicate_cache_records: Gauge,
 }
 
+/// Tracks how much of one file — or one byte range of a file — a scan has
+/// finished with, crediting [`ParquetFileMetrics::bytes_processed`] as it goes.
+///
+/// Every credit is clamped to the bytes left in the budget, and whatever is
+/// left over is credited on drop. The counter therefore advances by exactly the
+/// size of the range being scanned however the scan ends: normally, at a
+/// `LIMIT`, when a dynamic filter proves the rest of the file irrelevant, or on
+/// an error — including one that stops the file being opened at all, which is
+/// why the guard is created before the fallible stages of opening rather than
+/// alongside the decoder. That total is what makes the metric usable as a
+/// completion fraction rather than just another counter.
+///
+/// The clamp and the final top-up also absorb two small inexactnesses in
+/// crediting by row group: a file is slightly larger than the sum of its row
+/// groups (the footer, the page index and any padding belong to no row group),
+/// and a row group is assigned to a byte range by the offset of its first page,
+/// so a range's row groups do not add up to precisely its length.
+///
+/// The budget is held as a `usize` because [`Count`] is, so the two cannot
+/// disagree: on a 32-bit target a range longer than `usize::MAX` saturates once,
+/// here, rather than letting the remaining-byte arithmetic run ahead of what the
+/// counter can record. [`ParquetFileMetrics::bytes_scanned`] has the same
+/// ceiling.
+#[derive(Debug)]
+pub(crate) struct ByteProgress {
+    /// Bytes of the scanned range not yet credited.
+    remaining: usize,
+    bytes_processed: Count,
+}
+
+impl ByteProgress {
+    /// Start tracking a range of `total` bytes.
+    pub(crate) fn new(total: u64, bytes_processed: Count) -> Self {
+        Self {
+            remaining: saturating_usize(total),
+            bytes_processed,
+        }
+    }
+
+    /// Record that the scan is finished with `bytes` more of the range.
+    pub(crate) fn credit(&mut self, bytes: u64) {
+        let bytes = saturating_usize(bytes).min(self.remaining);
+        self.remaining -= bytes;
+        self.bytes_processed.add(bytes);
+    }
+}
+
+/// Narrow a byte count to the width [`Count`] stores, saturating rather than
+/// wrapping. Lossless on 64-bit targets.
+fn saturating_usize(bytes: u64) -> usize {
+    usize::try_from(bytes).unwrap_or(usize::MAX)
+}
+
+impl Drop for ByteProgress {
+    fn drop(&mut self) {
+        let remaining = self.remaining;
+        self.remaining = 0;
+        self.bytes_processed.add(remaining);
+    }
+}
+
 impl ParquetFileMetrics {
     /// Create new metrics
     pub fn new(
@@ -234,6 +295,39 @@ impl ParquetFileMetrics {
         }
     }
 
+    /// The `bytes_processed` counter for one file: the total number of bytes the
+    /// scan is finished with, whether they were read or skipped.
+    ///
+    /// Where [`Self::bytes_scanned`] counts only the bytes fetched from the
+    /// object store, this counts every byte the scan has resolved: the bytes it
+    /// read, plus the bytes of the row groups (and whole files) that pruning
+    /// proved cannot contribute. Over the lifetime of a file it therefore sums
+    /// to that file's size — or, for a file split into byte ranges for
+    /// parallelism, to the size of the range — so
+    /// `bytes_processed / total file bytes` is a scan completion fraction,
+    /// which `bytes_scanned` on its own is not: it understates progress by
+    /// however much pruning and projection pushdown saved.
+    ///
+    /// Credited at row-group granularity: a row group's bytes land when the
+    /// scan is done with it. Crediting a row group's bytes progressively as its
+    /// rows are decoded is left to a follow-up.
+    ///
+    /// Built on demand rather than held on [`ParquetFileMetrics`] because only
+    /// [`ByteProgress`] ever touches it, and a public field would make every
+    /// future metric added here a breaking change for anyone constructing the
+    /// struct with a literal.
+    pub(crate) fn bytes_processed_counter(
+        metrics: &ExecutionPlanMetricsSet,
+        partition: usize,
+        filename: &str,
+    ) -> Count {
+        MetricBuilder::new(metrics)
+            .with_new_label("filename", filename.to_string())
+            .with_type(MetricType::Summary)
+            .with_category(MetricCategory::Bytes)
+            .counter("bytes_processed", partition)
+    }
+
     /// Record pages whose page-index pruning was skipped because the containing
     /// row group was fully matched by row-group statistics.
     ///

datafusion/datasource-parquet/src/opener/early_stop.rs

@@ -38,8 +38,15 @@ pub(super) struct EarlyStoppingStream<S> {
     done: bool,
     file_pruner: FilePruner,
     files_ranges_pruned_statistics: PruningMetrics,
-    /// The inner stream
-    inner: S,
+    /// The inner stream, dropped as soon as this stream is done with it.
+    ///
+    /// Held as an `Option` so finishing releases the decoder — and the buffers
+    /// and per-file metric state it owns — at the moment we stop reading, not
+    /// whenever the caller gets around to dropping this wrapper. Notably the
+    /// scan's byte-progress accounting is completed by that drop, so deferring
+    /// it would leave the file reading as partially scanned after the scan had
+    /// demonstrably finished with it.
+    inner: Option<S>,
 }
 
 impl<S> EarlyStoppingStream<S> {
@@ -50,11 +57,17 @@ impl<S> EarlyStoppingStream<S> {
     ) -> Self {
         Self {
             done: false,
-            inner: stream,
+            inner: Some(stream),
             file_pruner,
             files_ranges_pruned_statistics,
         }
     }
+
+    /// Mark the stream finished and release the inner stream.
+    fn finish(&mut self) {
+        self.done = true;
+        self.inner = None;
+    }
 }
 
 impl<S> EarlyStoppingStream<S>
@@ -70,7 +83,7 @@ where
             self.files_ranges_pruned_statistics.add_pruned(1);
             // Previously this file range has been counted as matched
             self.files_ranges_pruned_statistics.subtract_matched(1);
-            self.done = true;
+            self.finish();
             Ok(None)
         } else {
             // Return the adapted batch
@@ -92,10 +105,13 @@ where
         if self.done {
             return Poll::Ready(None);
         }
-        match ready!(self.inner.poll_next_unpin(cx)) {
+        let Some(inner) = self.inner.as_mut() else {
+            return Poll::Ready(None);
+        };
+        match ready!(inner.poll_next_unpin(cx)) {
             None => {
                 // input done
-                self.done = true;
+                self.finish();
                 Poll::Ready(None)
             }
             Some(input_batch) => {
@@ -105,3 +121,147 @@ where
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use std::sync::Arc;
+    use std::sync::atomic::{AtomicBool, Ordering};
+
+    use arrow::array::{Int32Array, RecordBatch};
+    use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+    use datafusion_common::{
+        ColumnStatistics, ScalarValue, Statistics, stats::Precision,
+    };
+    use datafusion_datasource::PartitionedFile;
+    use datafusion_physical_expr::PhysicalExpr;
+    use datafusion_physical_expr::expressions::{
+        BinaryExpr, Column, DynamicFilterPhysicalExpr, Literal,
+    };
+    use datafusion_physical_plan::metrics::Count;
+    use futures::stream;
+
+    /// An inner stream that records when it is dropped, standing in for the
+    /// decoder whose drop completes the scan's byte accounting.
+    struct DropRecordingStream<S> {
+        inner: S,
+        dropped: Arc<AtomicBool>,
+    }
+
+    impl<S: Stream + Unpin> Stream for DropRecordingStream<S> {
+        type Item = S::Item;
+
+        fn poll_next(
+            mut self: Pin<&mut Self>,
+            cx: &mut Context<'_>,
+        ) -> Poll<Option<Self::Item>> {
+            self.inner.poll_next_unpin(cx)
+        }
+    }
+
+    impl<S> Drop for DropRecordingStream<S> {
+        fn drop(&mut self) {
+            self.dropped.store(true, Ordering::Relaxed);
+        }
+    }
+
+    fn schema() -> SchemaRef {
+        Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]))
+    }
+
+    /// A file whose only column holds values 1..=9, so a predicate demanding
+    /// larger values prunes it outright.
+    fn file_with_stats() -> PartitionedFile {
+        // Built field by field rather than from `Statistics::new_unknown`, which
+        // already seeds one entry per column, so that column 0 carries these
+        // bounds rather than an unknown placeholder.
+        let statistics = Statistics {
+            num_rows: Precision::Absent,
+            total_byte_size: Precision::Absent,
+            column_statistics: vec![
+                ColumnStatistics::new_unknown()
+                    .with_min_value(Precision::Exact(ScalarValue::Int32(Some(1))))
+                    .with_max_value(Precision::Exact(ScalarValue::Int32(Some(9))))
+                    .with_null_count(Precision::Exact(0)),
+            ],
+        };
+        PartitionedFile::new("test.parquet".to_string(), 1_000)
+            .with_statistics(Arc::new(statistics))
+    }
+
+    fn pruning_filter(schema: &SchemaRef) -> FilePruner {
+        let expr: Arc<dyn PhysicalExpr> = Arc::new(BinaryExpr::new(
+            Arc::new(Column::new("a", 0)),
+            datafusion_expr::Operator::Gt,
+            Arc::new(Literal::new(ScalarValue::Int32(Some(100)))),
+        ));
+        let dynamic: Arc<dyn PhysicalExpr> = Arc::new(DynamicFilterPhysicalExpr::new(
+            expr.children().into_iter().map(Arc::clone).collect(),
+            expr,
+        ));
+        FilePruner::try_new(dynamic, schema, &file_with_stats(), Count::new())
+            .expect("file has statistics, so a pruner can be built")
+    }
+
+    fn batch(schema: &SchemaRef) -> RecordBatch {
+        RecordBatch::try_new(
+            Arc::clone(schema),
+            vec![Arc::new(Int32Array::from(vec![1]))],
+        )
+        .unwrap()
+    }
+
+    /// Stopping early must release the inner stream there and then. The decoder's
+    /// drop is what completes this file's byte-progress accounting, so holding it
+    /// until the caller drops the wrapper would leave the scan reporting a file it
+    /// has finished with as still partly unread.
+    #[tokio::test]
+    async fn stopping_early_releases_the_inner_stream() {
+        let schema = schema();
+        let dropped = Arc::new(AtomicBool::new(false));
+        let inner = DropRecordingStream {
+            inner: stream::iter(vec![Ok(batch(&schema)), Ok(batch(&schema))]),
+            dropped: Arc::clone(&dropped),
+        };
+
+        let mut early_stopping = EarlyStoppingStream::new(
+            inner,
+            pruning_filter(&schema),
+            PruningMetrics::new(),
+        );
+
+        assert!(
+            early_stopping.next().await.is_none(),
+            "the filter prunes every row, so the first batch must end the stream",
+        );
+        assert!(
+            dropped.load(Ordering::Relaxed),
+            "the inner stream must be released when the scan stops, not when the \
+             wrapper is eventually dropped",
+        );
+    }
+
+    /// The same must hold when the inner stream simply runs out.
+    #[tokio::test]
+    async fn exhausting_the_inner_stream_releases_it() {
+        let schema = schema();
+        let dropped = Arc::new(AtomicBool::new(false));
+        let inner = DropRecordingStream {
+            inner: stream::iter(Vec::<Result<RecordBatch>>::new()),
+            dropped: Arc::clone(&dropped),
+        };
+
+        let mut early_stopping = EarlyStoppingStream::new(
+            inner,
+            pruning_filter(&schema),
+            PruningMetrics::new(),
+        );
+
+        assert!(early_stopping.next().await.is_none());
+        assert!(
+            dropped.load(Ordering::Relaxed),
+            "an exhausted inner stream must be released too",
+        );
+    }
+}