feat(parquet): report decompression time and bytes in ArrowReaderMetrics

parquet/src/arrow/arrow_reader/metrics.rs

@@ -19,6 +19,9 @@
 
 use std::sync::Arc;
 use std::sync::atomic::AtomicUsize;
+use std::time::Duration;
+
+use crate::file::metrics::DecompressionMetrics;
 
 /// This enum represents the state of Arrow reader metrics collection.
 ///
@@ -48,6 +51,67 @@ impl ArrowReaderMetrics {
         Self::Enabled(Arc::new(ArrowReaderMetricsInner::new()))
     }
 
+    /// The decompression counters.
+    ///
+    /// Read the four cells of the page kind by phase matrix through this, for
+    /// example the dictionary time paid while skipping, which is work done for
+    /// rows that were never returned. The totals are also forwarded directly by
+    /// [`Self::dictionary_decompress_time`] and its siblings.
+    ///
+    /// Also the handle to give to a page reader built by hand, via
+    /// [`SerializedPageReader::with_decompression_metrics`]. Disabled when
+    /// these metrics are, so nothing downstream reads a clock.
+    ///
+    /// [`SerializedPageReader::with_decompression_metrics`]:
+    ///     crate::file::serialized_reader::SerializedPageReader::with_decompression_metrics
+    pub fn decompression(&self) -> DecompressionMetrics {
+        match self {
+            Self::Disabled => DecompressionMetrics::disabled(),
+            Self::Enabled(inner) => inner.decompression.clone(),
+        }
+    }
+
+    /// Time spent decompressing dictionary pages, across every column read.
+    ///
+    /// A dictionary page is decompressed once per column chunk, however few of
+    /// its values are later decoded, so this is the cost a selective read pays
+    /// before it looks at anything. Decompression happens inside decoding and
+    /// inside skipping, so this is a share of the scan's time rather than a
+    /// stage to add beside it.
+    ///
+    /// Returns None if metrics are disabled.
+    pub fn dictionary_decompress_time(&self) -> Option<Duration> {
+        self.decompression_view()?.dictionary_decompress_time()
+    }
+
+    /// Time spent decompressing data pages, across every column read.
+    ///
+    /// Returns None if metrics are disabled.
+    pub fn data_decompress_time(&self) -> Option<Duration> {
+        self.decompression_view()?.data_decompress_time()
+    }
+
+    /// Uncompressed bytes produced from dictionary pages.
+    ///
+    /// Returns None if metrics are disabled.
+    pub fn dictionary_decompressed_bytes(&self) -> Option<u64> {
+        self.decompression_view()?.dictionary_decompressed_bytes()
+    }
+
+    /// Uncompressed bytes produced from data pages.
+    ///
+    /// Returns None if metrics are disabled.
+    pub fn data_decompressed_bytes(&self) -> Option<u64> {
+        self.decompression_view()?.data_decompressed_bytes()
+    }
+
+    fn decompression_view(&self) -> Option<&DecompressionMetrics> {
+        match self {
+            Self::Disabled => None,
+            Self::Enabled(inner) => Some(&inner.decompression),
+        }
+    }
+
     /// Predicate Cache: number of records read directly from the inner reader
     ///
     /// This is the total number of records read from the inner reader (that is
@@ -122,6 +186,8 @@ pub struct ArrowReaderMetricsInner {
     records_read_from_inner: AtomicUsize,
     /// Total number of records read from previously cached pages
     records_read_from_cache: AtomicUsize,
+    /// Decompression timings, shared with every page reader of this scan
+    decompression: DecompressionMetrics,
 }
 
 impl ArrowReaderMetricsInner {
@@ -130,6 +196,7 @@ impl ArrowReaderMetricsInner {
         Self {
             records_read_from_inner: AtomicUsize::new(0),
             records_read_from_cache: AtomicUsize::new(0),
+            decompression: DecompressionMetrics::enabled(),
         }
     }
 }

parquet/src/arrow/arrow_reader/mod.rs

@@ -1437,6 +1437,7 @@ impl<T: ChunkReader + 'static> ParquetRecordBatchReaderBuilder<T> {
             reader: Arc::new(input.0),
             metadata,
             row_groups,
+            decompression: metrics.decompression(),
         };
 
         let mut plan_builder = ReadPlanBuilder::new(batch_size)
@@ -1480,6 +1481,8 @@ impl<T: ChunkReader + 'static> ParquetRecordBatchReaderBuilder<T> {
 }
 
 struct ReaderRowGroups<T: ChunkReader> {
+    /// Shared with every column's page reader, so one read totals across columns
+    decompression: crate::file::metrics::DecompressionMetrics,
     reader: Arc<T>,
 
     metadata: Arc<ParquetMetaData>,
@@ -1502,6 +1505,7 @@ impl<T: ChunkReader + 'static> RowGroups for ReaderRowGroups<T> {
             reader: self.reader.clone(),
             metadata: self.metadata.clone(),
             row_groups: self.row_groups.clone().into_iter(),
+            decompression: self.decompression.clone(),
         }))
     }
 
@@ -1519,6 +1523,7 @@ impl<T: ChunkReader + 'static> RowGroups for ReaderRowGroups<T> {
 }
 
 struct ReaderPageIterator<T: ChunkReader> {
+    decompression: crate::file::metrics::DecompressionMetrics,
     reader: Arc<T>,
     column_idx: usize,
     row_groups: std::vec::IntoIter<usize>,
@@ -1539,6 +1544,7 @@ impl<T: ChunkReader + 'static> ReaderPageIterator<T> {
         let reader = self.reader.clone();
 
         SerializedPageReader::new(reader, column_chunk_metadata, total_rows, page_locations)?
+            .with_decompression_metrics(self.decompression.clone())
             .add_crypto_context(
                 rg_idx,
                 self.column_idx,

parquet/src/arrow/in_memory_row_group.rs

@@ -22,6 +22,7 @@ use crate::column::page::{PageIterator, PageReader};
 use crate::errors::ParquetError;
 use crate::file::metadata::page_index::RowGroupPageIndex;
 use crate::file::metadata::{ParquetMetaData, RowGroupMetaData};
+use crate::file::metrics::DecompressionMetrics;
 use crate::file::page_index::offset_index::PageLocation;
 use crate::file::reader::{ChunkReader, Length, SerializedPageReader};
 use bytes::{Buf, Bytes};
@@ -37,6 +38,8 @@ pub(crate) struct InMemoryRowGroup<'a> {
     pub(crate) row_count: usize,
     pub(crate) row_group_idx: usize,
     pub(crate) metadata: &'a ParquetMetaData,
+    /// Handed to every column's page reader, so one scan totals across columns
+    pub(crate) decompression: DecompressionMetrics,
 }
 
 /// What ranges to fetch for the columns in this row group
@@ -50,6 +53,12 @@ pub(crate) struct FetchRanges {
 }
 
 impl InMemoryRowGroup<'_> {
+    /// Measures decompression in the page readers this row group hands out.
+    pub(crate) fn with_decompression_metrics(mut self, metrics: DecompressionMetrics) -> Self {
+        self.decompression = metrics;
+        self
+    }
+
     /// Returns the byte ranges to fetch for the columns specified in
     /// `projection` and `selection`.
     ///
@@ -318,7 +327,8 @@ impl RowGroups for InMemoryRowGroup<'_> {
                     column_chunk_metadata,
                     self.row_count,
                     page_locations,
-                )?;
+                )?
+                .with_decompression_metrics(self.decompression.clone());
                 let page_reader = page_reader.add_crypto_context(
                     self.row_group_idx,
                     i,

parquet/src/arrow/push_decoder/reader_builder/data.rs

@@ -22,6 +22,7 @@ use crate::arrow::arrow_reader::RowSelection;
 use crate::arrow::in_memory_row_group::{ColumnChunkData, FetchRanges, InMemoryRowGroup};
 use crate::errors::ParquetError;
 use crate::file::metadata::ParquetMetaData;
+use crate::file::metrics::DecompressionMetrics;
 use crate::file::reader::ChunkReader;
 use crate::util::push_buffers::PushBuffers;
 use bytes::Bytes;
@@ -109,6 +110,7 @@ impl DataRequest {
             page_index,
             row_group_idx,
             metadata: parquet_metadata,
+            decompression: DecompressionMetrics::disabled(),
         };
 
         in_memory_row_group.fill_column_chunks(projection, page_start_offsets, chunks);
@@ -222,6 +224,7 @@ impl<'a> DataRequestBuilder<'a> {
             page_index,
             row_group_idx,
             metadata: parquet_metadata,
+            decompression: DecompressionMetrics::disabled(),
         };
 
         let FetchRanges {

parquet/src/arrow/push_decoder/reader_builder/mod.rs

@@ -559,13 +559,15 @@ impl RowGroupReaderBuilder {
 
                 let predicate = filter_info.current();
 
-                let row_group = data_request.try_into_in_memory_row_group(
-                    row_group_idx,
-                    row_count,
-                    &self.metadata,
-                    predicate.projection(),
-                    &mut self.buffers,
-                )?;
+                let row_group = data_request
+                    .try_into_in_memory_row_group(
+                        row_group_idx,
+                        row_count,
+                        &self.metadata,
+                        predicate.projection(),
+                        &mut self.buffers,
+                    )?
+                    .with_decompression_metrics(self.metrics.decompression());
 
                 let cache_options = filter_info.cache_builder().producer();
 
@@ -736,13 +738,15 @@ impl RowGroupReaderBuilder {
                     budget,
                 } = row_group_info;
 
-                let row_group = data_request.try_into_in_memory_row_group(
-                    row_group_idx,
-                    row_count,
-                    &self.metadata,
-                    &self.projection,
-                    &mut self.buffers,
-                )?;
+                let row_group = data_request
+                    .try_into_in_memory_row_group(
+                        row_group_idx,
+                        row_count,
+                        &self.metadata,
+                        &self.projection,
+                        &mut self.buffers,
+                    )?
+                    .with_decompression_metrics(self.metrics.decompression());
 
                 let plan = plan_builder.build();
 

parquet/src/column/page.rs

@@ -24,6 +24,7 @@ use crate::errors::{ParquetError, Result};
 use crate::file::metadata::thrift::{
     DataPageHeader, DataPageHeaderV2, DictionaryPageHeader, PageHeader,
 };
+use crate::file::metrics::DecompressionPhase;
 use crate::file::statistics::{Statistics, page_stats_to_thrift};
 
 /// Parquet Page definition.
@@ -404,6 +405,20 @@ pub trait PageReader: Iterator<Item = Result<Page>> + Send {
     /// column index information
     fn skip_next_page(&mut self) -> Result<()>;
 
+    /// Tell the reader whether the pages it decompresses next are being read
+    /// or skipped, so decompression cost can be attributed to the right one.
+    ///
+    /// Default is a no-op: an implementation that collects no metrics, and
+    /// every implementation outside this crate, needs no change.
+    fn set_decompression_phase(&mut self, _phase: DecompressionPhase) {}
+
+    /// Tell the reader that values are about to be read from the page it
+    /// loaded last, so a page (and dictionary) decompressed while skipping
+    /// was not decompressed for nothing.
+    ///
+    /// Default is a no-op, like [`Self::set_decompression_phase`].
+    fn current_page_read(&mut self) {}
+
     /// Returns `true` if the next page can be assumed to contain the start of a new record
     ///
     /// Prior to parquet V2 the specification was ambiguous as to whether a single record

parquet/src/column/reader.rs

@@ -27,6 +27,7 @@ use crate::column::reader::decoder::{
 };
 use crate::data_type::*;
 use crate::errors::{ParquetError, Result};
+use crate::file::metrics::DecompressionPhase;
 use crate::schema::types::ColumnDescPtr;
 use crate::util::bit_util::{ceil, num_required_bits, read_num_bytes};
 
@@ -230,6 +231,14 @@ where
     where
         F: FnMut(&mut V::Buffer, usize, usize, Option<&D::Buffer>) -> Result<()>,
     {
+        // Pages decompressed from here on are being decoded, not walked past.
+        self.page_reader
+            .set_decompression_phase(DecompressionPhase::Read);
+        // A page loaded by an earlier skip that still has values is about to be
+        // read from, so its decompression was not wasted.
+        if self.num_decoded_values < self.num_buffered_values {
+            self.page_reader.current_page_read();
+        }
         let mut total_records_read = 0;
         let mut total_levels_read = 0;
         let mut total_values_read = 0;
@@ -310,6 +319,11 @@ where
     ///
     /// Returns the number of records skipped
     pub fn skip_records(&mut self, num_records: usize) -> Result<usize> {
+        // A page skipped whole costs nothing, but a page walked through to
+        // reach a selected row inside it is decompressed and then thrown away.
+        // Attribute that to skipping rather than to reading.
+        self.page_reader
+            .set_decompression_phase(DecompressionPhase::Skip);
         let mut remaining_records = num_records;
         while remaining_records != 0 {
             if self.num_buffered_values == self.num_decoded_values {

parquet/src/file/metadata/dictionary.rs

@@ -35,6 +35,7 @@ use crate::errors::{ParquetError, Result};
 #[cfg(feature = "encryption")]
 use crate::file::metadata::ColumnChunkMetaData;
 use crate::file::metadata::ParquetMetaData;
+use crate::file::metrics::DecompressionRecorder;
 use crate::file::serialized_reader::{
     SerializedPageReaderContext, decode_page, read_page_header_len_from_bytes, verify_page_size,
 };
@@ -117,11 +118,13 @@ pub(crate) fn decode_dictionary_page(
     let page_buf = page_context.decrypt_page_data(page_buf, 0, true)?;
 
     let mut decompressor = create_codec(column_metadata.compression(), &CodecOptions::default())?;
+    // Reading a dictionary for metadata is not part of a scan, so it is not counted.
     let page = decode_page(
         header,
         page_buf,
         column_descriptor.physical_type(),
         decompressor.as_mut(),
+        &mut DecompressionRecorder::default(),
     )?;
     let Page::DictionaryPage {
         buf,