Skip to content

perf: skip wasted per-batch work in TopK insert_batch - #25657

Merged
neilconway merged 11 commits into
apache:mainfrom
dwsmith1983:perf/topk-batch-overhead
Sep 28, 2026
Merged

neilconway merged 11 commits into
apache:mainfrom
dwsmith1983:perf/topk-batch-overhead

Conversation

@dwsmith1983

@dwsmith1983 dwsmith1983 commented Sep 23, 2026 •

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

Once the TopK heap is full, most input batches are rejected outright by the threshold filter. insert_batch still evaluated every sort expression over the whole batch before checking the filter, then evaluated the same expressions again inside the filter, and cloned the batch even when nothing from it was kept. With an expression sort key that doubles the per-batch cost for no output.

What changes are included in this PR?

  • The threshold filter runs before the sort keys are evaluated; fully rejected batches skip sort-key evaluation.
  • The batch moves into the heap entry instead of being cloned. The last row's common prefix is encoded before the move so early completion still compares against the same boundary at the same point.
  • The per-batch timer clones one metric instead of the whole baseline set, here and in PartitionedTopK, PartitionedTopKRank, and PartitionedTopKDenseRank.

SELECT l_orderkey, l_extendedprice FROM lineitem ORDER BY l_extendedprice + 0 DESC LIMIT 1 on TPC-H SF1, TopK elapsed_compute, median of five: 152 ms before, 94 ms after. sort-tpch --limit 10 shows no change.

One nuance: a sort-key expression that errors only on rows the threshold already rejects no longer raises, because those batches skip sort-key evaluation. The filter evaluates the same expression over the batch first, so this can only differ for a secondary key under short-circuit evaluation.

What is the testing strategy for this PR?

New unit tests in datafusion/physical-plan/src/topk: test_batch_rejected_by_threshold_leaves_heap_untouched and test_expression_sort_key_matches_column_sort_key. Existing TopK tests and the topk, window and limit sqllogictest files pass unchanged.

Are there any user-facing changes?

No.

Evaluate the dynamic threshold filter before computing sort keys, so a
batch the filter rejects outright no longer pays for them. With an
expression key such as `l_extendedprice + 0` those keys were otherwise
computed twice per batch, once directly and once inside the filter.

Also move the input batch into the heap entry instead of cloning it,
encoding the last-row prefix needed for early completion before the
heap takes ownership, and clone only the elapsed_compute timer rather
than the whole BaselineMetrics in TopK and PartitionedTopK.

TopK elapsed_compute on TPC-H SF1 lineitem, ORDER BY l_extendedprice + 0
DESC LIMIT 1, median of 5: 151.97 ms before, 93.58 ms after (-38%).
sort-tpch --limit 10 shows no regression.
@github-actions github-actions Bot added the physical-plan Changes to the physical-plan crate label Sep 23, 2026
@codecov-commenter

codecov-commenter commented Sep 26, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 84.26966% with 28 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.55%. Comparing base (991fd23) to head (2c25298).

Files with missing lines Patch % Lines
datafusion/physical-plan/src/topk/mod.rs 84.26% 4 Missing and 24 partials ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main   #25657      +/-   ##
==========================================
- Coverage   82.55%   82.55%   -0.01%     
==========================================
  Files        1141     1141              
  Lines      440209   440365     +156     
  Branches   440209   440365     +156     
==========================================
+ Hits       363409   363533     +124     
- Misses      54829    54837       +8     
- Partials    21971    21995      +24     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@neilconway
neilconway self-requested a review September 27, 2026 16:50

@neilconway neilconway left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@dwsmith1983 Nice PR. Overall lgtm! One minor suggestion for consistency.

value.into_array(batch.num_rows())
})
.collect::<Result<Vec<_>>>()?;
let elapsed_compute = self.metrics.baseline.elapsed_compute().clone();

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we make the same fix in PartitionedTopKRank and PartitionedTopKDenseRank?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 2c25298: PartitionedTopKRank and PartitionedTopKDenseRank now take the same elapsed_compute handle as the other two insert paths. Also merged main.

PartitionedTopKRank and PartitionedTopKDenseRank still copied the whole
BaselineMetrics set on every insert_batch call just to start a scoped
timer. Take a handle to elapsed_compute alone, matching TopK and
PartitionedTopK.
@neilconway

Copy link
Copy Markdown
Contributor

Thanks @dwsmith1983 !

@neilconway
neilconway added this pull request to the merge queue Sep 28, 2026
Merged via the queue into apache:main with commit 8bb2758 Sep 28, 2026
43 checks passed
Omega359 pushed a commit to Omega359/arrow-datafusion that referenced this pull request Oct 11, 2026
## Which issue does this PR close?

- Related to apache#17098.

## Rationale for this change

Once the TopK heap is full, most input batches are rejected outright by
the threshold filter. `insert_batch` still evaluated every sort
expression over the whole batch before checking the filter, then
evaluated the same expressions again inside the filter, and cloned the
batch even when nothing from it was kept. With an expression sort key
that doubles the per-batch cost for no output.

## What changes are included in this PR?

- The threshold filter runs before the sort keys are evaluated; fully
rejected batches skip sort-key evaluation.
- The batch moves into the heap entry instead of being cloned. The last
row's common prefix is encoded before the move so early completion still
compares against the same boundary at the same point.
- The per-batch timer clones one metric instead of the whole baseline
set, here and in `PartitionedTopK`, `PartitionedTopKRank`, and
`PartitionedTopKDenseRank`.

`SELECT l_orderkey, l_extendedprice FROM lineitem ORDER BY
l_extendedprice + 0 DESC LIMIT 1` on TPC-H SF1, TopK `elapsed_compute`,
median of five: 152 ms before, 94 ms after. `sort-tpch --limit 10` shows
no change.

One nuance: a sort-key expression that errors only on rows the threshold
already rejects no longer raises, because those batches skip sort-key
evaluation. The filter evaluates the same expression over the batch
first, so this can only differ for a secondary key under short-circuit
evaluation.

## What is the testing strategy for this PR?

New unit tests in `datafusion/physical-plan/src/topk`:
`test_batch_rejected_by_threshold_leaves_heap_untouched` and
`test_expression_sort_key_matches_column_sort_key`. Existing TopK tests
and the `topk`, `window` and `limit` sqllogictest files pass unchanged.

## Are there any user-facing changes?

No.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

physical-plan Changes to the physical-plan crate v56.0.0

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants