Skip to content

Fix shared TopK early exit with shared prefix threshold - #22991

Merged
kosiew merged 6 commits into
apache:mainfrom
geoffreyclaude:fix/topk-global-prefix-early-exit
Jun 22, 2026
Merged

kosiew merged 6 commits into
apache:mainfrom
geoffreyclaude:fix/topk-global-prefix-early-exit

Conversation

@geoffreyclaude

@geoffreyclaude geoffreyclaude commented Jun 17, 2026 •

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Closes #22874.

Follow-up to fix(topk): call attempt_early_completion when filter rejects entire batch.

Rationale for this change

TopK dynamic filter pushdown attempt 2 lets SortExec tighten scan-side predicates while a TopK heap finds better rows. Once the current TopK threshold is known, scans can skip data that cannot enter the final ORDER BY ... LIMIT result.

That works well for a single partition. Partitioned SortExec has one extra case to handle:

  • each output partition has its own local TopK heap
  • those local heaps share one TopKDynamicFilters instance
  • one partition can tighten the shared filter before another partition has enough rows to fill its local heap

fix(topk): call attempt_early_completion when filter rejects entire batch fixed the local case where a heap already has a max row and the dynamic filter rejects a whole batch.

This PR fixes the remaining shared-filter case. A lagging partition can now use the shared prefix threshold to stop early even when its local heap is still empty. If there is no shared threshold yet, it falls back to the existing local heap prefix check.

The shared prefix check is not treating another partition's threshold as this partition's local heap boundary. It uses the same threshold that already drives the shared dynamic filter. Once a partition's ordered input has moved past that shared prefix threshold, later batches from that partition cannot add rows that survive the shared filter. The local heap still emits the candidates it has already kept; this only stops pulling input that can no longer add candidates.

Single-partition behavior is unchanged.

How the TopK optimizations fit together

There are two existing optimizations involved here:

  • Dynamic filter pushdown: once a TopK heap has K rows, its worst kept row becomes a threshold. That threshold tightens a scan-side filter so later data that cannot enter the final ORDER BY ... LIMIT result can be skipped.
  • Prefix early exit: when the input is ordered by a prefix of the requested sort, TopK can stop pulling once the last row in a batch is past a known TopK boundary on that shared prefix.

For a partition-preserving SortExec, those optimizations meet in one shared place. Each output partition has a local TopK heap, but all of those local heaps publish into one shared dynamic filter. A partition that fills first can tighten the shared filter for everyone else. Lagging partitions then need to use that same shared prefix threshold to stop pulling once their ordered input has moved past it.

This PR makes that composition explicit: the shared filter stays alive until every local TopK has emitted, and the shared threshold carries its common-prefix row so lagging partitions can apply the same prefix early-exit check.

What changes are included in this PR?

  • Check early completion when a batch passes the dynamic filter but produces zero heap replacements.
  • Track local TopK emitters so a shared filter completes only after the last emitter has produced output.
  • Store the shared threshold and its common-prefix row together in TopKDynamicFilters.
  • Check the shared prefix in attempt_early_completion before falling back to the local heap prefix.
  • Add focused TopK and SortExec tests for the local zero-replacement path, early completion before local heap fill, equal-prefix non-completion, DESC/null prefix ordering, and shared-filter completion.

How is this split for review?

The commits are ordered so each one has a narrow job:

  1. Check TopK early completion after zero-replacement batches
    Handles the local case where a batch passes the dynamic filter, produces zero heap replacements, but still proves later rows cannot enter the TopK.

  2. Complete shared TopK filters after all emitters
    Keeps a shared dynamic filter watchable until every local TopK emitter has emitted. This includes the SortExec wiring for preserved partitioning.

  3. Use shared TopK prefix thresholds for early exit
    Carries the shared threshold's common-prefix row so lagging partitions can stop before their local heap is full.

Are these changes tested?

Correctness is covered by targeted tests for:

  • the local zero-replacement early-completion path
  • early completion before local heap fill from a shared prefix threshold
  • equal-prefix non-completion
  • DESC and NULLS LAST prefix row ordering
  • shared-filter completion after all TopK emitters, including preserved SortExec partitioning

Relevant background:

Benchmark command:

dfbench sort-tpch --sorted --limit 10 --iterations 5 \
  --path /tmp/df-topk-bench-data/tpch_sf1 \
  -o /tmp/topk-shared-prefix-followup.json

Both sides were rebuilt with fresh isolated release-nonlto target directories before running the benchmark. This is not the regular topk_tpch script: this case needs --sorted --limit 10 to exercise the prefix early-exit path.

Clean rerun against the PR base commit 7bb6e152b. Times are milliseconds, using the average of 5 iterations for each row.

scope PR base this PR change
all sort-tpch queries 760.59 348.84 -54.1%
Q8 49.70 7.66 -84.6%
Q9 92.13 10.27 -88.8%
Q10 89.66 14.35 -84.0%

The DataSourceExec counters show the less noisy part of the result. Q8/Q9/Q10 now emit only the first batch from each partition instead of continuing to drain millions of rows.

query PR base DataSourceExec output_rows this PR DataSourceExec output_rows PR base bytes_scanned this PR bytes_scanned
Q8 3.66M 81.92K 56.81M 15.79M
Q9 3.66M 81.92K 75.19M 20.89M
Q10 3.10M 81.92K 110.9M 34.69M

Are there any user-facing changes?

No. This is an internal physical execution optimization fix.

@github-actions github-actions Bot added the physical-plan Changes to the physical-plan crate label Jun 17, 2026
@geoffreyclaude
geoffreyclaude marked this pull request as draft June 17, 2026 08:20
@geoffreyclaude
geoffreyclaude force-pushed the fix/topk-global-prefix-early-exit branch 5 times, most recently from 5bf6d35 to 0f69ac9 Compare June 17, 2026 11:49
@geoffreyclaude geoffreyclaude changed the title Fix shared TopK early exit with global prefix threshold Fix shared TopK early exit with shared prefix threshold Jun 17, 2026
@geoffreyclaude
geoffreyclaude force-pushed the fix/topk-global-prefix-early-exit branch from 0f69ac9 to 96ba61c Compare June 17, 2026 12:01
@geoffreyclaude

geoffreyclaude commented Jun 17, 2026 •

Copy link
Copy Markdown
Contributor Author

run benchmark partitioned_topk
EDIT: I meant topk_tpch...

@adriangbot

Copy link
Copy Markdown

🤖 Benchmark running (GKE) | trigger
Instance: c4a-highmem-16 (12 vCPU / 65 GiB) | Linux bench-c4730000666-584-rlxkd 6.12.68+ #1 SMP Sat May 2 07:49:07 UTC 2026 aarch64 GNU/Linux

CPU Details (lscpu)
Architecture:                            aarch64
CPU op-mode(s):                          64-bit
Byte Order:                              Little Endian
CPU(s):                                  16
On-line CPU(s) list:                     0-15
Vendor ID:                               ARM
Model name:                              Neoverse-V2
Model:                                   1
Thread(s) per core:                      1
Core(s) per cluster:                     16
Socket(s):                               -
Cluster(s):                              1
Stepping:                                r0p1
BogoMIPS:                                2000.00
Flags:                                   fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache:                               1 MiB (16 instances)
L1i cache:                               1 MiB (16 instances)
L2 cache:                                32 MiB (16 instances)
L3 cache:                                80 MiB (1 instance)
NUMA node(s):                            1
NUMA node0 CPU(s):                       0-15
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Not affected
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; CSV2, BHB
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Not affected

Comparing fix/topk-global-prefix-early-exit (96ba61c) to 7bb6e15 (merge-base) diff using: partitioned_topk
Results will be posted here when complete


File an issue against this benchmark runner

@adriangbot

Copy link
Copy Markdown

Benchmark for this request failed.

Last 20 lines of output:

Click to expand
# Micro-Benchmarks (specific operators and features)
cancellation:           How long cancelling a query takes
nlj:                    Benchmark for simple nested loop joins, testing various join scenarios
hj:                     Benchmark for simple hash joins, testing various join scenarios
smj:                    Benchmark for simple sort merge joins, testing various join scenarios
dict:                   Benchmark for dictionary-encoded group-by scenarios
compile_profile:        Compile and execute TPC-H across selected Cargo profiles, reporting timing and binary size


━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Supported Configuration (Environment Variables)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
DATA_DIR            directory to store datasets
CARGO_COMMAND       command that runs the benchmark binary
DATAFUSION_DIR      directory to use (default /workspace/datafusion-base)
RESULTS_NAME        folder where the benchmark files are stored
PREFER_HASH_JOIN    Prefer hash join algorithm (default true)
SIMULATE_LATENCY    Simulate object store latency to mimic S3 (default false)
DATAFUSION_*        Set the given datafusion configuration


File an issue against this benchmark runner

@geoffreyclaude

Copy link
Copy Markdown
Contributor Author

run benchmark topk_tpch

@adriangbot

Copy link
Copy Markdown

🤖 Benchmark running (GKE) | trigger
Instance: c4a-highmem-16 (12 vCPU / 65 GiB) | Linux bench-c4730553583-585-mm5xj 6.12.68+ #1 SMP Sat May 2 07:49:07 UTC 2026 aarch64 GNU/Linux

CPU Details (lscpu)
Architecture:                            aarch64
CPU op-mode(s):                          64-bit
Byte Order:                              Little Endian
CPU(s):                                  16
On-line CPU(s) list:                     0-15
Vendor ID:                               ARM
Model name:                              Neoverse-V2
Model:                                   1
Thread(s) per core:                      1
Core(s) per cluster:                     16
Socket(s):                               -
Cluster(s):                              1
Stepping:                                r0p1
BogoMIPS:                                2000.00
Flags:                                   fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache:                               1 MiB (16 instances)
L1i cache:                               1 MiB (16 instances)
L2 cache:                                32 MiB (16 instances)
L3 cache:                                80 MiB (1 instance)
NUMA node(s):                            1
NUMA node0 CPU(s):                       0-15
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Not affected
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; CSV2, BHB
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Not affected

Comparing fix/topk-global-prefix-early-exit (96ba61c) to 7bb6e15 (merge-base) diff using: topk_tpch
Results will be posted here when complete


File an issue against this benchmark runner

@adriangbot

Copy link
Copy Markdown

🤖 Benchmark completed (GKE) | trigger

Instance: c4a-highmem-16 (12 vCPU / 65 GiB)

CPU Details (lscpu)
Architecture:                            aarch64
CPU op-mode(s):                          64-bit
Byte Order:                              Little Endian
CPU(s):                                  16
On-line CPU(s) list:                     0-15
Vendor ID:                               ARM
Model name:                              Neoverse-V2
Model:                                   1
Thread(s) per core:                      1
Core(s) per cluster:                     16
Socket(s):                               -
Cluster(s):                              1
Stepping:                                r0p1
BogoMIPS:                                2000.00
Flags:                                   fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache:                               1 MiB (16 instances)
L1i cache:                               1 MiB (16 instances)
L2 cache:                                32 MiB (16 instances)
L3 cache:                                80 MiB (1 instance)
NUMA node(s):                            1
NUMA node0 CPU(s):                       0-15
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Not affected
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; CSV2, BHB
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Not affected
Details

Comparing HEAD and fix_topk-global-prefix-early-exit
--------------------
Benchmark run_topk_tpch.json
--------------------
┏━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Query ┃                           HEAD ┃ fix_topk-global-prefix-early-exit ┃        Change ┃
┡━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ Q1    │    2.25 / 2.76 ±0.73 / 4.19 ms │       2.03 / 2.55 ±0.85 / 4.24 ms │ +1.08x faster │
│ Q2    │ 10.79 / 11.44 ±0.70 / 12.69 ms │    10.52 / 10.87 ±0.49 / 11.83 ms │     no change │
│ Q3    │ 31.99 / 32.86 ±0.62 / 33.75 ms │    31.91 / 32.59 ±0.79 / 34.10 ms │     no change │
│ Q4    │ 11.65 / 12.59 ±0.78 / 13.54 ms │    11.55 / 11.64 ±0.09 / 11.77 ms │ +1.08x faster │
│ Q5    │    9.25 / 9.37 ±0.08 / 9.47 ms │      9.71 / 9.83 ±0.11 / 10.01 ms │     no change │
│ Q6    │ 16.55 / 16.61 ±0.06 / 16.69 ms │    16.89 / 17.08 ±0.21 / 17.48 ms │     no change │
│ Q7    │ 36.49 / 37.04 ±0.86 / 38.74 ms │    37.38 / 37.97 ±0.50 / 38.55 ms │     no change │
│ Q8    │ 27.68 / 28.84 ±1.27 / 30.55 ms │    28.11 / 29.79 ±0.93 / 30.70 ms │     no change │
│ Q9    │ 35.35 / 36.57 ±1.61 / 39.70 ms │    35.09 / 35.45 ±0.33 / 36.07 ms │     no change │
│ Q10   │ 53.78 / 56.24 ±2.70 / 60.92 ms │    54.25 / 55.71 ±1.33 / 57.37 ms │     no change │
│ Q11   │    3.82 / 3.96 ±0.13 / 4.21 ms │       3.72 / 4.75 ±1.12 / 6.15 ms │  1.20x slower │
└───────┴────────────────────────────────┴───────────────────────────────────┴───────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓
┃ Benchmark Summary                                ┃          ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩
│ Total Time (HEAD)                                │ 248.29ms │
│ Total Time (fix_topk-global-prefix-early-exit)   │ 248.25ms │
│ Average Time (HEAD)                              │  22.57ms │
│ Average Time (fix_topk-global-prefix-early-exit) │  22.57ms │
│ Queries Faster                                   │        2 │
│ Queries Slower                                   │        1 │
│ Queries with No Change                           │        8 │
│ Queries with Failure                             │        0 │
└──────────────────────────────────────────────────┴──────────┘

Resource Usage

topk_tpch — base (merge-base)

Metric Value
Wall time 5.0s
Peak memory 590.2 MiB
Avg memory 98.4 MiB
CPU user 0.0s
CPU sys 0.0s
Peak spill 0 B

topk_tpch — branch

Metric Value
Wall time 5.0s
Peak memory 609.4 MiB
Avg memory 101.6 MiB
CPU user 0.0s
CPU sys 0.0s
Peak spill 0 B

File an issue against this benchmark runner

@geoffreyclaude
geoffreyclaude marked this pull request as ready for review June 17, 2026 16:35
@geoffreyclaude

Copy link
Copy Markdown
Contributor Author

@kosiew for info, since you reviewed the previous PR this one builds upon

@geoffreyclaude

Copy link
Copy Markdown
Contributor Author

I opened #23003 to add a dedicated TopK benchmark over sorted TPCH data

@kosiew kosiew 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.

@geoffreyclaude

Thanks for the update. The changes look good overall. I just have one small suggestion that could help reduce duplication in the test helpers.

]))
}

fn make_topk_filter() -> Arc<RwLock<TopKDynamicFilters>> {

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.

Small cleanup suggestion: make_topk_filter and make_shared_topk_filter currently build the same wrapper. It might be a little easier to maintain if the single-emitter helper delegates to the shared helper so any future fixture changes only need to be made in one place.

fn make_topk_filter() -> Arc<RwLock<TopKDynamicFilters>> {
    make_shared_topk_filter(1)
}

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.

Thanks, factorized in c05846f

@LiaCastaneda LiaCastaneda 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.

LGTM, just some minor comments

/// The dynamic filter expression is preserved, but wrapper state such as the
/// shared threshold and remaining emitter count is reset for the new
/// partitioning.
fn rebuild_filter_for_current_partitioning(&mut self) {

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.

nit: Is it harmless to reconstruct the filter from scratch? I guess this isn't an issue if it's only called at planning time, but if this function were called during execution we could end up discarding an active shared_threshold.

@geoffreyclaude geoffreyclaude Jun 18, 2026 •

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.

Yes, exactly: this is plan-construction / plan-rewrite state, not runtime execution state. Rebuilding the wrapper is intentional when the SortExec properties change, because the wrapper owns plan-derived state such as the shared threshold and expected emitter count. We should not carry an active runtime threshold across that kind of rebuild.

)
// `with_fetch` creates the TopK filter; preserving partitioning after
// that must rebuild it with one emitter per output partition.
.with_fetch(Some(2))

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.

something codex highlighted was that in with_fetch, if a dynamic filter already exists it's cloned and reused, so it keeps the emitter count it was built with, but not necessarily the current partition count. I'm not sure if it's actually possible that there the emitter count and partition count can differ, but it seems that every other place that rebuilds the sort calls rebuild_filter_for_current_partitioning, should we call it there as well?

@geoffreyclaude geoffreyclaude Jun 18, 2026 •

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.

Good catch. I changed with_fetch in 4c32f8a so an existing dynamic filter keeps its expression, but the TopKDynamicFilters wrapper is rebuilt against the current output partitioning instead of being cloned with possibly stale emitter state. I also added test_with_fetch_rebuilds_existing_topk_filter for the ordering where the filter exists before fetch is set.

geoffreyclaude and others added 5 commits June 19, 2026 16:58
Unlike the fully rejected batch path fixed by apache#22852, this path has already
passed the dynamic filter. `find_new_topk_items` can still produce zero heap
replacements, and that unchanged heap can still combine with the batch prefix
to prove that later rows cannot enter the TopK.
A partition-preserving `SortExec` creates one local `TopK` per output
partition. The shared dynamic filter should complete only after all of them
emit, so this commit tracks the remaining emitters and rebuilds that state when
`SortExec` partitioning changes.
The full encoded sort key still decides whether a shared threshold is more
selective. The encoded common-prefix row travels with that same heap row, so
lagging partitions can stop before their local heap is full when the shared
prefix proves no later rows can enter the TopK.
@geoffreyclaude
geoffreyclaude force-pushed the fix/topk-global-prefix-early-exit branch from 4c32f8a to 560cc5b Compare June 19, 2026 15:02
@geoffreyclaude

Copy link
Copy Markdown
Contributor Author

run benchmark topk_sorted_tpch

@adriangbot

Copy link
Copy Markdown

🤖 Benchmark running (GKE) | trigger
Instance: c4a-highmem-16 (12 vCPU / 65 GiB) | Linux bench-c4752736859-596-p2hmr 6.12.68+ #1 SMP Sat May 2 07:49:07 UTC 2026 aarch64 GNU/Linux

CPU Details (lscpu)
Architecture:                            aarch64
CPU op-mode(s):                          64-bit
Byte Order:                              Little Endian
CPU(s):                                  16
On-line CPU(s) list:                     0-15
Vendor ID:                               ARM
Model name:                              Neoverse-V2
Model:                                   1
Thread(s) per core:                      1
Core(s) per cluster:                     16
Socket(s):                               -
Cluster(s):                              1
Stepping:                                r0p1
BogoMIPS:                                2000.00
Flags:                                   fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache:                               1 MiB (16 instances)
L1i cache:                               1 MiB (16 instances)
L2 cache:                                32 MiB (16 instances)
L3 cache:                                80 MiB (1 instance)
NUMA node(s):                            1
NUMA node0 CPU(s):                       0-15
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Not affected
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; CSV2, BHB
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Not affected

Comparing fix/topk-global-prefix-early-exit (560cc5b) to 88fb94d (merge-base) diff using: topk_sorted_tpch
Results will be posted here when complete


File an issue against this benchmark runner

@adriangbot

Copy link
Copy Markdown

🤖 Benchmark completed (GKE) | trigger

Instance: c4a-highmem-16 (12 vCPU / 65 GiB)

CPU Details (lscpu)
Architecture:                            aarch64
CPU op-mode(s):                          64-bit
Byte Order:                              Little Endian
CPU(s):                                  16
On-line CPU(s) list:                     0-15
Vendor ID:                               ARM
Model name:                              Neoverse-V2
Model:                                   1
Thread(s) per core:                      1
Core(s) per cluster:                     16
Socket(s):                               -
Cluster(s):                              1
Stepping:                                r0p1
BogoMIPS:                                2000.00
Flags:                                   fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache:                               1 MiB (16 instances)
L1i cache:                               1 MiB (16 instances)
L2 cache:                                32 MiB (16 instances)
L3 cache:                                80 MiB (1 instance)
NUMA node(s):                            1
NUMA node0 CPU(s):                       0-15
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Not affected
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; CSV2, BHB
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Not affected
Details

Comparing HEAD and fix_topk-global-prefix-early-exit
--------------------
Benchmark run_topk_sorted_tpch.json
--------------------
┏━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Query ┃                           HEAD ┃ fix_topk-global-prefix-early-exit ┃        Change ┃
┡━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ Q1    │    2.12 / 2.98 ±0.76 / 4.21 ms │       2.14 / 2.83 ±0.83 / 4.31 ms │ +1.06x faster │
│ Q2    │    2.27 / 2.68 ±0.55 / 3.78 ms │       2.33 / 2.43 ±0.10 / 2.60 ms │ +1.11x faster │
│ Q3    │ 31.59 / 32.71 ±1.52 / 35.61 ms │    31.76 / 32.87 ±1.86 / 36.55 ms │     no change │
│ Q4    │ 11.43 / 11.58 ±0.13 / 11.81 ms │       2.56 / 2.67 ±0.14 / 2.95 ms │ +4.33x faster │
│ Q5    │   9.52 / 9.88 ±0.40 / 10.61 ms │     9.99 / 10.18 ±0.21 / 10.57 ms │     no change │
│ Q6    │ 16.67 / 16.81 ±0.09 / 16.93 ms │    17.15 / 17.33 ±0.23 / 17.77 ms │     no change │
│ Q7    │ 37.14 / 37.97 ±0.70 / 38.89 ms │    37.07 / 37.65 ±0.69 / 38.92 ms │     no change │
│ Q8    │ 3.35 / 11.89 ±10.47 / 27.14 ms │       3.33 / 4.55 ±2.34 / 9.24 ms │ +2.61x faster │
│ Q9    │    4.73 / 4.77 ±0.03 / 4.82 ms │       4.75 / 4.78 ±0.02 / 4.81 ms │     no change │
│ Q10   │ 6.85 / 32.12 ±20.97 / 53.02 ms │       6.88 / 6.90 ±0.01 / 6.91 ms │ +4.66x faster │
│ Q11   │    3.83 / 4.01 ±0.30 / 4.61 ms │       3.80 / 4.38 ±0.87 / 6.10 ms │  1.09x slower │
└───────┴────────────────────────────────┴───────────────────────────────────┴───────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓
┃ Benchmark Summary                                ┃          ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩
│ Total Time (HEAD)                                │ 167.41ms │
│ Total Time (fix_topk-global-prefix-early-exit)   │ 126.56ms │
│ Average Time (HEAD)                              │  15.22ms │
│ Average Time (fix_topk-global-prefix-early-exit) │  11.51ms │
│ Queries Faster                                   │        5 │
│ Queries Slower                                   │        1 │
│ Queries with No Change                           │        5 │
│ Queries with Failure                             │        0 │
└──────────────────────────────────────────────────┴──────────┘

Resource Usage

topk_sorted_tpch — base (merge-base)

Metric Value
Wall time 5.0s
Peak memory 640.4 MiB
Avg memory 106.8 MiB
CPU user 0.0s
CPU sys 0.0s
Peak spill 0 B

topk_sorted_tpch — branch

Metric Value
Wall time 5.0s
Peak memory 224.0 KiB
Avg memory 37.3 KiB
CPU user 0.0s
CPU sys 0.0s
Peak spill 0 B

File an issue against this benchmark runner

@kosiew

kosiew commented Jun 20, 2026

Copy link
Copy Markdown
Contributor

Merging with main to get #23053

@kosiew
kosiew added this pull request to the merge queue Jun 22, 2026
Merged via the queue into apache:main with commit ba67bb4 Jun 22, 2026
38 checks passed
gh-worker-dd-mergequeue-cf854d Bot pushed a commit to DataDog/datafusion that referenced this pull request Sep 4, 2026
## Which issue does this PR close?

Closes apache#22874.

Follow-up to [fix(topk): call attempt_early_completion when filter
rejects entire batch].

## Rationale for this change

[TopK dynamic filter pushdown attempt 2] lets `SortExec` tighten
scan-side predicates while a TopK heap finds better rows. Once the
current TopK threshold is known, scans can skip data that cannot enter
the final `ORDER BY ... LIMIT` result.

That works well for a single partition. Partitioned `SortExec` has one
extra case to handle:

- each output partition has its own local `TopK` heap
- those local heaps share one `TopKDynamicFilters` instance
- one partition can tighten the shared filter before another partition
has enough rows to fill its local heap

[fix(topk): call attempt_early_completion when filter rejects entire
batch] fixed the local case where a heap already has a max row and the
dynamic filter rejects a whole batch.

This PR fixes the remaining shared-filter case. A lagging partition can
now use the shared prefix threshold to stop early even when its local
heap is still empty. If there is no shared threshold yet, it falls back
to the existing local heap prefix check.

The shared prefix check is not treating another partition's threshold as
this partition's local heap boundary. It uses the same threshold that
already drives the shared dynamic filter. Once a partition's ordered
input has moved past that shared prefix threshold, later batches from
that partition cannot add rows that survive the shared filter. The local
heap still emits the candidates it has already kept; this only stops
pulling input that can no longer add candidates.

Single-partition behavior is unchanged.

## How the TopK optimizations fit together

There are two existing optimizations involved here:

- Dynamic filter pushdown: once a `TopK` heap has K rows, its worst kept
row becomes a threshold. That threshold tightens a scan-side filter so
later data that cannot enter the final `ORDER BY ... LIMIT` result can
be skipped.
- Prefix early exit: when the input is ordered by a prefix of the
requested sort, `TopK` can stop pulling once the last row in a batch is
past a known TopK boundary on that shared prefix.

For a partition-preserving `SortExec`, those optimizations meet in one
shared place. Each output partition has a local `TopK` heap, but all of
those local heaps publish into one shared dynamic filter. A partition
that fills first can tighten the shared filter for everyone else.
Lagging partitions then need to use that same shared prefix threshold to
stop pulling once their ordered input has moved past it.

This PR makes that composition explicit: the shared filter stays alive
until every local `TopK` has emitted, and the shared threshold carries
its common-prefix row so lagging partitions can apply the same prefix
early-exit check.

## What changes are included in this PR?

- Check early completion when a batch passes the dynamic filter but
produces zero heap replacements.
- Track local TopK emitters so a shared filter completes only after the
last emitter has produced output.
- Store the shared threshold and its common-prefix row together in
`TopKDynamicFilters`.
- Check the shared prefix in `attempt_early_completion` before falling
back to the local heap prefix.
- Add focused TopK and `SortExec` tests for the local zero-replacement
path, early completion before local heap fill, equal-prefix
non-completion, DESC/null prefix ordering, and shared-filter completion.

## How is this split for review?

The commits are ordered so each one has a narrow job:

1. `Check TopK early completion after zero-replacement batches`
Handles the local case where a batch passes the dynamic filter, produces
zero heap replacements, but still proves later rows cannot enter the
TopK.

2. `Complete shared TopK filters after all emitters`
Keeps a shared dynamic filter watchable until every local TopK emitter
has emitted. This includes the `SortExec` wiring for preserved
partitioning.

3. `Use shared TopK prefix thresholds for early exit`
Carries the shared threshold's common-prefix row so lagging partitions
can stop before their local heap is full.

## Are these changes tested?

Correctness is covered by targeted tests for:

- the local zero-replacement early-completion path
- early completion before local heap fill from a shared prefix threshold
- equal-prefix non-completion
- DESC and NULLS LAST prefix row ordering
- shared-filter completion after all TopK emitters, including preserved
`SortExec` partitioning

Relevant background:

- [perf: Add TopK benchmarks as variation over the `sort_tpch`
benchmarks] added the benchmark setup used here.
- [perf: Introduce sort prefix computation for early TopK exit
optimization on partially sorted input (10x speedup on top10 bench)]
added common-prefix TopK early termination.
- [TopK dynamic filter pushdown attempt 2] added scan-side dynamic
filter pushdown and exposed the shared-filter / local-heap interaction.
- [fix(topk): call attempt_early_completion when filter rejects entire
batch] fixed the local all-filtered-batch case.
- This PR fixes the remaining partitioned shared-filter case.

Benchmark command:

```bash
dfbench sort-tpch --sorted --limit 10 --iterations 5 \
  --path /tmp/df-topk-bench-data/tpch_sf1 \
  -o /tmp/topk-shared-prefix-followup.json
```

Both sides were rebuilt with fresh isolated `release-nonlto` target
directories before running the benchmark. This is not the regular
`topk_tpch` script: this case needs `--sorted --limit 10` to exercise
the prefix early-exit path.

Clean rerun against the PR base commit `7bb6e152b`. Times are
milliseconds, using the average of 5 iterations for each row.

| scope | PR base | this PR | change |
|---|---:|---:|---:|
| all sort-tpch queries | 760.59 | 348.84 | -54.1% |
| Q8 | 49.70 | 7.66 | -84.6% |
| Q9 | 92.13 | 10.27 | -88.8% |
| Q10 | 89.66 | 14.35 | -84.0% |

The `DataSourceExec` counters show the less noisy part of the result.
Q8/Q9/Q10 now emit only the first batch from each partition instead of
continuing to drain millions of rows.

| query | PR base `DataSourceExec output_rows` | this PR `DataSourceExec
output_rows` | PR base `bytes_scanned` | this PR `bytes_scanned` |
|---|---:|---:|---:|---:|
| Q8 | 3.66M | 81.92K | 56.81M | 15.79M |
| Q9 | 3.66M | 81.92K | 75.19M | 20.89M |
| Q10 | 3.10M | 81.92K | 110.9M | 34.69M |

## Are there any user-facing changes?

No. This is an internal physical execution optimization fix.

[perf: Add TopK benchmarks as variation over the `sort_tpch`
benchmarks]: apache#15560
[perf: Introduce sort prefix computation for early TopK exit
optimization on partially sorted input (10x speedup on top10 bench)]:
apache#15563
[TopK dynamic filter pushdown attempt 2]:
apache#15770
[fix(topk): call attempt_early_completion when filter rejects entire
batch]: apache#22852

---------

Co-authored-by: kosiew <kosiew@gmail.com>
(cherry picked from commit ba67bb4)
gh-worker-dd-mergequeue-cf854d Bot pushed a commit to DataDog/datafusion that referenced this pull request Sep 4, 2026
gh-worker-dd-mergequeue-cf854d Bot added a commit to DataDog/datafusion that referenced this pull request Sep 4, 2026
…ck/apache-pr-22991-branch54-20260904

[branch-54] Cherry-pick apache#22991

Co-authored-by: geoffreyclaude <geoffrey.claude@datadoghq.com>
mohitgurav20 pushed a commit to mohitgurav20/datafusion that referenced this pull request Oct 24, 2026
…che#26455)

## Which issue does this PR close?

- Closes apache#26454.

## Rationale for this change

`ORDER BY ... LIMIT` over a `UNION ALL` returns wrong rows when a
leading sort key is a constant with a different value in each arm:

```sql
SELECT 5 AS c0, id AS c1 FROM t
UNION ALL
SELECT 1 AS c0, id AS c1 FROM t WHERE g = 0
ORDER BY c0, c1 LIMIT 2
-- returns (1, 0), (5, 0); expected (1, 0), (1, 1)
```

The optimizer plans a partitioned TopK under a `SortPreservingMergeExec`
and drops `c0` from the TopK key, because `c0` is constant within each
partition. All partitions of one `SortExec` share a single TopK dynamic
filter threshold. The `c0 = 5` partition publishes a threshold on `c1`
alone, and the `c0 = 1` partition then discards a row it needs. The
issue has the full walk-through.

## What changes are included in this PR?

`add_sort_above` and `add_sort_above_with_distribution` in
`datafusion/physical-optimizer/src/utils.rs` both remove sort keys that
are constant in the input. They now share one helper,
`ordering_without_constants(plan, sort_requirements,
keep_heterogeneous)`, which still removes them, except that a constant
whose value differs between partitions
(`AcrossPartitions::Heterogeneous`) is kept when `keep_heterogeneous` is
set. Each caller sets it when an operator compares rows from different
partitions on the new ordering:

- **Partitioned TopK** (both helpers): the new `SortExec` has a `fetch`
and its input has more than one partition. The shared threshold then
compares the full key, which is comparable across partitions.
- **Merge above the sort** (`add_sort_above_with_distribution` only):
when the parent requires a single partition, the helper wraps the sort
in a `SortPreservingMergeExec` built from the same ordering. That merge
compares rows across partitions even without a fetch, and since it
clears per-partition constants from its output, a shortened key would no
longer guarantee the required ordering. The constant is now kept there
too.

The helper preserves these existing behaviors:

- `Uniform` constants are still removed.
- A partitioned sort without a `fetch` and without an added merge still
removes `Heterogeneous` constants. Each partition sorts on its own, and
any merge above it sorts on the full key.
- With a single input partition nothing is compared across partitions.

With the change, the plan for the query above is:

```
SortPreservingMergeExec: [c0@0 ASC NULLS LAST, c1@1 ASC NULLS LAST], fetch=2
  SortExec: TopK(fetch=2), expr=[c0@0 ASC NULLS LAST, c1@1 ASC NULLS LAST], preserve_partitioning=[true], sort_prefix=[c0@0 ASC NULLS LAST]
```

The longer key gives the TopK a `sort_prefix` of `[c0]`, because each
partition's input is trivially ordered by its constant `c0`.

Why the fix lives here: `EXPLAIN VERBOSE` shows the shortened `SortExec:
TopK(fetch=2), expr=[c1@1 ...]` already in the plan right after
`EnsureRequirements`, with the fetch attached. No later rule changes it.
A temporary print in the two helpers showed `parallelize_sorts` building
it through `add_sort_above_with_check` → `add_sort_above`, and
`pushdown_sorts` building it again through
`add_sort_above_with_distribution`. Both calls had `fetch = Some(2)`,
two input partitions, and `c0` reported as `Heterogeneous`. So both
helpers need the change, and fixing only one would leave the bug.

apache#23528 fixed the same `Uniform`/`Heterogeneous` confusion in
`EquivalenceGroup`'s sort normalization. This PR applies it to the TopK
key and to the merge that these two helpers build.

Open PRs apache#25408 and apache#25065 change the callers of these helpers
(`enforce_sorting/sort_pushdown.rs` and `enforce_sorting/mod.rs`) but
not `utils.rs`. This PR should not conflict with them textually, but
whichever merges second should re-run `topk.slt`.

Two separate validation prerequisites are included: three narrowly
scoped `mutable_key_type` Clippy annotations matching existing optimizer
conventions, and lockfile patch updates for `rustls-webpki` (0.103.16)
and `webpki-root-certs` (1.0.10). The exact all-features Clippy command
failed on the unchanged base, and the security audit exposed
RUSTSEC-2026-0382 / RUSTSEC-2026-0383 in the existing lockfile. Neither
prerequisite changes the TopK implementation.

## What is the testing strategy for this PR?

- **`topk.slt`**, on `VALUES` tables, each with an `EXPLAIN` that pins
the TopK key and a result query:
  - the query above, and the same query without the filter;
  - `ORDER BY c0 DESC, c1`;
- `ORDER BY c1, c0`, where the constant is not the leading key and only
breaks ties on `c1`;
- `LIMIT ... OFFSET`, where the TopK needs three rows from the `c0 = 1`
arm;
- a three-arm union, so two arms can publish a threshold that cuts the
third;
- controls: `LIMIT 3`, a `Uniform` constant (end-to-end control), and no
`fetch` (sort pushed into the arms).

Whether a result query fails without the fix depends on which arm
publishes its threshold first (the `DESC` query never failed in 8 runs),
so the `EXPLAIN`s are what catch a regression reliably.
- **Unit tests in `utils.rs`**, one per branch: a partitioned TopK keeps
a `Heterogeneous` constant; the added merge keeps it with and without a
fetch; a partitioned sort without a fetch or merge drops it; a `Uniform`
constant is always dropped.
- **Unit tests in `topk/mod.rs`**: two TopKs share one
`TopKDynamicFilters`, with key `[c0, c1]` and `sort_prefix` `[c0]`. Both
feeding orders return the right top 2. A second test feeds several
batches per partition and checks the early-exit state: the better
partition keeps reading until it has `k` rows, and only then does the
worse one finish.

Ablation, with all tests in place:

- Dropping `Heterogeneous` constants again in
`ordering_without_constants` fails both `keep` unit tests in `utils.rs`
and 10 `topk.slt` assertions, including every new `EXPLAIN`.
- Removing only the merge part (`|| add_merge`) fails
`merge_keeps_a_constant_that_differs_between_partitions` on its `fetch =
None` case. No `.slt` test fails, see below.

No existing `.slt` plan changes.

Local validation on the final tree: formatting, `cargo clippy
--all-targets --all-features -- -D warnings`, and `uv run
./dev/rust_lint.sh` pass. The required extended workspace test command
passes with 12,733 Rust tests (8 ignored) and all 535 SQL files.
Independent mutations also confirm the new TopK tests fail when the
leading filter key is omitted, and the multi-batch test fails when the
shared-prefix comparison is reversed. Benchmark datasets were
unavailable; no performance claim is made.

## Are there any user-facing changes?

The wrong results are fixed. Plans that hit this case now show the
constant in the TopK `expr`. There are no API changes.

## Unknowns / assumptions

- **Shared prefix threshold.** The new key gives the TopK a
`sort_prefix` of `[c0]`. That enables the shared-prefix early exit from
apache#22991 for these plans. It stays correct because it compares on the full
key prefix, which now includes `c0`, and a threshold is only published
once the publishing partition has `k` rows. The multi-batch TopK unit
test exercises this directly; no `.slt` test forces a particular
partition to finish first.
- **No tested SQL reached the no-fetch merge branch.** I instrumented
`add_sort_above_with_distribution` for `fetch = None` with a
single-partition requirement and more than one input partition. It
recorded no hits across the whole sqllogictest suite, for candidate
queries (un-partitioned windows over union, filter, join, aggregate and
distinct), or with `optimizer.repartition_sorts = false` /
`optimizer.enable_sort_pushdown = false`. Distribution enforcement and
`parallelize_sorts` build the merge with the full ordering before
`pushdown_sorts` runs. That is an observation, not a proof, so that
branch is covered by the unit test only.
- **Fetch attached after key simplification (accepted risk;
follow-up).** The helper decides whether to retain heterogeneous
constants when it builds the sort. A partitioned sort without a fetch or
helper-added merge may omit `c0`; if `LimitPushdown` or
`LimitPushPastWindows` subsequently calls `SortExec::with_fetch`, that
method enables shared TopK filtering without restoring `c0`. This PR
does not guard that transition. No SQL query tried so far reaches the
entire sequence under the current rule order: the tested shapes already
have a fetch at construction, push sorts into individual union arms,
retain the full key below a merge, hash-repartition first, or eliminate
the inner sort. That evidence does not establish unreachability. The
proposed follow-up is a hand-built plan regression plus an executor
guard that isolates or disables unsafe shared filtering, including
dynamic-filter pushdown. The reproduction is proposed, not yet executed.
`preserve_fetch_above_ordered_plan` remains an additional path to audit.
- **Alternative fix.** The other option is to stop sharing the TopK
threshold across partitions when the input has a `Heterogeneous`
constant. That would be a larger change inside TopK and would give up
the shared filter. I went with the smaller planner change, but I am
happy to switch if maintainers prefer it.

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Jay Zhan <jayzhan211@gmail.com>
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 v55.0.0

Projects

None yet

Development

Successfully merging this pull request may close these issues.

TopK early-exit doesn't fire when a partition's local heap is empty but the shared dynamic filter is already tight

5 participants