Skip to content

Commit ec110ce

Browse files
authored
fix(parquet): don't runtime-prune row groups while a page-index RowSelection is live (#24355) (#24359)
## Which issue does this PR close? - Closes #24355. ## Rationale for this change With `pushdown_filters = true` + dynamic filter pushdown (on by default), a query `SELECT b FROM t WHERE <predicate on a> ORDER BY b LIMIT k` can silently return **wrong results** — rows satisfying the predicate are dropped and replaced by later ones, no error. Root cause (thanks to @adriangb's report + fixture in #24355): the push decoder carries one flat `RowSelection` over the concatenation of the *remaining* row groups. At a row-group boundary the runtime pruner drops row groups the dynamic predicate proves unwinnable and rebuilds the decoder: ```rust decoder.into_builder()?.with_row_groups(new_indices).build() ``` `with_row_groups(new_indices)` removes row groups **without slicing the carried `RowSelection` to match**, so the selectors intended for a dropped RG are applied to the next surviving one. In the fixture, page-index pruning leaves RG 1 with `skip 50, select 50`; after the TopK threshold prunes RG 1 and RG 2, the survivor RG 3 is decoded under RG 1's selection and its first 50 rows (`b = 0..49`, the correct answer) are wrongly skipped. This is a second, independent instance of the drift family in #24352/#24354; it is **not** fixed by #24354. ## What changes are included in this PR? - `opener/mod.rs`: **decline to build the runtime `RowGroupPruner` when a page-index `RowSelection` is present.** With no pruner there is no boundary rebuild, so the carried selection is never applied to the wrong row groups. This mirrors `PreparedAccessPlan::reorder_by_statistics`, which already bails when a row selection is present (`"Skipping RG reorder: row_selection present"`) because remapping the selection is too complex. This is the minimal, DataFusion-side stop-the-bleeding fix. The proper fix is upstream in arrow-rs: apache/arrow-rs#10624 proposes letting the push decoder carry **row-group-local** `RowSelection`s (`with_row_group_selections`) that are preserved across rebuilds, so dropping a row group keeps every survivor's selection aligned by construction — no global-selection slicing to get wrong, and no parallel `rg_plan` to drift (the #24352 path). DataFusion tracks that migration in #24358; this guard is removed once it lands. ## Are these changes tested? - Adds an slt regression test in `dynamic_row_group_pruning.slt` (the reporter's fixture via `generate_series` + `COPY`). It **fails on `main`** (returns `50..54` instead of `0..4`) and passes with this change. - Updates the existing rust integration test that previously asserted the runtime pruner **coexists** with a page-index selection (`dynamic_rg_pruning_coexists_with_page_index_row_selection`, `row_groups_pruned_dynamic_filter >= 1`). Since this PR intentionally disables the pruner in that case, it is renamed to `dynamic_rg_pruning_disabled_when_page_index_row_selection_present` and now asserts `row_groups_pruned_dynamic_filter == 0` while results stay correct and page-index pruning still runs. (That old test passed only because its scenario happened not to expose the bug — the misapplied selection fell outside the top-k.) - The other dynamic-prune tests are unaffected — they have no row selection, so the pruner is created as before. ## Are there any user-facing changes? Fixes silently-wrong results. Runtime row-group pruning is skipped for scans that also have a page-index row selection (correctness over a pruning optimization); this is undone once #24358 lands. cc @alamb @adriangb @hhhizzz
1 parent 1b67f2e commit ec110ce

3 files changed

Lines changed: 176 additions & 47 deletions

File tree

datafusion/core/tests/parquet/dynamic_row_group_pruning.rs

Lines changed: 45 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,8 @@ use std::sync::Arc;
3535
use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray};
3636
use arrow_schema::{DataType, Field, Schema};
3737

38+
use datafusion::prelude::SessionConfig;
39+
3840
use crate::parquet::Unit::RowGroup;
3941
use crate::parquet::{ContextWithParquet, Scenario};
4042

@@ -297,9 +299,18 @@ fn build_five_thousand_row_rgs(schema: &Arc<Schema>) -> Vec<RecordBatch> {
297299
.collect()
298300
}
299301

300-
/// Co-existence test for **page-index `RowSelection`** + dynamic RG
301-
/// pruning. Tests that the `into_builder` rebuild preserves the
302-
/// `RowSelection` derived from page-index pruning across RG drops.
302+
/// Regression test for <https://github.com/apache/datafusion/issues/24355>:
303+
/// when a page-index `RowSelection` is live, the runtime dynamic row-group
304+
/// pruner is intentionally **not built**, so its `into_builder` rebuild can
305+
/// never drop a row group without slicing the carried selection (which would
306+
/// silently return wrong rows). Correctness is bought at the cost of the
307+
/// dynamic-pruning optimization for this scan.
308+
///
309+
/// The behavior asserted below (pruner disabled →
310+
/// `row_groups_pruned_dynamic_filter == 0`) is expected to change once the
311+
/// proper upstream fix lands, which keeps both mechanisms:
312+
/// <https://github.com/apache/arrow-rs/issues/10624> (tracked on the
313+
/// DataFusion side in <https://github.com/apache/datafusion/issues/24358>).
303314
///
304315
/// Layout: 5 RGs × 1000 rows, with `data_page_row_count_limit=100` so
305316
/// each RG has 10 pages of 100 rows.
@@ -309,17 +320,14 @@ fn build_five_thousand_row_rgs(schema: &Arc<Schema>) -> Vec<RecordBatch> {
309320
/// first 5 pages (values 0..500) are pruned, the last 5 (500..1000)
310321
/// are scanned. RGs 1..4 keep all their pages (every page has
311322
/// `max >= 500`). The decoder receives a `RowSelection` that masks
312-
/// out those first 5 pages of RG 0.
313-
/// - `ORDER BY v DESC LIMIT 5` fills the TopK heap from RG 4
314-
/// (`max=4999`); the tightened threshold (≥ 4995) then proves RGs
315-
/// 0..3 unreachable and the runtime pruner drops them in one
316-
/// `into_builder` rebuild.
317-
///
318-
/// If `into_builder` did **not** preserve the row selection (or
319-
/// truncated / shifted it incorrectly), either the result rows would
320-
/// drift or the count of pruned pages would drop to zero.
323+
/// out those first 5 pages of RG 0 — its presence is what suppresses
324+
/// the runtime pruner.
325+
/// - `ORDER BY v DESC LIMIT 5` would let the tightened TopK threshold
326+
/// (≥ 4995) prune RGs 0..3, but because a row selection is present the
327+
/// runtime pruner is never created, so `row_groups_pruned_dynamic_filter`
328+
/// stays 0. Results are still correct and page-index pruning still runs.
321329
#[tokio::test]
322-
async fn dynamic_rg_pruning_coexists_with_page_index_row_selection() {
330+
async fn dynamic_rg_pruning_disabled_when_page_index_row_selection_present() {
323331
let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)]));
324332
let batches = build_five_thousand_row_rgs(&schema);
325333

@@ -348,12 +356,9 @@ async fn dynamic_rg_pruning_coexists_with_page_index_row_selection() {
348356
);
349357
}
350358

351-
// Page-index pruning must have engaged: RG 0's first 5 pages are
352-
// entirely < 500. If `into_builder` dropped the row-selection state,
353-
// this metric would still report the original count (it is captured
354-
// at file open). Combined with the dynamic-pruner assertion below it
355-
// proves both mechanisms were active and that the rebuild left the
356-
// selection coherent — otherwise the result rows above would drift.
359+
// Page-index pruning still engages: RG 0's first 5 pages are entirely
360+
// < 500. #24355 only suppresses the *runtime* row-group pruner, not
361+
// page-index pruning, so this must remain non-zero.
357362
let pages_pruned = output.metric_value("page_index_pages_pruned").unwrap_or(0);
358363
assert!(
359364
pages_pruned >= 5,
@@ -362,13 +367,18 @@ async fn dynamic_rg_pruning_coexists_with_page_index_row_selection() {
362367
output.description(),
363368
);
364369

370+
// The runtime dynamic pruner must be disabled while a page-index row
371+
// selection is live (#24355): with no pruner there is no rebuild that
372+
// could misapply the carried selection. Before the fix the pruner ran
373+
// and this metric was >= 1.
365374
let pruned = output
366375
.row_groups_pruned_dynamic_filter()
367376
.expect("`row_groups_pruned_dynamic_filter` metric must be registered");
368-
assert!(
369-
pruned >= 1,
370-
"with TopK + tight threshold the runtime pruner must skip at least \
371-
one row group; pruned={pruned}\n{}",
377+
assert_eq!(
378+
pruned,
379+
0,
380+
"runtime row-group pruning must be skipped when a page-index row \
381+
selection is present; pruned={pruned}\n{}",
372382
output.description(),
373383
);
374384
}
@@ -647,12 +657,20 @@ async fn topk_pushdown_does_not_reread_delivered_row_group() {
647657

648658
// `RowGroup(2048)` writes one row group per 2048-row batch (4 RGs) and
649659
// enables `pushdown_filters`, required for the dynamic filter to reach the
650-
// parquet scan.
651-
let mut ctx = ContextWithParquet::with_custom_data(
660+
// parquet scan. Page-index reading is disabled: this test exercises the
661+
// #24352 empty-row-group / rg_plan-sync path, which is row-filter-driven and
662+
// does not need the page index. With the page index on, `search_phrase <> ''`
663+
// produces an intra-row-group `RowSelection`, and #24355 disables the runtime
664+
// pruner whenever a row selection is present — which would stop this test
665+
// from exercising the dynamic pruner at all.
666+
let mut config = SessionConfig::new();
667+
config.options_mut().execution.parquet.enable_page_index = false;
668+
let mut ctx = ContextWithParquet::with_config(
652669
Scenario::Int,
653670
RowGroup(2048),
654-
Arc::clone(&schema),
655-
batches,
671+
config,
672+
Some(Arc::clone(&schema)),
673+
Some(batches),
656674
)
657675
.await;
658676

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

Lines changed: 38 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1435,7 +1435,7 @@ impl RowGroupsPrunedParquetOpen {
14351435
prepared.virtual_state.as_deref(),
14361436
)?;
14371437

1438-
let (decoder, rg_plan) = {
1438+
let (decoder, rg_plan, has_row_selection) = {
14391439
let pushdown_predicate = prepared
14401440
.pushdown_filters
14411441
.then_some(prepared.predicate.as_ref())
@@ -1464,6 +1464,18 @@ impl RowGroupsPrunedParquetOpen {
14641464
};
14651465

14661466
let prepared_access_plan = prepare_access_plan(access_plan)?;
1467+
// #24355: a row selection (from page-index pruning, or an externally
1468+
// supplied `ParquetRowSelection`) is carried by the decoder as one
1469+
// flat selection over the concatenation of the remaining row groups.
1470+
// The runtime pruner's `into_builder().with_row_groups(...)` rebuild
1471+
// drops row groups without slicing that selection to match, so record
1472+
// whether a selection is present and disable runtime pruning below
1473+
// when it is (mirroring `reorder_by_statistics`, which also bails when
1474+
// a row selection is present). The proper fix that keeps pruning
1475+
// under a live selection is tracked in
1476+
// https://github.com/apache/arrow-rs/issues/10624 /
1477+
// https://github.com/apache/datafusion/issues/24358.
1478+
let has_row_selection = prepared_access_plan.row_selection.is_some();
14671479
let rg_plan: VecDeque<RgPlanEntry> = prepared_access_plan
14681480
.row_group_indexes
14691481
.iter()
@@ -1482,7 +1494,7 @@ impl RowGroupsPrunedParquetOpen {
14821494
}
14831495
}
14841496

1485-
(builder.build()?, rg_plan)
1497+
(builder.build()?, rg_plan, has_row_selection)
14861498
};
14871499

14881500
let predicate_cache_inner_records =
@@ -1504,24 +1516,30 @@ impl RowGroupsPrunedParquetOpen {
15041516
// via the `DynamicFilterTracker` watch channel (#22460), so detecting
15051517
// a threshold change is a single atomic load — not a tree walk per
15061518
// RG check.
1507-
let row_group_pruner = match (&prepared.predicate, rg_plan.len() > 1) {
1508-
(Some(predicate), true)
1509-
if matches!(
1510-
DynamicFilterTracking::classify(predicate),
1511-
DynamicFilterTracking::Watching(_)
1512-
) =>
1513-
{
1514-
Some(RowGroupPruner::new(
1515-
Arc::clone(predicate),
1516-
Arc::clone(&prepared.physical_file_schema),
1517-
Arc::clone(reader_metadata.metadata()),
1518-
prepared.predicate_creation_errors.clone(),
1519-
prepared.file_metrics.predicate_evaluation_errors.clone(),
1520-
prepared.max_in_list_size,
1521-
))
1522-
}
1523-
_ => None,
1524-
};
1519+
// Also disabled when a row selection is live (#24355) — page-index
1520+
// pruning is the common source: the pruner rebuilds the decoder via
1521+
// `with_row_groups(...)`, which drops row groups without slicing the
1522+
// carried selection to match, so pruning under a live selection returns
1523+
// wrong results. Decline to prune in that case.
1524+
let row_group_pruner =
1525+
match (&prepared.predicate, rg_plan.len() > 1, has_row_selection) {
1526+
(Some(predicate), true, false)
1527+
if matches!(
1528+
DynamicFilterTracking::classify(predicate),
1529+
DynamicFilterTracking::Watching(_)
1530+
) =>
1531+
{
1532+
Some(RowGroupPruner::new(
1533+
Arc::clone(predicate),
1534+
Arc::clone(&prepared.physical_file_schema),
1535+
Arc::clone(reader_metadata.metadata()),
1536+
prepared.predicate_creation_errors.clone(),
1537+
prepared.file_metrics.predicate_evaluation_errors.clone(),
1538+
prepared.max_in_list_size,
1539+
))
1540+
}
1541+
_ => None,
1542+
};
15251543
let row_groups_pruned_dynamic = prepared
15261544
.file_metrics
15271545
.row_groups_pruned_dynamic_filter

datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -194,3 +194,96 @@ RESET datafusion.optimizer.enable_dynamic_filter_pushdown;
194194

195195
statement ok
196196
RESET datafusion.optimizer.enable_topk_dynamic_filter_pushdown;
197+
198+
# Regression test for a scan where two pruning mechanisms are live at once:
199+
# page-index pruning leaves an intra-row-group `RowSelection`, and a TopK
200+
# dynamic filter prunes row groups at runtime. The property under test is that
201+
# a `WHERE a >= 50 ORDER BY b ASC LIMIT 5` query returns the correct top-5 by
202+
# `b` while the dynamic predicate prunes a row group during the application of
203+
# multiple predicates. Layout (RG size 100):
204+
# RG 0: b=1000..1099, a=100..199 (a>=50 keeps all)
205+
# RG 1: b=2000..2099, a=0..99 (a>=50 keeps rows 50..99 — page-index prunes
206+
# the first 5 pages, leaving `skip 50, select 50`)
207+
# RG 2: b=3000..3099, a=100..199 (keeps all)
208+
# RG 3: b=0..99, a=100..199 (keeps all)
209+
# The correct top-5 by `b` (0..4) lives entirely in RG 3.
210+
# `data_page_row_count_limit`/`write_batch_size` force multiple pages per RG so
211+
# page-index pruning can produce the intra-RG selection.
212+
# Tracking issue for the behavior change (keeping both mechanisms):
213+
# https://github.com/apache/arrow-rs/issues/10624 /
214+
# https://github.com/apache/datafusion/issues/24358
215+
statement ok
216+
set datafusion.execution.target_partitions = 1;
217+
218+
statement ok
219+
set datafusion.execution.parquet.pushdown_filters = true;
220+
221+
statement ok
222+
CREATE TABLE rgsel_src AS
223+
SELECT
224+
CAST(CASE WHEN i / 100 = 1 THEN i % 100 ELSE 100 + (i % 100) END AS BIGINT) AS a,
225+
CAST(CASE
226+
WHEN i < 100 THEN 1000 + i
227+
WHEN i < 200 THEN 2000 + (i - 100)
228+
WHEN i < 300 THEN 3000 + (i - 200)
229+
ELSE (i - 300)
230+
END AS BIGINT) AS b
231+
FROM generate_series(0, 399) AS t(i);
232+
233+
statement ok
234+
COPY (SELECT * FROM rgsel_src)
235+
TO 'test_files/scratch/dynamic_row_group_pruning/rgsel.parquet'
236+
STORED AS PARQUET
237+
OPTIONS (
238+
'format.max_row_group_size' '100',
239+
'format.data_page_row_count_limit' '10',
240+
'format.write_batch_size' '10'
241+
);
242+
243+
statement ok
244+
drop table rgsel_src;
245+
246+
statement ok
247+
CREATE EXTERNAL TABLE rgsel (a BIGINT NOT NULL, b BIGINT NOT NULL)
248+
STORED AS PARQUET
249+
LOCATION 'test_files/scratch/dynamic_row_group_pruning/rgsel.parquet';
250+
251+
# The correct top-5 by `b` among rows with `a >= 50` is b = 0..4 (they live in
252+
# RG 3, all of whose rows satisfy `a >= 50`).
253+
query I
254+
SELECT b FROM rgsel WHERE a >= 50 ORDER BY b ASC LIMIT 5;
255+
----
256+
0
257+
1
258+
2
259+
3
260+
4
261+
262+
# The same query without filter pushdown never engages the runtime pruner, so
263+
# its answer is the ground truth the pushdown path above must match.
264+
statement ok
265+
set datafusion.execution.parquet.pushdown_filters = false;
266+
267+
query I
268+
SELECT b FROM rgsel WHERE a >= 50 ORDER BY b ASC LIMIT 5;
269+
----
270+
0
271+
1
272+
2
273+
3
274+
4
275+
276+
statement ok
277+
set datafusion.execution.parquet.pushdown_filters = true;
278+
279+
statement ok
280+
drop table rgsel;
281+
282+
# The SLT runner sets `target_partitions` to 4 instead of using the default, so
283+
# restore it explicitly rather than RESET (which would revert to the system
284+
# default = num_cpus and leak modified config out of this file).
285+
statement ok
286+
set datafusion.execution.target_partitions = 4;
287+
288+
statement ok
289+
RESET datafusion.execution.parquet.pushdown_filters;

0 commit comments

Comments
 (0)