fix(parquet): don't runtime-prune row groups while a page-index RowSelection is live (#24355)

datafusion/core/tests/parquet/dynamic_row_group_pruning.rs

@@ -35,6 +35,8 @@ use std::sync::Arc;
 use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray};
 use arrow_schema::{DataType, Field, Schema};
 
+use datafusion::prelude::SessionConfig;
+
 use crate::parquet::Unit::RowGroup;
 use crate::parquet::{ContextWithParquet, Scenario};
 
@@ -297,9 +299,18 @@ fn build_five_thousand_row_rgs(schema: &Arc<Schema>) -> Vec<RecordBatch> {
         .collect()
 }
 
-/// Co-existence test for **page-index `RowSelection`** + dynamic RG
-/// pruning. Tests that the `into_builder` rebuild preserves the
-/// `RowSelection` derived from page-index pruning across RG drops.
+/// Regression test for <https://github.com/apache/datafusion/issues/24355>:
+/// when a page-index `RowSelection` is live, the runtime dynamic row-group
+/// pruner is intentionally **not built**, so its `into_builder` rebuild can
+/// never drop a row group without slicing the carried selection (which would
+/// silently return wrong rows). Correctness is bought at the cost of the
+/// dynamic-pruning optimization for this scan.
+///
+/// The behavior asserted below (pruner disabled →
+/// `row_groups_pruned_dynamic_filter == 0`) is expected to change once the
+/// proper upstream fix lands, which keeps both mechanisms:
+/// <https://github.com/apache/arrow-rs/issues/10624> (tracked on the
+/// DataFusion side in <https://github.com/apache/datafusion/issues/24358>).
 ///
 /// Layout: 5 RGs × 1000 rows, with `data_page_row_count_limit=100` so
 /// each RG has 10 pages of 100 rows.
@@ -309,17 +320,14 @@ fn build_five_thousand_row_rgs(schema: &Arc<Schema>) -> Vec<RecordBatch> {
 ///   first 5 pages (values 0..500) are pruned, the last 5 (500..1000)
 ///   are scanned. RGs 1..4 keep all their pages (every page has
 ///   `max >= 500`). The decoder receives a `RowSelection` that masks
-///   out those first 5 pages of RG 0.
-/// - `ORDER BY v DESC LIMIT 5` fills the TopK heap from RG 4
-///   (`max=4999`); the tightened threshold (≥ 4995) then proves RGs
-///   0..3 unreachable and the runtime pruner drops them in one
-///   `into_builder` rebuild.
-///
-/// If `into_builder` did **not** preserve the row selection (or
-/// truncated / shifted it incorrectly), either the result rows would
-/// drift or the count of pruned pages would drop to zero.
+///   out those first 5 pages of RG 0 — its presence is what suppresses
+///   the runtime pruner.
+/// - `ORDER BY v DESC LIMIT 5` would let the tightened TopK threshold
+///   (≥ 4995) prune RGs 0..3, but because a row selection is present the
+///   runtime pruner is never created, so `row_groups_pruned_dynamic_filter`
+///   stays 0. Results are still correct and page-index pruning still runs.
 #[tokio::test]
-async fn dynamic_rg_pruning_coexists_with_page_index_row_selection() {
+async fn dynamic_rg_pruning_disabled_when_page_index_row_selection_present() {
     let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)]));
     let batches = build_five_thousand_row_rgs(&schema);
 
@@ -348,12 +356,9 @@ async fn dynamic_rg_pruning_coexists_with_page_index_row_selection() {
         );
     }
 
-    // Page-index pruning must have engaged: RG 0's first 5 pages are
-    // entirely < 500. If `into_builder` dropped the row-selection state,
-    // this metric would still report the original count (it is captured
-    // at file open). Combined with the dynamic-pruner assertion below it
-    // proves both mechanisms were active and that the rebuild left the
-    // selection coherent — otherwise the result rows above would drift.
+    // Page-index pruning still engages: RG 0's first 5 pages are entirely
+    // < 500. #24355 only suppresses the *runtime* row-group pruner, not
+    // page-index pruning, so this must remain non-zero.
     let pages_pruned = output.metric_value("page_index_pages_pruned").unwrap_or(0);
     assert!(
         pages_pruned >= 5,
@@ -362,13 +367,18 @@ async fn dynamic_rg_pruning_coexists_with_page_index_row_selection() {
         output.description(),
     );
 
+    // The runtime dynamic pruner must be disabled while a page-index row
+    // selection is live (#24355): with no pruner there is no rebuild that
+    // could misapply the carried selection. Before the fix the pruner ran
+    // and this metric was >= 1.
     let pruned = output
         .row_groups_pruned_dynamic_filter()
         .expect("`row_groups_pruned_dynamic_filter` metric must be registered");
-    assert!(
-        pruned >= 1,
-        "with TopK + tight threshold the runtime pruner must skip at least \
-         one row group; pruned={pruned}\n{}",
+    assert_eq!(
+        pruned,
+        0,
+        "runtime row-group pruning must be skipped when a page-index row \
+         selection is present; pruned={pruned}\n{}",
         output.description(),
     );
 }
@@ -647,12 +657,20 @@ async fn topk_pushdown_does_not_reread_delivered_row_group() {
 
     // `RowGroup(2048)` writes one row group per 2048-row batch (4 RGs) and
     // enables `pushdown_filters`, required for the dynamic filter to reach the
-    // parquet scan.
-    let mut ctx = ContextWithParquet::with_custom_data(
+    // parquet scan. Page-index reading is disabled: this test exercises the
+    // #24352 empty-row-group / rg_plan-sync path, which is row-filter-driven and
+    // does not need the page index. With the page index on, `search_phrase <> ''`
+    // produces an intra-row-group `RowSelection`, and #24355 disables the runtime
+    // pruner whenever a row selection is present — which would stop this test
+    // from exercising the dynamic pruner at all.
+    let mut config = SessionConfig::new();
+    config.options_mut().execution.parquet.enable_page_index = false;
+    let mut ctx = ContextWithParquet::with_config(
         Scenario::Int,
         RowGroup(2048),
-        Arc::clone(&schema),
-        batches,
+        config,
+        Some(Arc::clone(&schema)),
+        Some(batches),
     )
     .await;
 

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

@@ -1435,7 +1435,7 @@ impl RowGroupsPrunedParquetOpen {
             prepared.virtual_state.as_deref(),
         )?;
 
-        let (decoder, rg_plan) = {
+        let (decoder, rg_plan, has_row_selection) = {
             let pushdown_predicate = prepared
                 .pushdown_filters
                 .then_some(prepared.predicate.as_ref())
@@ -1464,6 +1464,18 @@ impl RowGroupsPrunedParquetOpen {
             };
 
             let prepared_access_plan = prepare_access_plan(access_plan)?;
+            // #24355: a row selection (from page-index pruning, or an externally
+            // supplied `ParquetRowSelection`) is carried by the decoder as one
+            // flat selection over the concatenation of the remaining row groups.
+            // The runtime pruner's `into_builder().with_row_groups(...)` rebuild
+            // drops row groups without slicing that selection to match, so record
+            // whether a selection is present and disable runtime pruning below
+            // when it is (mirroring `reorder_by_statistics`, which also bails when
+            // a row selection is present). The proper fix that keeps pruning
+            // under a live selection is tracked in
+            // https://github.com/apache/arrow-rs/issues/10624 /
+            // https://github.com/apache/datafusion/issues/24358.
+            let has_row_selection = prepared_access_plan.row_selection.is_some();
             let rg_plan: VecDeque<RgPlanEntry> = prepared_access_plan
                 .row_group_indexes
                 .iter()
@@ -1482,7 +1494,7 @@ impl RowGroupsPrunedParquetOpen {
                 }
             }
 
-            (builder.build()?, rg_plan)
+            (builder.build()?, rg_plan, has_row_selection)
         };
 
         let predicate_cache_inner_records =
@@ -1504,24 +1516,30 @@ impl RowGroupsPrunedParquetOpen {
         // via the `DynamicFilterTracker` watch channel (#22460), so detecting
         // a threshold change is a single atomic load — not a tree walk per
         // RG check.
-        let row_group_pruner = match (&prepared.predicate, rg_plan.len() > 1) {
-            (Some(predicate), true)
-                if matches!(
-                    DynamicFilterTracking::classify(predicate),
-                    DynamicFilterTracking::Watching(_)
-                ) =>
-            {
-                Some(RowGroupPruner::new(
-                    Arc::clone(predicate),
-                    Arc::clone(&prepared.physical_file_schema),
-                    Arc::clone(reader_metadata.metadata()),
-                    prepared.predicate_creation_errors.clone(),
-                    prepared.file_metrics.predicate_evaluation_errors.clone(),
-                    prepared.max_in_list_size,
-                ))
-            }
-            _ => None,
-        };
+        // Also disabled when a row selection is live (#24355) — page-index
+        // pruning is the common source: the pruner rebuilds the decoder via
+        // `with_row_groups(...)`, which drops row groups without slicing the
+        // carried selection to match, so pruning under a live selection returns
+        // wrong results. Decline to prune in that case.
+        let row_group_pruner =
+            match (&prepared.predicate, rg_plan.len() > 1, has_row_selection) {
+                (Some(predicate), true, false)
+                    if matches!(
+                        DynamicFilterTracking::classify(predicate),
+                        DynamicFilterTracking::Watching(_)
+                    ) =>
+                {
+                    Some(RowGroupPruner::new(
+                        Arc::clone(predicate),
+                        Arc::clone(&prepared.physical_file_schema),
+                        Arc::clone(reader_metadata.metadata()),
+                        prepared.predicate_creation_errors.clone(),
+                        prepared.file_metrics.predicate_evaluation_errors.clone(),
+                        prepared.max_in_list_size,
+                    ))
+                }
+                _ => None,
+            };
         let row_groups_pruned_dynamic = prepared
             .file_metrics
             .row_groups_pruned_dynamic_filter

datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt

@@ -194,3 +194,96 @@ RESET datafusion.optimizer.enable_dynamic_filter_pushdown;
 
 statement ok
 RESET datafusion.optimizer.enable_topk_dynamic_filter_pushdown;
+
+# Regression test for a scan where two pruning mechanisms are live at once:
+# page-index pruning leaves an intra-row-group `RowSelection`, and a TopK
+# dynamic filter prunes row groups at runtime. The property under test is that
+# a `WHERE a >= 50 ORDER BY b ASC LIMIT 5` query returns the correct top-5 by
+# `b` while the dynamic predicate prunes a row group during the application of
+# multiple predicates. Layout (RG size 100):
+#   RG 0: b=1000..1099, a=100..199  (a>=50 keeps all)
+#   RG 1: b=2000..2099, a=0..99     (a>=50 keeps rows 50..99 — page-index prunes
+#                                     the first 5 pages, leaving `skip 50, select 50`)
+#   RG 2: b=3000..3099, a=100..199  (keeps all)
+#   RG 3: b=0..99,      a=100..199  (keeps all)
+# The correct top-5 by `b` (0..4) lives entirely in RG 3.
+# `data_page_row_count_limit`/`write_batch_size` force multiple pages per RG so
+# page-index pruning can produce the intra-RG selection.
+# Tracking issue for the behavior change (keeping both mechanisms):
+# https://github.com/apache/arrow-rs/issues/10624 /
+# https://github.com/apache/datafusion/issues/24358
+statement ok
+set datafusion.execution.target_partitions = 1;
+
+statement ok
+set datafusion.execution.parquet.pushdown_filters = true;
+
+statement ok
+CREATE TABLE rgsel_src AS
+SELECT
+  CAST(CASE WHEN i / 100 = 1 THEN i % 100 ELSE 100 + (i % 100) END AS BIGINT) AS a,
+  CAST(CASE
+    WHEN i < 100 THEN 1000 + i
+    WHEN i < 200 THEN 2000 + (i - 100)
+    WHEN i < 300 THEN 3000 + (i - 200)
+    ELSE (i - 300)
+  END AS BIGINT) AS b
+FROM generate_series(0, 399) AS t(i);
+
+statement ok
+COPY (SELECT * FROM rgsel_src)
+TO 'test_files/scratch/dynamic_row_group_pruning/rgsel.parquet'
+STORED AS PARQUET
+OPTIONS (
+  'format.max_row_group_size' '100',
+  'format.data_page_row_count_limit' '10',
+  'format.write_batch_size' '10'
+);
+
+statement ok
+drop table rgsel_src;
+
+statement ok
+CREATE EXTERNAL TABLE rgsel (a BIGINT NOT NULL, b BIGINT NOT NULL)
+STORED AS PARQUET
+LOCATION 'test_files/scratch/dynamic_row_group_pruning/rgsel.parquet';
+
+# The correct top-5 by `b` among rows with `a >= 50` is b = 0..4 (they live in
+# RG 3, all of whose rows satisfy `a >= 50`).
+query I
+SELECT b FROM rgsel WHERE a >= 50 ORDER BY b ASC LIMIT 5;
+----
+0
+1
+2
+3
+4
+
+# The same query without filter pushdown never engages the runtime pruner, so
+# its answer is the ground truth the pushdown path above must match.
+statement ok
+set datafusion.execution.parquet.pushdown_filters = false;
+
+query I
+SELECT b FROM rgsel WHERE a >= 50 ORDER BY b ASC LIMIT 5;
+----
+0
+1
+2
+3
+4
+
+statement ok
+set datafusion.execution.parquet.pushdown_filters = true;
+
+statement ok
+drop table rgsel;
+
+# The SLT runner sets `target_partitions` to 4 instead of using the default, so
+# restore it explicitly rather than RESET (which would revert to the system
+# default = num_cpus and leak modified config out of this file).
+statement ok
+set datafusion.execution.target_partitions = 4;
+
+statement ok
+RESET datafusion.execution.parquet.pushdown_filters;