Lazy register parquet file metrics

Author: xudong963

datafusion/core/tests/parquet/external_access_plan.rs

@@ -84,7 +84,7 @@ async fn skip_all() {
     .await;
 
     // Verify that skipping all row groups skips reading any data at all
-    let bytes_scanned = metric_value(&parquet_metrics, "bytes_scanned").unwrap();
+    let bytes_scanned = metric_value(&parquet_metrics, "bytes_scanned").unwrap_or(0);
     assert_eq!(bytes_scanned, 0, "metrics : {parquet_metrics:#?}",);
 }
 

datafusion/core/tests/parquet/filter_pushdown.rs

@@ -544,9 +544,17 @@ impl<'a> TestCase<'a> {
             PushdownExpected::None
         };
 
-        let pushdown_rows_pruned = get_value(&metrics, "pushdown_rows_pruned");
+        let pushdown_rows_pruned = match pushdown_expected {
+            PushdownExpected::None => {
+                metric_value_or_zero(&metrics, "pushdown_rows_pruned")
+            }
+            PushdownExpected::Some => {
+                expect_metric_value(&metrics, "pushdown_rows_pruned")
+            }
+        };
+        let pushdown_rows_matched =
+            metric_value_or_zero(&metrics, "pushdown_rows_matched");
         println!("  pushdown_rows_pruned: {pushdown_rows_pruned}");
-        let pushdown_rows_matched = get_value(&metrics, "pushdown_rows_matched");
         println!("  pushdown_rows_matched: {pushdown_rows_matched}");
 
         match pushdown_expected {
@@ -562,11 +570,6 @@ impl<'a> TestCase<'a> {
             }
         };
 
-        let (page_index_rows_pruned, page_index_rows_matched) =
-            get_pruning_metrics(&metrics, "page_index_rows_pruned");
-        println!(" page_index_rows_pruned: {page_index_rows_pruned}");
-        println!(" page_index_rows_matched: {page_index_rows_matched}");
-
         let page_index_filtering_expected = if scan_options.enable_page_index {
             self.page_index_filtering_expected
         } else {
@@ -577,9 +580,17 @@ impl<'a> TestCase<'a> {
 
         match page_index_filtering_expected {
             PageIndexFilteringExpected::None => {
+                let (page_index_rows_pruned, page_index_rows_matched) =
+                    pruning_metrics_or_zero(&metrics, "page_index_rows_pruned");
+                println!(" page_index_rows_pruned: {page_index_rows_pruned}");
+                println!(" page_index_rows_matched: {page_index_rows_matched}");
                 assert_eq!(page_index_rows_pruned, 0);
             }
             PageIndexFilteringExpected::Some => {
+                let (page_index_rows_pruned, page_index_rows_matched) =
+                    expect_pruning_metrics(&metrics, "page_index_rows_pruned");
+                println!(" page_index_rows_pruned: {page_index_rows_pruned}");
+                println!(" page_index_rows_matched: {page_index_rows_matched}");
                 assert!(
                     page_index_rows_pruned > 0,
                     "Expected to filter rows via page index but none were",
@@ -591,29 +602,55 @@ impl<'a> TestCase<'a> {
     }
 }
 
-fn get_pruning_metrics(metrics: &MetricsSet, metric_name: &str) -> (usize, usize) {
+fn pruning_metrics_or_zero(metrics: &MetricsSet, metric_name: &str) -> (usize, usize) {
+    get_pruning_metrics(metrics, metric_name, true)
+}
+
+fn expect_pruning_metrics(metrics: &MetricsSet, metric_name: &str) -> (usize, usize) {
+    get_pruning_metrics(metrics, metric_name, false)
+}
+
+fn get_pruning_metrics(
+    metrics: &MetricsSet,
+    metric_name: &str,
+    allow_missing: bool,
+) -> (usize, usize) {
     match metrics.sum_by_name(metric_name) {
         Some(MetricValue::PruningMetrics {
             pruning_metrics, ..
         }) => (pruning_metrics.pruned(), pruning_metrics.matched()),
         Some(_) => {
             panic!("Metric '{metric_name}' is not a pruning metric in\n\n{metrics:#?}")
         }
-        None => panic!(
-            "Expected metric not found. Looking for '{metric_name}' in\n\n{metrics:#?}"
-        ),
+        None if allow_missing => (0, 0),
+        None => panic!("Expected metric '{metric_name}' not found in\n\n{metrics:#?}"),
     }
 }
 
-fn get_value(metrics: &MetricsSet, metric_name: &str) -> usize {
+fn metric_value_or_zero(metrics: &MetricsSet, metric_name: &str) -> usize {
+    get_value(metrics, metric_name, true)
+}
+
+fn expect_metric_value(metrics: &MetricsSet, metric_name: &str) -> usize {
+    get_value(metrics, metric_name, false)
+}
+
+fn metric_value_for_expected(
+    metrics: &MetricsSet,
+    metric_name: &str,
+    expected: usize,
+) -> usize {
+    get_value(metrics, metric_name, expected == 0)
+}
+
+fn get_value(metrics: &MetricsSet, metric_name: &str, allow_missing: bool) -> usize {
     match metrics.sum_by_name(metric_name) {
         Some(MetricValue::PruningMetrics {
             pruning_metrics, ..
         }) => pruning_metrics.pruned(),
         Some(v) => v.as_usize(),
-        None => panic!(
-            "Expected metric not found. Looking for '{metric_name}' in\n\n{metrics:#?}"
-        ),
+        None if allow_missing => 0,
+        None => panic!("Expected metric '{metric_name}' not found in\n\n{metrics:#?}"),
     }
 }
 
@@ -736,11 +773,19 @@ impl PredicateCacheTest {
 
         // verify the predicate cache metrics
         assert_eq!(
-            get_value(&metrics, "predicate_cache_inner_records"),
+            metric_value_for_expected(
+                &metrics,
+                "predicate_cache_inner_records",
+                expected_inner_records
+            ),
             expected_inner_records
         );
         assert_eq!(
-            get_value(&metrics, "predicate_cache_records"),
+            metric_value_for_expected(
+                &metrics,
+                "predicate_cache_records",
+                expected_records
+            ),
             expected_records
         );
         Ok(())

datafusion/core/tests/parquet/mod.rs

@@ -267,6 +267,20 @@ impl TestOutput {
     }
 }
 
+pub(crate) fn zero_if_metric_absent(
+    actual: Option<usize>,
+    expected: Option<usize>,
+) -> Option<usize> {
+    // Lazy parquet metrics may represent zero as an absent metric. Tests that
+    // specifically enforce lazy-skipping zero metrics live with the metrics
+    // implementation; these integration tests only compare scan behavior.
+    if expected == Some(0) {
+        Some(actual.unwrap_or(0))
+    } else {
+        actual
+    }
+}
+
 /// Creates an execution context that has an external table "t"
 /// registered pointing at a parquet file made with `make_test_file`
 /// and the appropriate scenario

datafusion/core/tests/parquet/page_pruning.rs

@@ -18,7 +18,7 @@
 use std::sync::Arc;
 
 use crate::parquet::Unit::Page;
-use crate::parquet::{ContextWithParquet, Scenario};
+use crate::parquet::{ContextWithParquet, Scenario, zero_if_metric_absent};
 
 use arrow::array::{Int32Array, RecordBatch};
 use arrow::datatypes::{DataType, Field, Schema};
@@ -238,8 +238,14 @@ async fn test_prune(
             .await;
 
     println!("{}", output.description());
-    assert_eq!(output.predicate_evaluation_errors(), expected_errors);
-    assert_eq!(output.row_pages_pruned(), expected_row_pages_pruned);
+    assert_eq!(
+        zero_if_metric_absent(output.predicate_evaluation_errors(), expected_errors),
+        expected_errors
+    );
+    assert_eq!(
+        zero_if_metric_absent(output.row_pages_pruned(), expected_row_pages_pruned),
+        expected_row_pages_pruned
+    );
     assert_eq!(
         output.result_rows,
         expected_results,
@@ -360,7 +366,10 @@ async fn prune_date64() {
 
     println!("{}", output.description());
     // This should prune out groups  without error
-    assert_eq!(output.predicate_evaluation_errors(), Some(0));
+    assert_eq!(
+        zero_if_metric_absent(output.predicate_evaluation_errors(), Some(0)),
+        Some(0)
+    );
     assert_eq!(output.row_pages_pruned(), Some(15));
     assert_eq!(output.result_rows, 1, "{}", output.description());
 }

datafusion/core/tests/parquet/row_group_pruning.rs

@@ -27,7 +27,7 @@ use datafusion_common::{DataFusionError, ScalarValue};
 use itertools::Itertools;
 
 use crate::parquet::Unit::RowGroup;
-use crate::parquet::{ContextWithParquet, Scenario};
+use crate::parquet::{ContextWithParquet, Scenario, zero_if_metric_absent};
 use datafusion_expr::{col, lit};
 struct RowGroupPruningTest {
     scenario: Scenario,
@@ -136,33 +136,51 @@ impl RowGroupPruningTest {
 
         println!("{}", output.description());
         assert_eq!(
-            output.predicate_evaluation_errors(),
+            zero_if_metric_absent(
+                output.predicate_evaluation_errors(),
+                self.expected_errors
+            ),
             self.expected_errors,
             "mismatched predicate_evaluation error"
         );
         assert_eq!(
-            output.row_groups_matched_statistics(),
+            zero_if_metric_absent(
+                output.row_groups_matched_statistics(),
+                self.expected_row_group_matched_by_statistics
+            ),
             self.expected_row_group_matched_by_statistics,
             "mismatched row_groups_matched_statistics",
         );
         assert_eq!(
-            output.row_groups_pruned_statistics(),
+            zero_if_metric_absent(
+                output.row_groups_pruned_statistics(),
+                self.expected_row_group_pruned_by_statistics
+            ),
             self.expected_row_group_pruned_by_statistics,
             "mismatched row_groups_pruned_statistics",
         );
         assert_eq!(
-            output.files_ranges_pruned_statistics(),
+            zero_if_metric_absent(
+                output.files_ranges_pruned_statistics(),
+                self.expected_files_pruned_by_statistics
+            ),
             self.expected_files_pruned_by_statistics,
             "mismatched files_ranges_pruned_statistics",
         );
         let bloom_filter_metrics = output.row_groups_bloom_filter();
         assert_eq!(
-            bloom_filter_metrics.as_ref().map(|pm| pm.total_matched()),
+            zero_if_metric_absent(
+                bloom_filter_metrics.as_ref().map(|pm| pm.total_matched()),
+                self.expected_row_group_matched_by_bloom_filter
+            ),
             self.expected_row_group_matched_by_bloom_filter,
             "mismatched row_groups_matched_bloom_filter",
         );
         assert_eq!(
-            bloom_filter_metrics.map(|pm| pm.total_pruned()),
+            zero_if_metric_absent(
+                bloom_filter_metrics.map(|pm| pm.total_pruned()),
+                self.expected_row_group_pruned_by_bloom_filter
+            ),
             self.expected_row_group_pruned_by_bloom_filter,
             "mismatched row_groups_pruned_bloom_filter",
         );
@@ -196,32 +214,50 @@ impl RowGroupPruningTest {
 
         println!("{}", output.description());
         assert_eq!(
-            output.predicate_evaluation_errors(),
+            zero_if_metric_absent(
+                output.predicate_evaluation_errors(),
+                self.expected_errors
+            ),
             self.expected_errors,
             "mismatched predicate_evaluation error"
         );
         assert_eq!(
-            output.row_groups_matched_statistics(),
+            zero_if_metric_absent(
+                output.row_groups_matched_statistics(),
+                self.expected_row_group_matched_by_statistics
+            ),
             self.expected_row_group_matched_by_statistics,
             "mismatched row_groups_matched_statistics",
         );
         assert_eq!(
-            output.row_groups_fully_matched_statistics(),
+            zero_if_metric_absent(
+                output.row_groups_fully_matched_statistics(),
+                self.expected_row_group_fully_matched_by_statistics
+            ),
             self.expected_row_group_fully_matched_by_statistics,
             "mismatched row_groups_fully_matched_statistics",
         );
         assert_eq!(
-            output.row_groups_pruned_statistics(),
+            zero_if_metric_absent(
+                output.row_groups_pruned_statistics(),
+                self.expected_row_group_pruned_by_statistics
+            ),
             self.expected_row_group_pruned_by_statistics,
             "mismatched row_groups_pruned_statistics",
         );
         assert_eq!(
-            output.files_ranges_pruned_statistics(),
+            zero_if_metric_absent(
+                output.files_ranges_pruned_statistics(),
+                self.expected_files_pruned_by_statistics
+            ),
             self.expected_files_pruned_by_statistics,
             "mismatched files_ranges_pruned_statistics",
         );
         assert_eq!(
-            output.limit_pruned_row_groups(),
+            zero_if_metric_absent(
+                output.limit_pruned_row_groups(),
+                self.expected_limit_pruned_row_groups
+            ),
             self.expected_limit_pruned_row_groups,
             "mismatched limit_pruned_row_groups",
         );
@@ -343,7 +379,10 @@ async fn prune_date64() {
 
     println!("{}", output.description());
     // This should prune out groups  without error
-    assert_eq!(output.predicate_evaluation_errors(), Some(0));
+    assert_eq!(
+        zero_if_metric_absent(output.predicate_evaluation_errors(), Some(0)),
+        Some(0)
+    );
     // 'dates' table has 4 row groups, and only the first one is matched by the predicate
     assert_eq!(output.row_groups_matched_statistics(), Some(1));
     assert_eq!(output.row_groups_pruned_statistics(), Some(3));
@@ -383,9 +422,15 @@ async fn prune_disabled() {
     println!("{}", output.description());
 
     // This should not prune any
-    assert_eq!(output.predicate_evaluation_errors(), Some(0));
+    assert_eq!(
+        zero_if_metric_absent(output.predicate_evaluation_errors(), Some(0)),
+        Some(0)
+    );
     assert_eq!(output.row_groups_matched(), Some(4));
-    assert_eq!(output.row_groups_pruned(), Some(0));
+    assert_eq!(
+        zero_if_metric_absent(output.row_groups_pruned(), Some(0)),
+        Some(0)
+    );
     assert_eq!(
         output.result_rows,
         expected_rows,