diff --git a/datafusion/core/tests/parquet/dynamic_row_group_pruning.rs b/datafusion/core/tests/parquet/dynamic_row_group_pruning.rs index 917faaaa5ce64..5ee42b30674bf 100644 --- a/datafusion/core/tests/parquet/dynamic_row_group_pruning.rs +++ b/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) -> Vec { .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 : +/// 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: +/// (tracked on the +/// DataFusion side in ). /// /// 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) -> Vec { /// 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; diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index a57f4695b55e3..693e9bd2cbf31 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/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 = 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 diff --git a/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt b/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt index 81d9511839725..c6700ebf0b97c 100644 --- a/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt +++ b/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;