fix(parquet): sync rg_plan to decoder frontier — fix wrong TopK results from re-reading already-delivered row groups (#24352)

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<Schema>) -> Vec<RecordBatch> {
+    (0..4i64)
+        .map(|rg| {
+            let mut event_time = Vec::with_capacity(2048);
+            let mut search_phrase: Vec<String> = 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}",
+    );
+}

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<RgPlanEntry>,
+        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<RgPlanEntry> = [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<_>>(),
+            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<RgPlanEntry> = [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}",
+        );
+    }
 }

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;