diff --git a/datafusion/core/tests/parquet/dynamic_row_group_pruning.rs b/datafusion/core/tests/parquet/dynamic_row_group_pruning.rs index d5d648be9b7aa..917faaaa5ce64 100644 --- a/datafusion/core/tests/parquet/dynamic_row_group_pruning.rs +++ b/datafusion/core/tests/parquet/dynamic_row_group_pruning.rs @@ -32,7 +32,7 @@ use std::sync::Arc; -use arrow::array::{ArrayRef, Int64Array, RecordBatch}; +use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray}; use arrow_schema::{DataType, Field, Schema}; use crate::parquet::Unit::RowGroup; @@ -585,3 +585,116 @@ async fn dynamic_rg_pruning_fires_for_multi_column_sort_leading_tied() { output.description(), ); } + +/// Build the #24352 fixture: four 2048-row row groups where the filter column +/// (`search_phrase`) differs from the sort column (`event_time`), and one row +/// group (the second) has an empty post-predicate selection invisible to +/// statistics — its only small `event_time` (50) sits on the row whose +/// `search_phrase` is `''`. +/// +/// RG 0: event_time = i*1000 (i in 0..2048) +/// RG 1: i=2048 -> (50, ''), else (20000+i, 'p'||i) (i in 2048..4096) +/// RG 2: event_time = 100 + (i-4096) (i in 4096..6144) +/// RG 3: event_time = 5000 + (i-6144) (i in 6144..8192) +fn build_q26_batches(schema: &Arc) -> Vec { + (0..4i64) + .map(|rg| { + let mut event_time = Vec::with_capacity(2048); + let mut search_phrase: Vec = Vec::with_capacity(2048); + for j in 0..2048i64 { + let i = rg * 2048 + j; + let (et, sp) = if i < 2048 { + (i * 1000, format!("p{i}")) + } else if i < 4096 { + if i == 2048 { + (50, String::new()) + } else { + (20000 + i, format!("p{i}")) + } + } else if i < 6144 { + (100 + (i - 4096), format!("p{i}")) + } else { + (5000 + (i - 6144), format!("p{i}")) + }; + event_time.push(et); + search_phrase.push(sp); + } + RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int64Array::from(event_time)) as ArrayRef, + Arc::new(StringArray::from(search_phrase)) as ArrayRef, + ], + ) + .unwrap() + }) + .collect() +} + +/// Regression for #24352: with `pushdown_filters` + TopK dynamic filter, a row +/// group whose post-predicate selection is empty is silently finished by +/// arrow-rs without handing back a reader. Before `rg_plan` was synced to the +/// decoder frontier (`peek_next_row_group`), it trailed the decoder by one, so +/// a later runtime prune rebuilt the decoder from a stale plan and re-read an +/// already-delivered row group — the duplicate rows displaced the true top-k. +#[tokio::test] +async fn topk_pushdown_does_not_reread_delivered_row_group() { + let schema = Arc::new(Schema::new(vec![ + Field::new("event_time", DataType::Int64, false), + Field::new("search_phrase", DataType::Utf8, false), + ])); + let batches = build_q26_batches(&schema); + + // `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( + Scenario::Int, + RowGroup(2048), + Arc::clone(&schema), + batches, + ) + .await; + + let output = ctx + .query( + "SELECT search_phrase FROM t \ + WHERE search_phrase <> '' ORDER BY event_time LIMIT 10", + ) + .await; + + // `search_phrase` is unique per row, so any repeated value is the same + // source row emitted twice. The correct answer is the 10 smallest- + // `event_time` non-empty phrases, matching DuckDB / pushdown-off. + assert_eq!(output.result_rows, 10, "{}", output.description()); + + // The test must actually exercise the runtime prune/rebuild path that + // caused #24352 (not just a happy-path scan), otherwise a future default or + // optimizer change could let it pass without the bug's precondition. Assert + // the dynamic filter pruned at least one row group. + let pruned = output + .row_groups_pruned_dynamic_filter() + .expect("`row_groups_pruned_dynamic_filter` metric must be registered"); + assert!( + pruned >= 1, + "test must exercise dynamic RG pruning (the #24352 path); pruned={pruned}\n{}", + output.description(), + ); + + let formatted = output.pretty_results(); + for p in [ + "p0", "p4096", "p4097", "p4098", "p4099", "p4100", "p4101", "p4102", "p4103", + "p4104", + ] { + assert!( + formatted.contains(&format!("| {p} ")), + "missing {p} from top-k; got:\n{formatted}", + ); + } + // The bug emitted p4096 twice (and dropped p4101..=p4104); assert no dup. + assert_eq!( + formatted.matches("| p4096 ").count(), + 1, + "p4096 emitted more than once — rg_plan/decoder desync; got:\n{formatted}", + ); +} diff --git a/datafusion/datasource-parquet/src/push_decoder.rs b/datafusion/datasource-parquet/src/push_decoder.rs index 14904bada2cfc..74d8997198872 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -53,7 +53,7 @@ use parquet::arrow::async_reader::AsyncFileReader; use parquet::arrow::push_decoder::{ParquetPushDecoder, ParquetPushDecoderBuilder}; use parquet::file::metadata::ParquetMetaData; -use datafusion_common::{DataFusionError, Result}; +use datafusion_common::{DataFusionError, Result, internal_err}; use datafusion_physical_expr::expressions::DynamicFilterTracking; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_plan::metrics::{BaselineMetrics, Count, Gauge}; @@ -342,6 +342,20 @@ impl PushDecoderStreamState { .as_ref() .expect("decoder present") .is_at_row_group_boundary(); + // Only the runtime pruner rebuilds the decoder from `rg_plan`, so + // only it needs `rg_plan` kept in sync with the decoder frontier. + // arrow-rs silently finishes row groups whose post-predicate + // selection is empty without handing back a reader, so without this + // sync `rg_plan` trails the decoder by one and a rebuild re-reads an + // already-delivered row group (#24352). Gating on the pruner also + // avoids the O(remaining row groups) cost of `peek_next_row_group()` + // on ordinary scans that never rebuild. + if at_boundary + && self.row_group_pruner.is_some() + && let Err(e) = self.sync_rg_plan_to_decoder_frontier() + { + return Some((Err(e), self)); + } if at_boundary && !self.rg_plan.is_empty() { let mut pruned_count = 0usize; if let Some(pruner) = self.row_group_pruner.as_mut() { @@ -414,6 +428,51 @@ impl PushDecoderStreamState { } } + /// Keep `rg_plan.front()` aligned with the row group the decoder will emit + /// next. `try_next_reader` silently finishes row groups whose post-predicate + /// selection is empty (no reader handed back), which would otherwise leave + /// `rg_plan` trailing the decoder by one — a later prune/rebuild would then + /// re-include an already-delivered row group (#24352). + fn sync_rg_plan_to_decoder_frontier(&mut self) -> Result<()> { + match self + .decoder + .as_ref() + .expect("decoder present") + .peek_next_row_group() + .map_err(DataFusionError::from)? + { + Some(actual) => Self::advance_rg_plan_to(&mut self.rg_plan, actual)?, + // Decoder has nothing left to emit — drain our plan so the stream + // finishes cleanly. + None => self.rg_plan.clear(), + } + Ok(()) + } + + /// Pop entries off `rg_plan` until its front is `target`. + /// + /// `target` is the RG the decoder will emit next and must still be in the + /// plan. A missing `target` means the decoder's frontier and `rg_plan` have + /// diverged; we surface that as an internal error rather than silently + /// draining the plan, which would truncate the scan. Kept free-standing on + /// `rg_plan` (rather than `&mut self`) so the pop/guard logic is + /// unit-testable without constructing a full stream state. + fn advance_rg_plan_to( + rg_plan: &mut VecDeque, + target: usize, + ) -> Result<()> { + while let Some(front) = rg_plan.front() { + if front.rg_index == target { + return Ok(()); + } + rg_plan.pop_front(); + } + internal_err!( + "push decoder frontier RG {target} is not in rg_plan; \ + decoder and plan have diverged" + ) + } + /// Copies metrics from ArrowReaderMetrics (the metrics collected by the /// arrow-rs parquet reader) to the parquet file metrics for DataFusion fn copy_arrow_reader_metrics(&self) { @@ -607,4 +666,32 @@ mod tests { assert!(!pruner.should_prune(&[1])); assert!(!pruner.should_prune(&[2])); } + + #[test] + fn advance_rg_plan_to_pops_up_to_target() { + let mut plan: VecDeque = [0usize, 1, 2, 3] + .into_iter() + .map(|rg_index| RgPlanEntry { rg_index }) + .collect(); + PushDecoderStreamState::advance_rg_plan_to(&mut plan, 2).unwrap(); + assert_eq!( + plan.iter().map(|e| e.rg_index).collect::>(), + vec![2, 3], + "must pop the entries before `target` and stop at it", + ); + } + + #[test] + fn advance_rg_plan_to_errors_when_target_absent() { + let mut plan: VecDeque = [0usize, 1, 2] + .into_iter() + .map(|rg_index| RgPlanEntry { rg_index }) + .collect(); + let err = PushDecoderStreamState::advance_rg_plan_to(&mut plan, 5) + .expect_err("a target absent from the plan must be an internal error"); + assert!( + err.to_string().contains("diverged"), + "expected a divergence internal error, got: {err}", + ); + } } diff --git a/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt b/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt index 2149cacfc0a55..81d9511839725 100644 --- a/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt +++ b/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt @@ -110,3 +110,87 @@ RESET datafusion.execution.parquet.pushdown_filters; statement ok RESET datafusion.explain.analyze_level; + +# Regression test for #24352: TopK dynamic filter + `pushdown_filters` must not +# re-read an already-delivered row group. The filter column (`search_phrase`) +# differs from the sort column (`event_time`), and one row group has an empty +# post-predicate selection that row-group statistics cannot see — its only small +# `event_time` (50) sits on the row where `search_phrase = ''`. arrow-rs finishes +# that RG without handing back a reader; without syncing `rg_plan` to the decoder +# frontier via `peek_next_row_group`, `rg_plan` trailed the decoder by one, so a +# later runtime prune rebuilt the decoder from a stale plan and re-read an +# already-delivered RG — duplicating rows and dropping the true top-k tail. +statement ok +set datafusion.execution.parquet.pushdown_filters = true; + +statement ok +set datafusion.execution.target_partitions = 1; + +# Both dynamic-filter switches are on by default; set them explicitly so this +# test keeps exercising the prune/rebuild path even if the defaults change. +statement ok +set datafusion.optimizer.enable_dynamic_filter_pushdown = true; + +statement ok +set datafusion.optimizer.enable_topk_dynamic_filter_pushdown = true; + +statement ok +CREATE TABLE q26_src AS +SELECT + CAST(CASE + WHEN i < 2048 THEN i * 1000 + WHEN i < 4096 THEN (CASE WHEN i = 2048 THEN 50 ELSE 20000 + i END) + WHEN i < 6144 THEN 100 + (i - 4096) + ELSE 5000 + (i - 6144) + END AS BIGINT) AS event_time, + CASE WHEN i = 2048 THEN '' ELSE 'p' || CAST(i AS VARCHAR) END AS search_phrase +FROM generate_series(0, 8191) AS t(i); + +statement ok +COPY (SELECT * FROM q26_src) +TO 'test_files/scratch/dynamic_row_group_pruning/q26.parquet' +STORED AS PARQUET +OPTIONS ('format.max_row_group_size' '2048'); + +statement ok +drop table q26_src; + +statement ok +CREATE EXTERNAL TABLE q26 (event_time BIGINT NOT NULL, search_phrase VARCHAR NOT NULL) +STORED AS PARQUET +LOCATION 'test_files/scratch/dynamic_row_group_pruning/q26.parquet'; + +# Each search_phrase is unique, so any repeated value would be the same source +# row emitted twice. The result must be the 10 smallest-`event_time` non-empty +# phrases with no duplicates (matches DuckDB and pushdown-off DataFusion). +query T +SELECT search_phrase FROM q26 WHERE search_phrase <> '' ORDER BY event_time LIMIT 10; +---- +p0 +p4096 +p4097 +p4098 +p4099 +p4100 +p4101 +p4102 +p4103 +p4104 + +statement ok +drop table q26; + +statement ok +RESET datafusion.execution.parquet.pushdown_filters; + +# 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.optimizer.enable_dynamic_filter_pushdown; + +statement ok +RESET datafusion.optimizer.enable_topk_dynamic_filter_pushdown;