Support limit pruning

Author: xudong963

datafusion-testing

@@ -178,6 +178,11 @@ impl TestOutput {
         self.metric_value("page_index_rows_pruned")
     }
 
+    /// The number of row groups pruned by limit pruning
+    fn limit_pruned_row_groups(&self) -> Option<usize> {
+        self.metric_value("limit_pruned_row_groups")
+    }
+
     fn description(&self) -> String {
         format!(
             "Input:\n{}\nQuery:\n{}\nOutput:\n{}\nMetrics:\n{}",
@@ -191,20 +196,40 @@ impl TestOutput {
 /// and the appropriate scenario
 impl ContextWithParquet {
     async fn new(scenario: Scenario, unit: Unit) -> Self {
-        Self::with_config(scenario, unit, SessionConfig::new()).await
+        Self::with_config(scenario, unit, SessionConfig::new(), None, None).await
+    }
+
+    /// Set custom schema and batches for the test
+    pub async fn with_custom_data(
+        scenario: Scenario,
+        unit: Unit,
+        schema: Arc<Schema>,
+        batches: Vec<RecordBatch>,
+    ) -> Self {
+        Self::with_config(
+            scenario,
+            unit,
+            SessionConfig::new(),
+            Some(schema),
+            Some(batches),
+        )
+        .await
     }
 
     async fn with_config(
         scenario: Scenario,
         unit: Unit,
         mut config: SessionConfig,
+        custom_schema: Option<Arc<Schema>>,
+        custom_batches: Option<Vec<RecordBatch>>,
     ) -> Self {
         // Use a single partition for deterministic results no matter how many CPUs the host has
         config = config.with_target_partitions(1);
         let file = match unit {
             Unit::RowGroup(row_per_group) => {
                 config = config.with_parquet_bloom_filter_pruning(true);
-                make_test_file_rg(scenario, row_per_group).await
+                make_test_file_rg(scenario, row_per_group, custom_schema, custom_batches)
+                    .await
             }
             Unit::Page(row_per_page) => {
                 config = config.with_parquet_page_index_pruning(true);
@@ -1030,7 +1055,12 @@ fn create_data_batch(scenario: Scenario) -> Vec<RecordBatch> {
 }
 
 /// Create a test parquet file with various data types
-async fn make_test_file_rg(scenario: Scenario, row_per_group: usize) -> NamedTempFile {
+async fn make_test_file_rg(
+    scenario: Scenario,
+    row_per_group: usize,
+    custom_schema: Option<Arc<Schema>>,
+    custom_batches: Option<Vec<RecordBatch>>,
+) -> NamedTempFile {
     let mut output_file = tempfile::Builder::new()
         .prefix("parquet_pruning")
         .suffix(".parquet")
@@ -1043,8 +1073,14 @@ async fn make_test_file_rg(scenario: Scenario, row_per_group: usize) -> NamedTem
         .set_statistics_enabled(EnabledStatistics::Page)
         .build();
 
-    let batches = create_data_batch(scenario);
-    let schema = batches[0].schema();
+    let (batches, schema) =
+        if let (Some(schema), Some(batches)) = (custom_schema, custom_batches) {
+            (batches, schema)
+        } else {
+            let batches = create_data_batch(scenario);
+            let schema = batches[0].schema();
+            (batches, schema)
+        };
 
     let mut writer = ArrowWriter::try_new(&mut output_file, schema, Some(props)).unwrap();
 

datafusion/core/tests/parquet/mod.rs

@@ -18,8 +18,12 @@
 //! This file contains an end to end test of parquet pruning. It writes
 //! data into a parquet file and then verifies row groups are pruned as
 //! expected.
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Int32Array, RecordBatch};
+use arrow_schema::{DataType, Field, Schema};
 use datafusion::prelude::SessionConfig;
-use datafusion_common::ScalarValue;
+use datafusion_common::{DataFusionError, ScalarValue};
 use itertools::Itertools;
 
 use crate::parquet::Unit::RowGroup;
@@ -34,6 +38,7 @@ struct RowGroupPruningTest {
     expected_files_pruned_by_statistics: Option<usize>,
     expected_row_group_matched_by_bloom_filter: Option<usize>,
     expected_row_group_pruned_by_bloom_filter: Option<usize>,
+    expected_limit_pruned_row_groups: Option<usize>,
     expected_rows: usize,
 }
 impl RowGroupPruningTest {
@@ -48,6 +53,7 @@ impl RowGroupPruningTest {
             expected_files_pruned_by_statistics: None,
             expected_row_group_matched_by_bloom_filter: None,
             expected_row_group_pruned_by_bloom_filter: None,
+            expected_limit_pruned_row_groups: None,
             expected_rows: 0,
         }
     }
@@ -99,6 +105,11 @@ impl RowGroupPruningTest {
         self
     }
 
+    fn with_limit_pruned_row_groups(mut self, pruned_by_limit: Option<usize>) -> Self {
+        self.expected_limit_pruned_row_groups = pruned_by_limit;
+        self
+    }
+
     /// Set the number of expected rows from the output of this test
     fn with_expected_rows(mut self, rows: usize) -> Self {
         self.expected_rows = rows;
@@ -143,6 +154,73 @@ impl RowGroupPruningTest {
             self.expected_row_group_pruned_by_bloom_filter,
             "mismatched row_groups_pruned_bloom_filter",
         );
+        assert_eq!(
+            output.limit_pruned_row_groups(),
+            self.expected_limit_pruned_row_groups,
+            "mismatched limit_pruned_row_groups",
+        );
+        assert_eq!(
+            output.result_rows,
+            self.expected_rows,
+            "Expected {} rows, got {}: {}",
+            output.result_rows,
+            self.expected_rows,
+            output.description(),
+        );
+    }
+
+    // Execute the test with the current configuration
+    async fn test_row_group_prune_with_custom_data(
+        self,
+        schema: Arc<Schema>,
+        batches: Vec<RecordBatch>,
+    ) {
+        let output = ContextWithParquet::with_custom_data(
+            self.scenario,
+            RowGroup(2),
+            schema,
+            batches,
+        )
+        .await
+        .query(&self.query)
+        .await;
+
+        println!("{}", output.description());
+        assert_eq!(
+            output.predicate_evaluation_errors(),
+            self.expected_errors,
+            "mismatched predicate_evaluation error"
+        );
+        assert_eq!(
+            output.row_groups_matched_statistics(),
+            self.expected_row_group_matched_by_statistics,
+            "mismatched row_groups_matched_statistics",
+        );
+        assert_eq!(
+            output.row_groups_pruned_statistics(),
+            self.expected_row_group_pruned_by_statistics,
+            "mismatched row_groups_pruned_statistics",
+        );
+        assert_eq!(
+            output.files_ranges_pruned_statistics(),
+            self.expected_files_pruned_by_statistics,
+            "mismatched files_ranges_pruned_statistics",
+        );
+        assert_eq!(
+            output.row_groups_matched_bloom_filter(),
+            self.expected_row_group_matched_by_bloom_filter,
+            "mismatched row_groups_matched_bloom_filter",
+        );
+        assert_eq!(
+            output.row_groups_pruned_bloom_filter(),
+            self.expected_row_group_pruned_by_bloom_filter,
+            "mismatched row_groups_pruned_bloom_filter",
+        );
+        assert_eq!(
+            output.limit_pruned_row_groups(),
+            self.expected_limit_pruned_row_groups,
+            "mismatched limit_pruned_row_groups",
+        );
         assert_eq!(
             output.result_rows,
             self.expected_rows,
@@ -287,11 +365,16 @@ async fn prune_disabled() {
     let expected_rows = 10;
     let config = SessionConfig::new().with_parquet_pruning(false);
 
-    let output =
-        ContextWithParquet::with_config(Scenario::Timestamps, RowGroup(5), config)
-            .await
-            .query(query)
-            .await;
+    let output = ContextWithParquet::with_config(
+        Scenario::Timestamps,
+        RowGroup(5),
+        config,
+        None,
+        None,
+    )
+    .await
+    .query(query)
+    .await;
     println!("{}", output.description());
 
     // This should not prune any
@@ -1634,3 +1717,56 @@ async fn test_bloom_filter_decimal_dict() {
         .test_row_group_prune()
         .await;
 }
+
+// Helper function to create a batch with a single Int32 column.
+fn make_i32_batch(
+    name: &str,
+    values: Vec<i32>,
+) -> datafusion_common::error::Result<RecordBatch> {
+    let schema = Arc::new(Schema::new(vec![Field::new(name, DataType::Int32, false)]));
+    let array: ArrayRef = Arc::new(Int32Array::from(values));
+    RecordBatch::try_new(schema, vec![array]).map_err(DataFusionError::from)
+}
+
+#[tokio::test]
+async fn test_limit_pruning() -> datafusion_common::error::Result<()> {
+    // Scenario: Simple integer column, multiple row groups
+    // Query: SELECT c1 FROM  t WHERE c1 > 0 LIMIT 2
+    // We expect 2 rows in total.
+
+    // Row Group 0: c1 = [1, 2] -> Fully matched, 2 rows
+    // Row Group 1: c1 = [3, 4] -> Fully matched, 2 rows
+    // Row Group 2: c1 = [5, 6] -> Fully matched, 2 rows
+    // Row Group 3: c1 = [-1, 0] -> Pruned by statistics, 0 rows
+
+    // If limit = 2, and RG0 is fully matched and has 2 rows, we should
+    // only scan RG0 and prune other row groups (RG1, RG2, RG3)
+    // RG3 is pruned by statistics. RG1 and RG2 are pruned by limit.
+    // So 3 row groups are effectively pruned due to limit pruning.
+
+    let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)]));
+    let query = "explain verbose SELECT c1 FROM t WHERE c1 > 0 LIMIT 2";
+
+    let batches = vec![
+        make_i32_batch("c1", vec![1, 2])?, // RG0: Fully matched, 2 rows
+        make_i32_batch("c1", vec![3, 4])?, // RG1: Fully matched, 2 rows
+        make_i32_batch("c1", vec![5, 6])?, // RG2: Fully matched, 2 rows
+        make_i32_batch("c1", vec![-1, 0])?, // RG3: Pruned by statistics, 0 rows
+    ];
+
+    RowGroupPruningTest::new()
+        .with_scenario(Scenario::Int) // Assuming Scenario::Int can handle this data
+        .with_query(query)
+        .with_expected_errors(Some(0))
+        .with_expected_rows(2)
+        .with_pruned_files(Some(0))
+        .with_matched_by_bloom_filter(Some(0))
+        .with_pruned_by_bloom_filter(Some(0))
+        .with_matched_by_stats(Some(3)) // RG0, RG1, RG2 are matched by stats (c1 > 0)
+        .with_pruned_by_stats(Some(1)) // RG3 is pruned by stats (c1 = [-1, 0] does not satisfy c1 > 0)
+        // .with_limit_pruned_row_groups(Some(2)) // RG1, RG2 are pruned by limit. (RG3 is already pruned by stats)
+        .test_row_group_prune_with_custom_data(schema, batches)
+        .await;
+
+    Ok(())
+}

datafusion/core/tests/parquet/row_group_pruning.rs

@@ -48,6 +48,8 @@ pub struct ParquetFileMetrics {
     pub row_groups_matched_bloom_filter: Count,
     /// Number of row groups pruned by bloom filters
     pub row_groups_pruned_bloom_filter: Count,
+    /// Number of row groups pruned due to limit pruning.
+    pub limit_pruned_row_groups: Count,
     /// Number of row groups whose statistics were checked and matched (not pruned)
     pub row_groups_matched_statistics: Count,
     /// Number of row groups pruned by statistics
@@ -93,6 +95,10 @@ impl ParquetFileMetrics {
             .with_new_label("filename", filename.to_string())
             .counter("row_groups_pruned_bloom_filter", partition);
 
+        let limit_pruned_row_groups = MetricBuilder::new(metrics)
+            .with_new_label("filename", filename.to_string())
+            .counter("limit_pruned_row_groups", partition);
+
         let row_groups_matched_statistics = MetricBuilder::new(metrics)
             .with_new_label("filename", filename.to_string())
             .counter("row_groups_matched_statistics", partition);
@@ -147,6 +153,7 @@ impl ParquetFileMetrics {
             row_groups_pruned_bloom_filter,
             row_groups_matched_statistics,
             row_groups_pruned_statistics,
+            limit_pruned_row_groups,
             bytes_scanned,
             pushdown_rows_pruned,
             pushdown_rows_matched,

datafusion/datasource-parquet/src/metrics.rs

@@ -373,8 +373,12 @@ impl FileOpener for ParquetOpener {
                 }
             }
 
-            let mut access_plan = row_groups.build();
+            // Prune by limit
+            if let Some(limit) = limit {
+                row_groups.prune_by_limit(limit, rg_metadata, &file_metrics);
+            }
 
+            let mut access_plan = row_groups.build();
             // page index pruning: if all data on individual pages can
             // be ruled using page metadata, rows from other columns
             // with that range can be skipped as well