Skip to content

Commit 574fe67

Browse files
authored
fix(parquet): sync rg_plan to decoder frontier — fix wrong TopK results from re-reading already-delivered row groups (#24352) (#24354)
## Which issue does this PR close? - Closes #24352. ## Rationale for this change With `datafusion.execution.parquet.pushdown_filters = true` and TopK dynamic filter pushdown (both on by default), a query of the shape `SELECT b FROM t WHERE <predicate on a> ORDER BY b LIMIT k` can silently return **wrong results** — one source row emitted several times and the true tail of the top-k missing — with no error or warning. Root cause (thanks to @hhhizzz's very detailed report + fixture in #24352): a row group whose post-predicate selection is empty is silently finished by arrow-rs **without handing back a reader**. `PushDecoderStreamState` pops its `rg_plan` **only** when a reader is returned, so after a silently-finished RG the plan trails the decoder by one. When the runtime row-group pruner then rebuilds the decoder (`into_builder().with_row_groups(...)`) from the stale `rg_plan`, it re-includes an already-delivered row group, whose rows are emitted a second time and displace the genuine top-k in the heap. ## What changes are included in this PR? - `push_decoder.rs`: before each boundary prune/rebuild, `rg_plan` is synced to the row group the decoder will actually emit next via `peek_next_row_group()` (`sync_rg_plan_to_decoder_frontier` / `advance_rg_plan_to`), dropping entries for silently-finished row groups so a rebuild can never re-include a delivered group. A rebuild frontier naming an RG not in the plan is now an internal error instead of a silent plan drain. ## Are these changes tested? - Adds @hhhizzz's fixture as an slt regression test in `dynamic_row_group_pruning.slt` (filter column `search_phrase` differs from the sort column `event_time`, one row group has an empty post-predicate selection invisible to statistics). It now returns the correct `p0 p4096 p4097 … p4104` (was the buggy `p0 p4096 p4096 …`). - clippy clean; `datasource-parquet` unit tests and the sqllogictest suite pass locally. ## Are there any user-facing changes? Fixes silently-wrong query results; no API change. ## Note This is the standalone bug fix extracted from #23696 (per review discussion in #24352): the same `rg_plan` ↔ decoder-frontier sync, on its own so it merges fast and is easy to backport. #23696 will rebase on top so it carries only the fully-matched `RowFilter` skip performance optimization. cc @alamb @adriangb @hhhizzz
1 parent 1f0615a commit 574fe67

3 files changed

Lines changed: 286 additions & 2 deletions

File tree

datafusion/core/tests/parquet/dynamic_row_group_pruning.rs

Lines changed: 114 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@
3232
3333
use std::sync::Arc;
3434

35-
use arrow::array::{ArrayRef, Int64Array, RecordBatch};
35+
use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray};
3636
use arrow_schema::{DataType, Field, Schema};
3737

3838
use crate::parquet::Unit::RowGroup;
@@ -585,3 +585,116 @@ async fn dynamic_rg_pruning_fires_for_multi_column_sort_leading_tied() {
585585
output.description(),
586586
);
587587
}
588+
589+
/// Build the #24352 fixture: four 2048-row row groups where the filter column
590+
/// (`search_phrase`) differs from the sort column (`event_time`), and one row
591+
/// group (the second) has an empty post-predicate selection invisible to
592+
/// statistics — its only small `event_time` (50) sits on the row whose
593+
/// `search_phrase` is `''`.
594+
///
595+
/// RG 0: event_time = i*1000 (i in 0..2048)
596+
/// RG 1: i=2048 -> (50, ''), else (20000+i, 'p'||i) (i in 2048..4096)
597+
/// RG 2: event_time = 100 + (i-4096) (i in 4096..6144)
598+
/// RG 3: event_time = 5000 + (i-6144) (i in 6144..8192)
599+
fn build_q26_batches(schema: &Arc<Schema>) -> Vec<RecordBatch> {
600+
(0..4i64)
601+
.map(|rg| {
602+
let mut event_time = Vec::with_capacity(2048);
603+
let mut search_phrase: Vec<String> = Vec::with_capacity(2048);
604+
for j in 0..2048i64 {
605+
let i = rg * 2048 + j;
606+
let (et, sp) = if i < 2048 {
607+
(i * 1000, format!("p{i}"))
608+
} else if i < 4096 {
609+
if i == 2048 {
610+
(50, String::new())
611+
} else {
612+
(20000 + i, format!("p{i}"))
613+
}
614+
} else if i < 6144 {
615+
(100 + (i - 4096), format!("p{i}"))
616+
} else {
617+
(5000 + (i - 6144), format!("p{i}"))
618+
};
619+
event_time.push(et);
620+
search_phrase.push(sp);
621+
}
622+
RecordBatch::try_new(
623+
Arc::clone(schema),
624+
vec![
625+
Arc::new(Int64Array::from(event_time)) as ArrayRef,
626+
Arc::new(StringArray::from(search_phrase)) as ArrayRef,
627+
],
628+
)
629+
.unwrap()
630+
})
631+
.collect()
632+
}
633+
634+
/// Regression for #24352: with `pushdown_filters` + TopK dynamic filter, a row
635+
/// group whose post-predicate selection is empty is silently finished by
636+
/// arrow-rs without handing back a reader. Before `rg_plan` was synced to the
637+
/// decoder frontier (`peek_next_row_group`), it trailed the decoder by one, so
638+
/// a later runtime prune rebuilt the decoder from a stale plan and re-read an
639+
/// already-delivered row group — the duplicate rows displaced the true top-k.
640+
#[tokio::test]
641+
async fn topk_pushdown_does_not_reread_delivered_row_group() {
642+
let schema = Arc::new(Schema::new(vec![
643+
Field::new("event_time", DataType::Int64, false),
644+
Field::new("search_phrase", DataType::Utf8, false),
645+
]));
646+
let batches = build_q26_batches(&schema);
647+
648+
// `RowGroup(2048)` writes one row group per 2048-row batch (4 RGs) and
649+
// enables `pushdown_filters`, required for the dynamic filter to reach the
650+
// parquet scan.
651+
let mut ctx = ContextWithParquet::with_custom_data(
652+
Scenario::Int,
653+
RowGroup(2048),
654+
Arc::clone(&schema),
655+
batches,
656+
)
657+
.await;
658+
659+
let output = ctx
660+
.query(
661+
"SELECT search_phrase FROM t \
662+
WHERE search_phrase <> '' ORDER BY event_time LIMIT 10",
663+
)
664+
.await;
665+
666+
// `search_phrase` is unique per row, so any repeated value is the same
667+
// source row emitted twice. The correct answer is the 10 smallest-
668+
// `event_time` non-empty phrases, matching DuckDB / pushdown-off.
669+
assert_eq!(output.result_rows, 10, "{}", output.description());
670+
671+
// The test must actually exercise the runtime prune/rebuild path that
672+
// caused #24352 (not just a happy-path scan), otherwise a future default or
673+
// optimizer change could let it pass without the bug's precondition. Assert
674+
// the dynamic filter pruned at least one row group.
675+
let pruned = output
676+
.row_groups_pruned_dynamic_filter()
677+
.expect("`row_groups_pruned_dynamic_filter` metric must be registered");
678+
assert!(
679+
pruned >= 1,
680+
"test must exercise dynamic RG pruning (the #24352 path); pruned={pruned}\n{}",
681+
output.description(),
682+
);
683+
684+
let formatted = output.pretty_results();
685+
for p in [
686+
"p0", "p4096", "p4097", "p4098", "p4099", "p4100", "p4101", "p4102", "p4103",
687+
"p4104",
688+
] {
689+
assert!(
690+
formatted.contains(&format!("| {p} ")),
691+
"missing {p} from top-k; got:\n{formatted}",
692+
);
693+
}
694+
// The bug emitted p4096 twice (and dropped p4101..=p4104); assert no dup.
695+
assert_eq!(
696+
formatted.matches("| p4096 ").count(),
697+
1,
698+
"p4096 emitted more than once — rg_plan/decoder desync; got:\n{formatted}",
699+
);
700+
}

datafusion/datasource-parquet/src/push_decoder.rs

Lines changed: 88 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ use parquet::arrow::async_reader::AsyncFileReader;
5353
use parquet::arrow::push_decoder::{ParquetPushDecoder, ParquetPushDecoderBuilder};
5454
use parquet::file::metadata::ParquetMetaData;
5555

56-
use datafusion_common::{DataFusionError, Result};
56+
use datafusion_common::{DataFusionError, Result, internal_err};
5757
use datafusion_physical_expr::expressions::DynamicFilterTracking;
5858
use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
5959
use datafusion_physical_plan::metrics::{BaselineMetrics, Count, Gauge};
@@ -342,6 +342,20 @@ impl PushDecoderStreamState {
342342
.as_ref()
343343
.expect("decoder present")
344344
.is_at_row_group_boundary();
345+
// Only the runtime pruner rebuilds the decoder from `rg_plan`, so
346+
// only it needs `rg_plan` kept in sync with the decoder frontier.
347+
// arrow-rs silently finishes row groups whose post-predicate
348+
// selection is empty without handing back a reader, so without this
349+
// sync `rg_plan` trails the decoder by one and a rebuild re-reads an
350+
// already-delivered row group (#24352). Gating on the pruner also
351+
// avoids the O(remaining row groups) cost of `peek_next_row_group()`
352+
// on ordinary scans that never rebuild.
353+
if at_boundary
354+
&& self.row_group_pruner.is_some()
355+
&& let Err(e) = self.sync_rg_plan_to_decoder_frontier()
356+
{
357+
return Some((Err(e), self));
358+
}
345359
if at_boundary && !self.rg_plan.is_empty() {
346360
let mut pruned_count = 0usize;
347361
if let Some(pruner) = self.row_group_pruner.as_mut() {
@@ -414,6 +428,51 @@ impl PushDecoderStreamState {
414428
}
415429
}
416430

431+
/// Keep `rg_plan.front()` aligned with the row group the decoder will emit
432+
/// next. `try_next_reader` silently finishes row groups whose post-predicate
433+
/// selection is empty (no reader handed back), which would otherwise leave
434+
/// `rg_plan` trailing the decoder by one — a later prune/rebuild would then
435+
/// re-include an already-delivered row group (#24352).
436+
fn sync_rg_plan_to_decoder_frontier(&mut self) -> Result<()> {
437+
match self
438+
.decoder
439+
.as_ref()
440+
.expect("decoder present")
441+
.peek_next_row_group()
442+
.map_err(DataFusionError::from)?
443+
{
444+
Some(actual) => Self::advance_rg_plan_to(&mut self.rg_plan, actual)?,
445+
// Decoder has nothing left to emit — drain our plan so the stream
446+
// finishes cleanly.
447+
None => self.rg_plan.clear(),
448+
}
449+
Ok(())
450+
}
451+
452+
/// Pop entries off `rg_plan` until its front is `target`.
453+
///
454+
/// `target` is the RG the decoder will emit next and must still be in the
455+
/// plan. A missing `target` means the decoder's frontier and `rg_plan` have
456+
/// diverged; we surface that as an internal error rather than silently
457+
/// draining the plan, which would truncate the scan. Kept free-standing on
458+
/// `rg_plan` (rather than `&mut self`) so the pop/guard logic is
459+
/// unit-testable without constructing a full stream state.
460+
fn advance_rg_plan_to(
461+
rg_plan: &mut VecDeque<RgPlanEntry>,
462+
target: usize,
463+
) -> Result<()> {
464+
while let Some(front) = rg_plan.front() {
465+
if front.rg_index == target {
466+
return Ok(());
467+
}
468+
rg_plan.pop_front();
469+
}
470+
internal_err!(
471+
"push decoder frontier RG {target} is not in rg_plan; \
472+
decoder and plan have diverged"
473+
)
474+
}
475+
417476
/// Copies metrics from ArrowReaderMetrics (the metrics collected by the
418477
/// arrow-rs parquet reader) to the parquet file metrics for DataFusion
419478
fn copy_arrow_reader_metrics(&self) {
@@ -607,4 +666,32 @@ mod tests {
607666
assert!(!pruner.should_prune(&[1]));
608667
assert!(!pruner.should_prune(&[2]));
609668
}
669+
670+
#[test]
671+
fn advance_rg_plan_to_pops_up_to_target() {
672+
let mut plan: VecDeque<RgPlanEntry> = [0usize, 1, 2, 3]
673+
.into_iter()
674+
.map(|rg_index| RgPlanEntry { rg_index })
675+
.collect();
676+
PushDecoderStreamState::advance_rg_plan_to(&mut plan, 2).unwrap();
677+
assert_eq!(
678+
plan.iter().map(|e| e.rg_index).collect::<Vec<_>>(),
679+
vec![2, 3],
680+
"must pop the entries before `target` and stop at it",
681+
);
682+
}
683+
684+
#[test]
685+
fn advance_rg_plan_to_errors_when_target_absent() {
686+
let mut plan: VecDeque<RgPlanEntry> = [0usize, 1, 2]
687+
.into_iter()
688+
.map(|rg_index| RgPlanEntry { rg_index })
689+
.collect();
690+
let err = PushDecoderStreamState::advance_rg_plan_to(&mut plan, 5)
691+
.expect_err("a target absent from the plan must be an internal error");
692+
assert!(
693+
err.to_string().contains("diverged"),
694+
"expected a divergence internal error, got: {err}",
695+
);
696+
}
610697
}

datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,3 +110,87 @@ RESET datafusion.execution.parquet.pushdown_filters;
110110

111111
statement ok
112112
RESET datafusion.explain.analyze_level;
113+
114+
# Regression test for #24352: TopK dynamic filter + `pushdown_filters` must not
115+
# re-read an already-delivered row group. The filter column (`search_phrase`)
116+
# differs from the sort column (`event_time`), and one row group has an empty
117+
# post-predicate selection that row-group statistics cannot see — its only small
118+
# `event_time` (50) sits on the row where `search_phrase = ''`. arrow-rs finishes
119+
# that RG without handing back a reader; without syncing `rg_plan` to the decoder
120+
# frontier via `peek_next_row_group`, `rg_plan` trailed the decoder by one, so a
121+
# later runtime prune rebuilt the decoder from a stale plan and re-read an
122+
# already-delivered RG — duplicating rows and dropping the true top-k tail.
123+
statement ok
124+
set datafusion.execution.parquet.pushdown_filters = true;
125+
126+
statement ok
127+
set datafusion.execution.target_partitions = 1;
128+
129+
# Both dynamic-filter switches are on by default; set them explicitly so this
130+
# test keeps exercising the prune/rebuild path even if the defaults change.
131+
statement ok
132+
set datafusion.optimizer.enable_dynamic_filter_pushdown = true;
133+
134+
statement ok
135+
set datafusion.optimizer.enable_topk_dynamic_filter_pushdown = true;
136+
137+
statement ok
138+
CREATE TABLE q26_src AS
139+
SELECT
140+
CAST(CASE
141+
WHEN i < 2048 THEN i * 1000
142+
WHEN i < 4096 THEN (CASE WHEN i = 2048 THEN 50 ELSE 20000 + i END)
143+
WHEN i < 6144 THEN 100 + (i - 4096)
144+
ELSE 5000 + (i - 6144)
145+
END AS BIGINT) AS event_time,
146+
CASE WHEN i = 2048 THEN '' ELSE 'p' || CAST(i AS VARCHAR) END AS search_phrase
147+
FROM generate_series(0, 8191) AS t(i);
148+
149+
statement ok
150+
COPY (SELECT * FROM q26_src)
151+
TO 'test_files/scratch/dynamic_row_group_pruning/q26.parquet'
152+
STORED AS PARQUET
153+
OPTIONS ('format.max_row_group_size' '2048');
154+
155+
statement ok
156+
drop table q26_src;
157+
158+
statement ok
159+
CREATE EXTERNAL TABLE q26 (event_time BIGINT NOT NULL, search_phrase VARCHAR NOT NULL)
160+
STORED AS PARQUET
161+
LOCATION 'test_files/scratch/dynamic_row_group_pruning/q26.parquet';
162+
163+
# Each search_phrase is unique, so any repeated value would be the same source
164+
# row emitted twice. The result must be the 10 smallest-`event_time` non-empty
165+
# phrases with no duplicates (matches DuckDB and pushdown-off DataFusion).
166+
query T
167+
SELECT search_phrase FROM q26 WHERE search_phrase <> '' ORDER BY event_time LIMIT 10;
168+
----
169+
p0
170+
p4096
171+
p4097
172+
p4098
173+
p4099
174+
p4100
175+
p4101
176+
p4102
177+
p4103
178+
p4104
179+
180+
statement ok
181+
drop table q26;
182+
183+
statement ok
184+
RESET datafusion.execution.parquet.pushdown_filters;
185+
186+
# The SLT runner sets `target_partitions` to 4 instead of using the default, so
187+
# restore it explicitly rather than RESET (which would revert to the system
188+
# default = num_cpus and leak modified config out of this file).
189+
statement ok
190+
set datafusion.execution.target_partitions = 4;
191+
192+
statement ok
193+
RESET datafusion.optimizer.enable_dynamic_filter_pushdown;
194+
195+
statement ok
196+
RESET datafusion.optimizer.enable_topk_dynamic_filter_pushdown;

0 commit comments

Comments
 (0)