diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 283d22767f9b..1211f969275e 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -1220,32 +1220,37 @@ fn range_join_key_alignment_sort_merge_join() -> Result<()> { } /// Four unbounded partitions declared as `Range([a, b])` whose rows are -/// ordered by `ordered_by` within each partition: partition membership and -/// row order are independent. -fn unbounded_range_stream(ordered_by: [&str; 2]) -> Result> { +/// ordered by each of `orderings` within each partition: partition membership +/// and row order are independent. +fn unbounded_range_stream(orderings: &[[&str; 2]]) -> Result> { ordered_stream( range_partitioning_multi( &[("a", SortOptions::default()), ("b", SortOptions::default())], &[&[10, 0], &[20, 0], &[30, 0]], )?, - ordered_by, + orderings, true, ) } -/// A stream table with `partitioning` whose rows are ordered by `ordered_by` -/// within each partition. +/// A stream table with `partitioning` whose rows are ordered by each of +/// `orderings` within each partition. fn ordered_stream( partitioning: Partitioning, - ordered_by: [&str; 2], + orderings: &[[&str; 2]], infinite: bool, ) -> Result> { let schema = schema(); - let ordering: LexOrdering = [ - PhysicalSortExpr::new(col(ordered_by[0], &schema)?, SortOptions::default()), - PhysicalSortExpr::new(col(ordered_by[1], &schema)?, SortOptions::default()), - ] - .into(); + let orderings = orderings + .iter() + .map(|[first, second]| { + Ok([ + PhysicalSortExpr::new(col(first, &schema)?, SortOptions::default()), + PhysicalSortExpr::new(col(second, &schema)?, SortOptions::default()), + ] + .into()) + }) + .collect::>>()?; let partitions = (0..partitioning.partition_count()) .map(|_| { Arc::new(TestStreamPartition { @@ -1257,7 +1262,7 @@ fn ordered_stream( Arc::clone(&schema), partitions, None, - vec![ordering], + orderings, infinite, None, )? @@ -1374,8 +1379,8 @@ fn range_join_key_alignment_keeps_streamable_sort_merge_join_keys() -> Result<() // A sort-merge join on (b, a) over unbounded Range([a, b]) inputs whose // rows are ordered by `ordered_by`. let smj_on_b_a = |ordered_by: [&str; 2]| -> Result> { - let left = unbounded_range_stream(ordered_by)?; - let right = unbounded_range_stream(ordered_by)?; + let left = unbounded_range_stream(&[ordered_by])?; + let right = unbounded_range_stream(&[ordered_by])?; let on = vec![ (col("b", &left.schema())?, col("b", &right.schema())?), (col("a", &left.schema())?, col("a", &right.schema())?), @@ -1437,11 +1442,11 @@ fn range_join_key_alignment_ignores_per_partition_constants() -> Result<()> { &[("a", SortOptions::default()), ("b", SortOptions::default())], &[&[10, 0], &[20, 0], &[30, 0]], )?, - ["b", "a"], + &[["b", "a"]], false, )?; let right: Arc = Arc::new(PerPartitionConstantExec::try_new( - ordered_stream(Partitioning::UnknownPartitioning(4), ["b", "a"], true)?, + ordered_stream(Partitioning::UnknownPartitioning(4), &[["b", "a"]], true)?, "a", )?); let join_on: JoinOn = vec![ @@ -1467,6 +1472,33 @@ fn range_join_key_alignment_ignores_per_partition_constants() -> Result<()> { }) } +#[test] +fn range_join_key_alignment_sees_every_ordering_of_an_unbounded_input() -> Result<()> { + // Unbounded Range([a, b]) inputs sorted by (c, d) and, separately, by + // (a, b). The second ordering satisfies the range order, so realigning + // the keys from (b, a) to (a, b) adds no sort and needs no exchange. + let left = unbounded_range_stream(&[["c", "d"], ["a", "b"]])?; + let right = unbounded_range_stream(&[["c", "d"], ["a", "b"]])?; + let join_on: JoinOn = vec![ + (col("b", &left.schema())?, col("b", &right.schema())?), + (col("a", &left.schema())?, col("a", &right.schema())?), + ]; + let join = sort_merge_join_exec(left, right, &join_on, &JoinType::Inner); + + check_range_join_key_alignment(join, |plan| { + insta::allow_duplicates! { + assert_plan!( + plan, + @r" + SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0), (b@1, b@1)] + StreamingTableExec: partition_sizes=4, projection=[a, b, c, d, e], infinite_source=true, output_partitioning=Range([a@0 ASC, b@1 ASC], [(10, 0), (20, 0), (30, 0)], 4), output_orderings=[[c@2 ASC, d@3 ASC], [a@0 ASC, b@1 ASC]] + StreamingTableExec: partition_sizes=4, projection=[a, b, c, d, e], infinite_source=true, output_partitioning=Range([a@0 ASC, b@1 ASC], [(10, 0), (20, 0), (30, 0)], 4), output_orderings=[[c@2 ASC, d@3 ASC], [a@0 ASC, b@1 ASC]] + " + ); + } + }) +} + #[test] fn range_join_key_alignment_skips_single_partition_range() -> Result<()> { // A single-partition Range layout has nothing to co-partition with, so @@ -1478,7 +1510,7 @@ fn range_join_key_alignment_skips_single_partition_range() -> Result<()> { &[("a", SortOptions::default()), ("b", SortOptions::default())], &[], )?, - ["b", "a"], + &[["b", "a"]], false, ) }; diff --git a/datafusion/sqllogictest/test_files/equivalence_properties.slt b/datafusion/sqllogictest/test_files/equivalence_properties.slt new file mode 100644 index 000000000000..a1484231e499 --- /dev/null +++ b/datafusion/sqllogictest/test_files/equivalence_properties.slt @@ -0,0 +1,194 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Planner decisions that depend on what each plan node knows about its +# output: its orderings, the columns that are constant, and the columns that +# are equal. Each case says which decision it pins. + +statement ok +set datafusion.explain.physical_plan_only = true; + +statement ok +set datafusion.execution.target_partitions = 4; + +# window_2.csv is sorted by (a, b) and, separately, by c. +statement ok +CREATE EXTERNAL TABLE t ( + a0 INTEGER, + a INTEGER, + b INTEGER, + c INTEGER, + d INTEGER +) +STORED AS CSV +WITH ORDER (a ASC, b ASC) +WITH ORDER (c ASC) +LOCATION '../core/tests/data/window_2.csv' +OPTIONS ('format.has_header' 'true'); + +statement ok +CREATE TABLE small(x INT) AS VALUES (1), (2); + +### +### Unions +### + +# Each side of the union is sorted by a and also by a cast of a, but the cast +# is in column x on the left side and in column y on the right side. The +# union must still report that its output is sorted by a, so ORDER BY a needs +# only a merge. +query TT +EXPLAIN SELECT a, CAST(a AS BIGINT) x, CAST(d AS BIGINT) y FROM t +UNION ALL +SELECT a, CAST(d AS BIGINT) x, CAST(a AS BIGINT) y FROM t +ORDER BY a; +---- +physical_plan +01)SortPreservingMergeExec: [a@0 ASC NULLS LAST] +02)--UnionExec +03)----ProjectionExec: expr=[a@0 as a, CAST(a@0 AS Int64) as x, CAST(d@1 AS Int64) as y] +04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1, maintains_sort_order=true +05)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a, d], output_ordering=[a@0 ASC NULLS LAST], file_type=csv, has_header=true +06)----ProjectionExec: expr=[a@0 as a, CAST(d@1 AS Int64) as x, CAST(a@0 AS Int64) as y] +07)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1, maintains_sort_order=true +08)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a, d], output_ordering=[a@0 ASC NULLS LAST], file_type=csv, has_header=true + +### +### GROUP BY keys in a different order than the input's orderings +### + +# The GROUP BY lists its keys as d, c, b, a, which matches neither input +# ordering. The aggregates still find that c, a and b are sorted (c from the +# ordering on c, then a and b from the ordering on (a, b)), so they run in +# partially sorted mode. The clustering mode lists those keys by their +# position in the GROUP BY, in the order they were found: [1, 3, 2]. +query TT +EXPLAIN SELECT d, c, b, a, count(*) FROM t GROUP BY d, c, b, a; +---- +physical_plan +01)ProjectionExec: expr=[d@0 as d, c@1 as c, b@2 as b, a@3 as a, count(Int64(1))@4 as count(*)] +02)--AggregateExec: mode=FinalPartitioned, gby=[d@0 as d, c@1 as c, b@2 as b, a@3 as a], aggr=[count(Int64(1))], group_clustering_mode=Partial([1, 3, 2]) +03)----RepartitionExec: partitioning=Hash([d@0, c@1, b@2, a@3], 4), input_partitions=4, preserve_order=true, sort_exprs=c@1 ASC NULLS LAST, a@3 ASC NULLS LAST, b@2 ASC NULLS LAST +04)------AggregateExec: mode=Partial, gby=[d@3 as d, c@2 as c, b@1 as b, a@0 as a], aggr=[count(Int64(1))], group_clustering_mode=Partial([1, 3, 2]) +05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1, maintains_sort_order=true +06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c, d], output_orderings=[[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], [c@2 ASC NULLS LAST]], file_type=csv, has_header=true + +### +### Constants +### + +# The filter makes b constant, and the inner ORDER BY c LIMIT 5 (planned as a +# merge with a fetch) keeps that fact. So the outer ORDER BY b, c is already +# satisfied by the order on c and the plan has no second sort. +query TT +EXPLAIN SELECT * FROM (SELECT b, c FROM t WHERE b = 1 ORDER BY c LIMIT 5) ORDER BY b, c; +---- +physical_plan +01)SortPreservingMergeExec: [c@1 ASC NULLS LAST], fetch=5 +02)--FilterExec: b@0 = 1, fetch=5 +03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1, maintains_sort_order=true +04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[b, c], output_ordering=[c@1 ASC NULLS LAST], file_type=csv, has_header=true + +query II +SELECT * FROM (SELECT b, c FROM t WHERE b = 1 ORDER BY c LIMIT 5) ORDER BY b, c; +---- +1 25 +1 26 +1 27 +1 28 +1 29 + +# A cross join keeps the constants of its left input: b = 1 still holds after +# the join, so ORDER BY l.b, l.c sorts on c only. +query TT +EXPLAIN SELECT l.b, l.c, s.x FROM (SELECT b, c FROM t WHERE b = 1) l CROSS JOIN small s ORDER BY l.b, l.c; +---- +physical_plan +01)SortExec: expr=[c@1 ASC NULLS LAST], preserve_partitioning=[false] +02)--CrossJoinExec +03)----CoalescePartitionsExec +04)------FilterExec: b@0 = 1 +05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1, maintains_sort_order=true +06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[b, c], output_ordering=[c@1 ASC NULLS LAST], file_type=csv, has_header=true +07)----DataSourceExec: partitions=1, partition_sizes=[1] + +# A cross join also keeps a right-side constant set by an equality filter: +# x = 1 still holds after the join, so ORDER BY l.c, s.x sorts on c only. +query TT +EXPLAIN SELECT l.c, s.x FROM (SELECT c FROM t) l CROSS JOIN (SELECT x FROM small WHERE x = 1) s ORDER BY l.c, s.x; +---- +physical_plan +01)SortPreservingMergeExec: [c@0 ASC NULLS LAST, x@1 ASC NULLS LAST] +02)--SortExec: expr=[c@0 ASC NULLS LAST], preserve_partitioning=[true] +03)----CrossJoinExec +04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[c], output_ordering=[c@0 ASC NULLS LAST], file_type=csv, has_header=true +05)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +06)--------FilterExec: x@0 = 1 +07)----------DataSourceExec: partitions=1, partition_sizes=[1] + +query II +SELECT l.c, s.x FROM (SELECT c FROM t) l CROSS JOIN (SELECT x FROM small WHERE x = 1) s ORDER BY l.c, s.x LIMIT 3; +---- +0 1 +1 1 +2 1 + +### +### Order-preserving repartition of an unbounded input +### + +# x = CAST(a AS BIGINT) is strictly increasing in a. The order-preserving +# repartition on x merges on its input's order (a, b, c, then x), and its +# output must still be known to be sorted by x, b as well. So the window's +# PARTITION BY x ORDER BY b runs in sorted mode without a SortExec. A sort +# here would make planning fail, because the input is unbounded. +statement ok +set datafusion.execution.target_partitions = 2; + +statement ok +CREATE UNBOUNDED EXTERNAL TABLE t_inf (a0 INTEGER, a INTEGER, b INTEGER, c INTEGER, d INTEGER) +STORED AS CSV +WITH ORDER (a ASC, b ASC, c ASC) +LOCATION '../core/tests/data/window_2.csv' +OPTIONS ('format.has_header' 'true'); + +query TT +EXPLAIN SELECT a, x, SUM(c) OVER (PARTITION BY x ORDER BY b) s +FROM (SELECT CAST(a AS BIGINT) x, a, b, c FROM t_inf); +---- +physical_plan +01)ProjectionExec: expr=[a@1 as a, x@0 as x, sum(t_inf.c) PARTITION BY [x] ORDER BY [t_inf.b ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@4 as s] +02)--BoundedWindowAggExec: wdw=[sum(t_inf.c) PARTITION BY [x] ORDER BY [t_inf.b ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(t_inf.c) PARTITION BY [x] ORDER BY [t_inf.b ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----RepartitionExec: partitioning=Hash([x@0], 2), input_partitions=2, preserve_order=true, sort_exprs=a@1 ASC NULLS LAST, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST, x@0 ASC NULLS LAST +04)------ProjectionExec: expr=[CAST(a@0 AS Int64) as x, a@0 as a, b@1 as b, c@2 as c] +05)--------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1, maintains_sort_order=true +06)----------StreamingTableExec: partition_sizes=1, projection=[a, b, c], infinite_source=true, output_ordering=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST] + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +DROP TABLE t_inf; + +statement ok +DROP TABLE small; + +statement ok +DROP TABLE t; + +statement ok +RESET datafusion.explain.physical_plan_only;