fix(topk): call attempt_early_completion when filter rejects entire b…

[ajegou]

…atch

Which issue does this PR close?

• Closes TopK early-exit doesn't fire when heap's dynamic filter rejects an entire batch #22849.

Rationale for this change

TopK::insert_batch short-circuits when the heap's dynamic filter rejects every row in a batch:

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.

What changes are included in this PR?

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.

Are these changes tested?

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).

Are there any user-facing changes?

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.

[kosiew]

@ajegou
Thanks for the update. This looks good to me overall. I left a couple of small suggestions that might make the new TopK prefix handling a bit easier to read and maintain.

[kosiew]

This comment is accurate, but it feels a bit longer than the control flow it explains. I think we could keep just the key invariant here: the heap is unchanged, but the rejected batch can still prove prefix completion.

For example:

// 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(());

[kosiew]

Nice regression coverage. One small maintainability thought: this repeats a lot of the setup from test_try_finish_marks_finished_with_prefix. A small helper for the (a, b) schema, sort expressions, and TopK construction could make future TopK prefix tests shorter and keep the main scenario easier to spot.