From 5a8804cd674dbcfa2711ad0239b2346a75803624 Mon Sep 17 00:00:00 2001 From: ajegou Date: Mon, 15 Jun 2026 10:04:34 +0200 Subject: [PATCH] fix(topk): call attempt_early_completion when filter rejects entire batch (#22852) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Closes #22849 - A related cross-partition starvation case is tracked separately in [discussion](https://github.com/apache/datafusion/pull/22852#issuecomment-4670382915) for details `TopK::insert_batch` short-circuits when the heap's dynamic filter rejects every row in a batch: ```rust if !filter.has_true() { // nothing to filter, so no need to update return Ok(()); } ``` The early-exit check `attempt_early_completion(&batch)` lives later in the same function, gated on `replacements > 0`. So a batch that the filter rejects entirely bypasses the check. The heap's dynamic filter is derived from the heap's worst row (via `update_filter`). A batch whose rows all come from a strictly worse sort prefix is exactly the batch the filter rejects entirely — i.e. the very signal `attempt_early_completion` is designed to detect ("the next batch is past the heap's boundary, we can stop") is what causes the function to short-circuit *before* the check runs. This is a feature-interaction regression between two PRs that were both correct in isolation. The `attempt_early_completion` mechanism was added by #15563 (closing #15529). At the time, there was no heap-derived dynamic filter on TopK, so the only sensible call site was right after a successful heap insertion. Two months later, #15770 added the dynamic-filter pushdown for TopK sorts, introducing the `!filter.has_true()` short-circuit. The two features address different problems and the new short-circuit didn't connect to the existing prefix-completion check — which is how this gap opened up. **Consequence**: on a TopK over an input ordered on the sort prefix, `finished = true` is never set once the heap stabilizes. Since `finished` is the signal `SortExec` uses to stop pulling from its input (via `Poll::Ready(None)` from the TopK stream, which cascades into dropping the source stream), the source keeps being polled long past the point where no further row can improve the heap. The LIMIT optimization effectively degrades to "heap saves memory but reads everything"; sources with cancellable streams (e.g. networked sources) never receive the cancellation signal. Single behavioral change in `datafusion/physical-plan/src/topk/mod.rs`: call `attempt_early_completion(&batch)` immediately before the `return Ok(())` in the `!filter.has_true()` branch. Why this scope, not a broader restructuring: - The existing `attempt_early_completion` call inside `if replacements > 0` is load-bearing for a related case: a batch containing a mix of "still valuable" rows and "past the boundary" rows. The existing `test_try_finish_marks_finished_with_prefix` test covers this case — Batch 2 with `a=[2,3], b=[10,20]` against a heap where `heap.max.a = 2`; the `(2, 10)` row must be inserted before the check on the `(3, 20)` last row triggers. Moving the call earlier would skip the insertion of valuable rows and break that test. - The bug is specifically that the *short-circuit* path doesn't call the check. The fix targets exactly that path. - A related but separate gap is not addressed here: when `filter.has_true() == true` but `replacements == 0` (the filter accepts some rows but `find_new_topk_items` ends up inserting none of them), the existing call inside `if replacements > 0` is also skipped. This requires a divergence between the heap's filter predicate and the row-byte comparison used inside `find_new_topk_items`, which shouldn't normally happen (the filter is derived from the heap's worst row using the same comparator). A deterministic synthetic repro would likely require concurrent heap updates from sibling partitions or boundary-value edge cases (NaN/NULL semantics, type coercion). Happy to send a follow-up if reviewers want it covered; the workload that motivated this fix was the filter-rejection case empirically. Yes. Added a regression test `test_try_finish_fires_when_filter_rejects_entire_batch`. The assertion target is `topk.finished` — the flag that signals "stop pulling from the source" to upstream consumers (read by `TopKExec::poll_next` to emit `Poll::Ready(None)`). Asserting that the flag transitions on the fully-filter-rejected batch is equivalent to asserting that the source-stopping mechanism activates. - Builds a TopK over a `(a, b)` sort with prefix `a`, k=3. - Inserts a batch that fills the heap with rows from `a ∈ {1, 2}`; `update_filter` tightens the filter to `a < 2 OR (a = 2 AND b < 30)`. - Inserts a second batch with all rows at `a = 3` — filter rejects every row. - Without the fix: `insert_batch` short-circuits, `topk.finished` stays `false`. Test fails. - With the fix: `attempt_early_completion` fires (last-row prefix `a = 3` > heap.max prefix `a = 2`), `topk.finished` becomes `true`. Test passes. The test also asserts the emitted top-K is unchanged from after batch 1, confirming no candidate row was incorrectly excluded by the early bail. All 28 existing `topk::` tests continue to pass (including `test_try_finish_marks_finished_with_prefix`, which exercises the mixed-prefix case). No public API or output changes. The fix only changes when TopK marks itself `finished = true` — specifically, it now fires `attempt_early_completion` for batches that are entirely rejected by the heap's dynamic filter, where previously it would silently skip the check. Output of TopK is unchanged; only the early-exit behavior improves. --------- Co-authored-by: Gabriel <45515538+gabotechs@users.noreply.github.com> (cherry picked from commit 6520315d41851d1fb31da0ae1b4f22e48a6b2705) --- datafusion/physical-plan/src/topk/mod.rs | 101 ++++++++++++++++++----- 1 file changed, 82 insertions(+), 19 deletions(-) diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs index e0b91f25161c0..a8e6078cc4f29 100644 --- a/datafusion/physical-plan/src/topk/mod.rs +++ b/datafusion/physical-plan/src/topk/mod.rs @@ -255,7 +255,9 @@ impl TopK { let mut filter = array.as_boolean().clone(); let true_count = filter.true_count(); if true_count == 0 { - // nothing to filter, so no need to update + // The heap is unchanged, but a fully rejected batch can still prove + // that the shared sort prefix has passed the heap boundary. + self.attempt_early_completion(&batch)?; return Ok(()); } // only update the keys / rows if the filter does not match all rows @@ -1098,20 +1100,15 @@ mod tests { assert_eq!(record_batch_store.batches_size, 0); } - /// This test validates that the `try_finish` method marks the TopK operator as finished - /// when the prefix (on column "a") of the last row in the current batch is strictly greater - /// than the max top‑k row. - /// The full sort expression is defined on both columns ("a", "b"), but the input ordering is only on "a". - #[tokio::test] - async fn test_try_finish_marks_finished_with_prefix() -> Result<()> { - // Create a schema with two columns. + /// Builds an `(a Int32, b Float64)` schema and a `TopK` with full sort + /// `(a ASC, b ASC)`, input prefix `[a]`, `k = 3`, `batch_size = 2`. Used by + /// the prefix-completion tests below to keep their per-scenario logic in focus. + fn build_ab_prefix_topk() -> Result<(Arc, TopK)> { let schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::Int32, false), Field::new("b", DataType::Float64, false), ])); - // Create sort expressions. - // Full sort: first by "a", then by "b". let sort_expr_a = PhysicalSortExpr { expr: col("a", schema.as_ref())?, options: SortOptions::default(), @@ -1121,28 +1118,33 @@ mod tests { options: SortOptions::default(), }; - // Input ordering uses only column "a" (a prefix of the full sort). + // Input ordering uses only column "a" (a prefix of the full sort on (a, b)). let prefix = vec![sort_expr_a.clone()]; let full_expr = LexOrdering::from([sort_expr_a, sort_expr_b]); - // Create a dummy runtime environment and metrics. - let runtime = Arc::new(RuntimeEnv::default()); - let metrics = ExecutionPlanMetricsSet::new(); - - // Create a TopK instance with k = 3 and batch_size = 2. - let mut topk = TopK::try_new( + let topk = TopK::try_new( 0, Arc::clone(&schema), prefix, full_expr, 3, 2, - runtime, - &metrics, + Arc::new(RuntimeEnv::default()), + &ExecutionPlanMetricsSet::new(), Arc::new(RwLock::new(TopKDynamicFilters::new(Arc::new( DynamicFilterPhysicalExpr::new(vec![], lit(true)), )))), )?; + Ok((schema, topk)) + } + + /// This test validates that the `try_finish` method marks the TopK operator as finished + /// when the prefix (on column "a") of the last row in the current batch is strictly greater + /// than the max top‑k row. + /// The full sort expression is defined on both columns ("a", "b"), but the input ordering is only on "a". + #[tokio::test] + async fn test_try_finish_marks_finished_with_prefix() -> Result<()> { + let (schema, mut topk) = build_ab_prefix_topk()?; // Create the first batch with two columns: // Column "a": [1, 1, 2], Column "b": [20.0, 15.0, 30.0]. @@ -1195,6 +1197,67 @@ mod tests { Ok(()) } + /// Regression test for #22849: a batch whose rows are entirely rejected by the + /// heap's dynamic filter must still trigger `attempt_early_completion` when its + /// last row's prefix is worse than the heap's worst. + /// + /// Before the fix, the `!filter.has_true()` short-circuit returned without calling + /// `attempt_early_completion`. Because the heap's filter is itself derived from the + /// heap's worst row, a batch from a strictly-worse prefix is exactly the case the + /// filter rejects entirely — i.e. the very signal the early-exit was designed to + /// detect was being silently dropped. + #[tokio::test] + async fn test_try_finish_fires_when_filter_rejects_entire_batch() -> Result<()> { + let (schema, mut topk) = build_ab_prefix_topk()?; + + // Batch 1 fills the heap with (1, 20.0), (1, 15.0), (2, 30.0). + // heap.max becomes (a=2, b=30.0); update_filter tightens the heap filter to + // a < 2 OR (a = 2 AND b < 30.0). + let array_a1: ArrayRef = + Arc::new(Int32Array::from(vec![Some(1), Some(1), Some(2)])); + let array_b1: ArrayRef = Arc::new(Float64Array::from(vec![20.0, 15.0, 30.0])); + let batch1 = RecordBatch::try_new(Arc::clone(&schema), vec![array_a1, array_b1])?; + topk.insert_batch(batch1)?; + assert!( + !topk.finished, + "Expected 'finished' to be false after batch 1 \ + (last row prefix a=2 equals heap.max prefix a=2, not strictly greater)." + ); + + // Batch 2: every row has a=3, so the heap's filter (a < 2 OR (a = 2 AND b < 30)) + // rejects every row. Before the fix, `insert_batch` would short-circuit on + // `!filter.has_true()` and return without checking the prefix; `finished` + // would stay false even though no future batch could improve the heap. + let array_a2: ArrayRef = Arc::new(Int32Array::from(vec![Some(3), Some(3)])); + let array_b2: ArrayRef = Arc::new(Float64Array::from(vec![10.0, 20.0])); + let batch2 = RecordBatch::try_new(Arc::clone(&schema), vec![array_a2, array_b2])?; + topk.insert_batch(batch2)?; + assert!( + topk.finished, + "Expected 'finished' to be true after batch 2 \ + (filter rejected every row, but the batch's last row prefix a=3 \ + is strictly greater than heap.max prefix a=2)." + ); + + // The emitted top-k is unchanged from after batch 1 since none of batch 2's + // rows could improve the heap. + let results: Vec<_> = topk.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+---+------+", + "| a | b |", + "+---+------+", + "| 1 | 15.0 |", + "| 1 | 20.0 |", + "| 2 | 30.0 |", + "+---+------+", + ], + &results + ); + + Ok(()) + } + /// This test verifies that the dynamic filter is marked as complete after TopK processing finishes. #[tokio::test] async fn test_topk_marks_filter_complete() -> Result<()> {