You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "cning112 (via GitHub)" <gi...@apache.org> on 2026/10/17 21:09:27 UTC

[PR] fix: keep per-partition constant sort keys in a partitioned TopK [datafusion]

cning112 opened a new pull request, #26455:
URL: https://github.com/apache/datafusion/pull/26455

   ## Which issue does this PR close?
   
   - Closes #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, `sort_ordering_above`, which still removes them, with one exception. When the new `SortExec` has a `fetch` and its input has more than one partition, a constant that differs between partitions (`AcrossPartitions::Heterogeneous`) stays in the key. The shared threshold then compares the full key, which is comparable across partitions.
   
   Nothing else changes:
   
   - `Uniform` constants are still removed.
   - Sorts without a `fetch` still remove `Heterogeneous` constants. A full per-partition sort has no shared threshold, and the merge above it sorts on the full key.
   - With a single input partition the sort is not partitioned, so nothing is shared.
   
   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.
   
   #23528 fixed the same `Uniform`/`Heterogeneous` confusion in `EquivalenceGroup`'s sort normalization. This PR deals only with the fetch case in these two helpers.
   
   Open PRs #25408 and #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`.
   
   ## What is the testing strategy for this PR?
   
   New cases at the end of `datafusion/sqllogictest/test_files/topk.slt`, on a `VALUES` table:
   
   - the query above with an `EXPLAIN` showing the TopK key, at `target_partitions = 1` and `4`;
   - the same query without the filter;
   - `LIMIT 3` and `ORDER BY c0 DESC, c1` as controls.
   
   On `main` without the fix, four of them fail: the `EXPLAIN`, `LIMIT 2` at both partition counts, and the no-filter query. The `LIMIT 3` and `DESC` controls pass both before and after.
   
   No existing `.slt` plan changes. `--complete` over the whole suite rewrote nothing but trailing newlines in three unrelated files, and I reverted those.
   
   ## 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 #22991 for these plans. My reasoning is that the exit stays correct because it compares on the full key prefix, which now includes `c0`. The `DESC` case exercises the opposite partition order, but I have not written a test that forces either partition to finish first.
   - **Merge above the sort.** `add_sort_above_with_distribution` wraps the sort in a `SortPreservingMergeExec` with the same shortened ordering when the parent requires a single partition. Without a fetch, that merge could also lose a `Heterogeneous` key. I could not build a query that reaches it, and this PR does not change it.
   - **Other places that attach a fetch.** I assume no other rule attaches a fetch to a partitioned `SortExec` whose key was already shortened. In the queries I tried, the fetch is already set when `EnsureRequirements` builds the sort: plain `LIMIT`, `LIMIT ... OFFSET`, and a filter over the union. `LimitPushdown` left those plans unchanged. I have not audited every caller of `SortExec::with_fetch`.
   - **The no-filter variant.** Its result on `main` depends on how the partitions run: wrong in the sqllogictest, right over a single-batch `MemTable` in my runs. I did not dig into why.
   - **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.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] fix: keep per-partition constant sort keys in a partitioned TopK [datafusion]

Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #26455:
URL: https://github.com/apache/datafusion/pull/26455#issuecomment-6359802011

   <!-- semver-check-comment -->
   Thank you for opening this pull request!
   
   Reviewer note: [cargo-semver-checks](https://github.com/obi1kenobi/cargo-semver-checks) reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch).
   
   <details>
   <summary>Details</summary>
   
   ```
        Cloning apache/main
       Building datafusion-physical-optimizer v55.1.0 (current)
          Built [  46.999s] (current)
        Parsing datafusion-physical-optimizer v55.1.0 (current)
         Parsed [   0.021s] (current)
       Building datafusion-physical-optimizer v55.1.0 (baseline)
          Built [  45.433s] (baseline)
        Parsing datafusion-physical-optimizer v55.1.0 (baseline)
         Parsed [   0.022s] (baseline)
       Checking datafusion-physical-optimizer v55.1.0 -> v55.1.0 (no change; assume patch)
        Checked [   0.118s] 229 checks: 229 pass, 31 skip
        Summary no semver update required
       Finished [  93.796s] datafusion-physical-optimizer
       Building datafusion-physical-plan v55.1.0 (current)
          Built [  43.769s] (current)
        Parsing datafusion-physical-plan v55.1.0 (current)
         Parsed [   0.231s] (current)
       Building datafusion-physical-plan v55.1.0 (baseline)
          Built [  43.823s] (baseline)
        Parsing datafusion-physical-plan v55.1.0 (baseline)
         Parsed [   0.233s] (baseline)
       Checking datafusion-physical-plan v55.1.0 -> v55.1.0 (no change; assume patch)
        Checked [   0.707s] 229 checks: 228 pass, 1 fail, 0 warn, 31 skip
   
   --- failure inherent_method_missing: pub method removed or renamed ---
   
   Description:
   A publicly-visible method or associated fn is no longer available under its prior name. It may have been renamed or removed entirely.
           ref: https://doc.rust-lang.org/cargo/reference/semver.html#item-remove
          impl: https://github.com/obi1kenobi/cargo-semver-checks/tree/v0.51.0/src/lints/inherent_method_missing.ron
   
   Failed in:
     AggregateExec::input_keeps_groups_local, previously in file /home/runner/work/datafusion/datafusion/target/semver-checks/git-apache_main/27c76af68376a4e491fb57fc721b6aece1ab12a3/datafusion/physical-plan/src/aggregates/mod.rs:1663
   
        Summary semver requires new major version: 1 major and 0 minor checks failed
       Finished [  90.641s] datafusion-physical-plan
   ```
   
   </details>


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] fix: keep per-partition constant sort keys in a partitioned TopK [datafusion]

Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 merged PR #26455:
URL: https://github.com/apache/datafusion/pull/26455


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] fix: keep per-partition constant sort keys in a partitioned TopK [datafusion]

Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #26455:
URL: https://github.com/apache/datafusion/pull/26455#discussion_r4306486284


##########
datafusion/sqllogictest/test_files/topk.slt:
##########
@@ -396,3 +396,92 @@ set datafusion.explain.physical_plan_only = false
 
 statement ok
 set datafusion.execution.batch_size = 8192
+
+# A partitioned TopK must not drop a sort key that is constant within each
+# partition but differs between partitions. All partitions of one `SortExec`
+# share a single TopK threshold, so a threshold on `c1` alone, published by the
+# `c0 = 5` arm, would discard rows that the `c0 = 1` arm needs.
+statement ok
+CREATE TABLE heterogeneous_const(id BIGINT, g BIGINT) AS VALUES (0, 0), (1, 0), (2, 1), (3, 1), (4, 1);
+
+statement ok
+set datafusion.explain.physical_plan_only = true
+
+statement ok
+set datafusion.execution.target_partitions = 1;
+
+query TT
+EXPLAIN SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
+UNION ALL
+SELECT 1 AS c0, id AS c1 FROM heterogeneous_const WHERE g = 0
+ORDER BY c0, c1 LIMIT 2;
+----
+physical_plan
+01)SortPreservingMergeExec: [c0@0 ASC NULLS LAST, c1@1 ASC NULLS LAST], fetch=2
+02)--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]
+03)----UnionExec
+04)------ProjectionExec: expr=[5 as c0, id@0 as c1]
+05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+06)------ProjectionExec: expr=[1 as c0, id@0 as c1]
+07)--------FilterExec: g@1 = 0, projection=[id@0]
+08)----------DataSourceExec: partitions=1, partition_sizes=[1]
+
+query II
+SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
+UNION ALL
+SELECT 1 AS c0, id AS c1 FROM heterogeneous_const WHERE g = 0
+ORDER BY c0, c1 LIMIT 2;
+----
+1 0
+1 1
+
+# Control: with LIMIT 3 the `c0 = 1` arm never needs a row the threshold cuts
+query II
+SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
+UNION ALL
+SELECT 1 AS c0, id AS c1 FROM heterogeneous_const WHERE g = 0
+ORDER BY c0, c1 LIMIT 3;
+----
+1 0
+1 1
+5 0
+
+# Same bug without the filter: both arms hold the same `c1` values
+query II
+SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
+UNION ALL
+SELECT 1 AS c0, id AS c1 FROM heterogeneous_const
+ORDER BY c0, c1 LIMIT 2;
+----
+1 0
+1 1
+
+# Descending on the constant key: the `c0 = 5` arm comes first
+query II
+SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
+UNION ALL
+SELECT 1 AS c0, id AS c1 FROM heterogeneous_const WHERE g = 0
+ORDER BY c0 DESC, c1 LIMIT 2;
+----
+5 0
+5 1
+
+# Same query with several partitions. The SLT runner uses 4, so this also
+# restores its setting.
+statement ok
+set datafusion.execution.target_partitions = 4;
+
+query II

Review Comment:
   nit: The fix also covers a heterogeneous constant that isn't the leading key, where it breaks ties. With the old helper this returned `(5, 0), (5, 1)` in 5 of my 6 runs. It might be worth adding (verified on this branch; the second arm uses `id + 1` so the `c1` values tie):
   
   ```sql
   query TT
   EXPLAIN SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
   UNION ALL
   SELECT 1 AS c0, id + 1 AS c1 FROM heterogeneous_const WHERE g = 0
   ORDER BY c1, c0 LIMIT 2;
   ----
   physical_plan
   01)SortPreservingMergeExec: [c1@1 ASC NULLS LAST, c0@0 ASC NULLS LAST], fetch=2
   02)--SortExec: TopK(fetch=2), expr=[c1@1 ASC NULLS LAST, c0@0 ASC NULLS LAST], preserve_partitioning=[true]
   03)----UnionExec
   04)------ProjectionExec: expr=[5 as c0, id@0 as c1]
   05)--------DataSourceExec: partitions=1, partition_sizes=[1]
   06)------ProjectionExec: expr=[1 as c0, id@0 + 1 as c1]
   07)--------FilterExec: g@1 = 0, projection=[id@0]
   08)----------DataSourceExec: partitions=1, partition_sizes=[1]
   
   query II
   SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
   UNION ALL
   SELECT 1 AS c0, id + 1 AS c1 FROM heterogeneous_const WHERE g = 0
   ORDER BY c1, c0 LIMIT 2;
   ----
   5 0
   1 1
   ```
   



##########
datafusion/sqllogictest/test_files/topk.slt:
##########
@@ -396,3 +396,92 @@ set datafusion.explain.physical_plan_only = false
 
 statement ok
 set datafusion.execution.batch_size = 8192
+
+# A partitioned TopK must not drop a sort key that is constant within each
+# partition but differs between partitions. All partitions of one `SortExec`
+# share a single TopK threshold, so a threshold on `c1` alone, published by the
+# `c0 = 5` arm, would discard rows that the `c0 = 1` arm needs.
+statement ok
+CREATE TABLE heterogeneous_const(id BIGINT, g BIGINT) AS VALUES (0, 0), (1, 0), (2, 1), (3, 1), (4, 1);
+
+statement ok
+set datafusion.explain.physical_plan_only = true
+
+statement ok
+set datafusion.execution.target_partitions = 1;
+
+query TT
+EXPLAIN SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
+UNION ALL
+SELECT 1 AS c0, id AS c1 FROM heterogeneous_const WHERE g = 0
+ORDER BY c0, c1 LIMIT 2;
+----
+physical_plan
+01)SortPreservingMergeExec: [c0@0 ASC NULLS LAST, c1@1 ASC NULLS LAST], fetch=2
+02)--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]
+03)----UnionExec
+04)------ProjectionExec: expr=[5 as c0, id@0 as c1]
+05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+06)------ProjectionExec: expr=[1 as c0, id@0 as c1]
+07)--------FilterExec: g@1 = 0, projection=[id@0]
+08)----------DataSourceExec: partitions=1, partition_sizes=[1]
+
+query II
+SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
+UNION ALL
+SELECT 1 AS c0, id AS c1 FROM heterogeneous_const WHERE g = 0
+ORDER BY c0, c1 LIMIT 2;
+----
+1 0
+1 1
+
+# Control: with LIMIT 3 the `c0 = 1` arm never needs a row the threshold cuts
+query II
+SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
+UNION ALL
+SELECT 1 AS c0, id AS c1 FROM heterogeneous_const WHERE g = 0
+ORDER BY c0, c1 LIMIT 3;
+----
+1 0
+1 1
+5 0
+
+# Same bug without the filter: both arms hold the same `c1` values
+query II
+SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
+UNION ALL
+SELECT 1 AS c0, id AS c1 FROM heterogeneous_const
+ORDER BY c0, c1 LIMIT 2;
+----
+1 0
+1 1
+
+# Descending on the constant key: the `c0 = 5` arm comes first
+query II

Review Comment:
   nit: I ran this file 6 times with the old `utils.rs` and this DESC query passed every time, since whether it fails depends on which arm publishes its threshold first. The EXPLAIN above is what reliably catches the regression. Could you add an EXPLAIN for the DESC case too, or drop this query?
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] fix: keep per-partition constant sort keys in a partitioned TopK [datafusion]

Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on PR #26455:
URL: https://github.com/apache/datafusion/pull/26455#issuecomment-6380561374

   conflicts!


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] fix: keep per-partition constant sort keys in a partitioned TopK [datafusion]

Posted by "cning112 (via GitHub)" <gi...@apache.org>.
cning112 commented on code in PR #26455:
URL: https://github.com/apache/datafusion/pull/26455#discussion_r4343005906


##########
datafusion/sqllogictest/test_files/topk.slt:
##########
@@ -396,3 +396,92 @@ set datafusion.explain.physical_plan_only = false
 
 statement ok
 set datafusion.execution.batch_size = 8192
+
+# A partitioned TopK must not drop a sort key that is constant within each
+# partition but differs between partitions. All partitions of one `SortExec`
+# share a single TopK threshold, so a threshold on `c1` alone, published by the
+# `c0 = 5` arm, would discard rows that the `c0 = 1` arm needs.
+statement ok
+CREATE TABLE heterogeneous_const(id BIGINT, g BIGINT) AS VALUES (0, 0), (1, 0), (2, 1), (3, 1), (4, 1);
+
+statement ok
+set datafusion.explain.physical_plan_only = true
+
+statement ok
+set datafusion.execution.target_partitions = 1;
+
+query TT
+EXPLAIN SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
+UNION ALL
+SELECT 1 AS c0, id AS c1 FROM heterogeneous_const WHERE g = 0
+ORDER BY c0, c1 LIMIT 2;
+----
+physical_plan
+01)SortPreservingMergeExec: [c0@0 ASC NULLS LAST, c1@1 ASC NULLS LAST], fetch=2
+02)--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]
+03)----UnionExec
+04)------ProjectionExec: expr=[5 as c0, id@0 as c1]
+05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+06)------ProjectionExec: expr=[1 as c0, id@0 as c1]
+07)--------FilterExec: g@1 = 0, projection=[id@0]
+08)----------DataSourceExec: partitions=1, partition_sizes=[1]
+
+query II
+SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
+UNION ALL
+SELECT 1 AS c0, id AS c1 FROM heterogeneous_const WHERE g = 0
+ORDER BY c0, c1 LIMIT 2;
+----
+1 0
+1 1
+
+# Control: with LIMIT 3 the `c0 = 1` arm never needs a row the threshold cuts
+query II
+SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
+UNION ALL
+SELECT 1 AS c0, id AS c1 FROM heterogeneous_const WHERE g = 0
+ORDER BY c0, c1 LIMIT 3;
+----
+1 0
+1 1
+5 0
+
+# Same bug without the filter: both arms hold the same `c1` values
+query II
+SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
+UNION ALL
+SELECT 1 AS c0, id AS c1 FROM heterogeneous_const
+ORDER BY c0, c1 LIMIT 2;
+----
+1 0
+1 1
+
+# Descending on the constant key: the `c0 = 5` arm comes first
+query II

Review Comment:
   Thanks, good catch. Added an `EXPLAIN` for the `DESC` case in aa9d17c49 and kept the result query. The `EXPLAIN` pins `expr=[c0@0 DESC, c1@1 ASC NULLS LAST]` with `sort_prefix=[c0@0 DESC]`, so it fails without the fix regardless of which arm publishes its threshold first.



##########
datafusion/sqllogictest/test_files/topk.slt:
##########
@@ -396,3 +396,92 @@ set datafusion.explain.physical_plan_only = false
 
 statement ok
 set datafusion.execution.batch_size = 8192
+
+# A partitioned TopK must not drop a sort key that is constant within each
+# partition but differs between partitions. All partitions of one `SortExec`
+# share a single TopK threshold, so a threshold on `c1` alone, published by the
+# `c0 = 5` arm, would discard rows that the `c0 = 1` arm needs.
+statement ok
+CREATE TABLE heterogeneous_const(id BIGINT, g BIGINT) AS VALUES (0, 0), (1, 0), (2, 1), (3, 1), (4, 1);
+
+statement ok
+set datafusion.explain.physical_plan_only = true
+
+statement ok
+set datafusion.execution.target_partitions = 1;
+
+query TT
+EXPLAIN SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
+UNION ALL
+SELECT 1 AS c0, id AS c1 FROM heterogeneous_const WHERE g = 0
+ORDER BY c0, c1 LIMIT 2;
+----
+physical_plan
+01)SortPreservingMergeExec: [c0@0 ASC NULLS LAST, c1@1 ASC NULLS LAST], fetch=2
+02)--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]
+03)----UnionExec
+04)------ProjectionExec: expr=[5 as c0, id@0 as c1]
+05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+06)------ProjectionExec: expr=[1 as c0, id@0 as c1]
+07)--------FilterExec: g@1 = 0, projection=[id@0]
+08)----------DataSourceExec: partitions=1, partition_sizes=[1]
+
+query II
+SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
+UNION ALL
+SELECT 1 AS c0, id AS c1 FROM heterogeneous_const WHERE g = 0
+ORDER BY c0, c1 LIMIT 2;
+----
+1 0
+1 1
+
+# Control: with LIMIT 3 the `c0 = 1` arm never needs a row the threshold cuts
+query II
+SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
+UNION ALL
+SELECT 1 AS c0, id AS c1 FROM heterogeneous_const WHERE g = 0
+ORDER BY c0, c1 LIMIT 3;
+----
+1 0
+1 1
+5 0
+
+# Same bug without the filter: both arms hold the same `c1` values
+query II
+SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
+UNION ALL
+SELECT 1 AS c0, id AS c1 FROM heterogeneous_const
+ORDER BY c0, c1 LIMIT 2;
+----
+1 0
+1 1
+
+# Descending on the constant key: the `c0 = 5` arm comes first
+query II
+SELECT 5 AS c0, id AS c1 FROM heterogeneous_const
+UNION ALL
+SELECT 1 AS c0, id AS c1 FROM heterogeneous_const WHERE g = 0
+ORDER BY c0 DESC, c1 LIMIT 2;
+----
+5 0
+5 1
+
+# Same query with several partitions. The SLT runner uses 4, so this also
+# restores its setting.
+statement ok
+set datafusion.execution.target_partitions = 4;
+
+query II

Review Comment:
   Added verbatim in aa9d17c49, thanks. I also added `LIMIT ... OFFSET`, three-arm, uniform-constant and no-fetch cases to `topk.slt`. Separately, 2d13b51ff keeps the constant in the `SortPreservingMergeExec` that `add_sort_above_with_distribution` adds, with unit tests for each branch of the helper. 272e15721 adds TopK-level tests for the shared threshold, including several batches per partition. Since these change behaviour beyond what you approved, would you mind taking another look?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org