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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
