From b42a16bae8df68b7701c433de10d08ed24e23885 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Thu, 5 Nov 2026 12:26:52 -0500 Subject: [PATCH 01/19] fix: skip a leading per-partition constant when finding ordering heads A normal ordering keeps a constant whose value differs between partitions (`AcrossPartitions::Heterogeneous`, since #23528), so such a constant can lead a normal ordering. `SortKeyWalk` assumed every ordering's head was at index 0 until the first key was pinned, so it treated the constant as the head and did not see the expression after it. A window ordered by `(total DESC, k)` over an input ordered by `[k, total]`, with `k` constant within each partition, then failed `SanityCheckPlan`. Compute the initial heads past leading fixed expressions, as the walk already does after each pin. Closes #26762 Co-Authored-By: Claude Fable 5.1 --- .../src/equivalence/properties/walk.rs | 56 ++++++++++++++++++- datafusion/sqllogictest/test_files/window.slt | 51 +++++++++++++++++ 2 files changed, 106 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-expr/src/equivalence/properties/walk.rs b/datafusion/physical-expr/src/equivalence/properties/walk.rs index e0487cdd23dfa..826e6cb222418 100644 --- a/datafusion/physical-expr/src/equivalence/properties/walk.rs +++ b/datafusion/physical-expr/src/equivalence/properties/walk.rs @@ -50,12 +50,32 @@ pub(super) struct SortKeyWalk<'a> { impl<'a> SortKeyWalk<'a> { pub(super) fn new(props: &'a EquivalenceProperties) -> Self { - Self { + let mut walk = Self { props, pinned: HashSet::new(), implied: vec![], head_idxs: vec![], + }; + // A normal ordering keeps a constant whose value differs between + // partitions (see `EquivalenceGroup::normalize_sort_exprs`), and such + // a constant may lead it. Within a partition it has one value, so it + // is fixed and the head is the first expression after it. + let normal = &props.oeq_cache.normal_cls; + if normal + .iter() + .any(|ordering| walk.is_fixed(&ordering[0].expr)) + { + walk.head_idxs = normal + .iter() + .map(|ordering| { + ordering + .iter() + .position(|sort_expr| !walk.is_fixed(&sort_expr.expr)) + .unwrap_or(ordering.len()) + }) + .collect(); } + walk } /// Returns how many leading keys (in normal form) the existing orderings @@ -328,6 +348,40 @@ mod tests { }; use datafusion_expr::{ColumnarValue, Operator, Volatility, create_udf}; + /// A normal ordering keeps a constant whose value differs between + /// partitions (`AcrossPartitions::Heterogeneous`), and such a constant + /// can lead it. Within a partition the constant has one value, so it is + /// fixed and the head is the expression after it. With `k` such a + /// constant and ordering `[k, b]`, the requirements `[b]` and `[b, k]` + /// hold before anything is pinned. + /// See . + #[test] + fn leading_per_partition_constant_is_not_a_head() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])); + let [k, b] = ["k", "b"].map(|name| col(name, &schema).unwrap()); + + let mut properties = EquivalenceProperties::new(Arc::clone(&schema)); + properties.add_constants([ConstExpr::new( + Arc::clone(&k), + AcrossPartitions::Heterogeneous, + )])?; + // Added after the constant, so that `k` stays in the normal ordering: + properties.add_ordering( + [&k, &b].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), + ); + + for required in [vec![&b], vec![&b, &k], vec![&k, &b]] { + let required = required + .into_iter() + .map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))); + assert!(properties.ordering_satisfy(required)?); + } + Ok(()) + } + /// Pinning a key can move the head of an ordering to an expression that /// is equal to a function of the expressions after it, which implies a /// new ordering. With ordering `[k, a, c]` and `CAST(c AS DATE) = a`, diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 8508ca6ff66bb..54dee349d8be8 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -7834,3 +7834,54 @@ DROP TABLE range_null_peers; statement ok DROP TABLE range_nan_peers; + +# A window ordered by `(total DESC, k)` over an input sorted by `[k, total]`, +# where `k` is constant within each partition (`lower(g)` under `WHERE g = 'a'`). +# Within a partition `k` has one value, so the input ordering satisfies the +# window's requirement and no sort is added. +# https://github.com/apache/datafusion/issues/26762 +statement ok +set datafusion.execution.target_partitions = 1; + +statement ok +CREATE TABLE per_partition_constant(g VARCHAR, v BIGINT) AS VALUES ('a', 1), ('b', 2); + +query TI +SELECT k, DENSE_RANK() OVER (ORDER BY total DESC, k) AS r +FROM ( + SELECT k, SUM(s) OVER (PARTITION BY k) AS total + FROM (SELECT lower(g) AS k, SUM(v) AS s FROM per_partition_constant WHERE g = 'a' GROUP BY lower(g)) +); +---- +a 1 + +statement ok +set datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT k, DENSE_RANK() OVER (ORDER BY total DESC, k) AS r +FROM ( + SELECT k, SUM(s) OVER (PARTITION BY k) AS total + FROM (SELECT lower(g) AS k, SUM(v) AS s FROM per_partition_constant WHERE g = 'a' GROUP BY lower(g)) +); +---- +physical_plan +01)ProjectionExec: expr=[k@0 as k, dense_rank() ORDER BY [total DESC NULLS FIRST, k ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as r] +02)--BoundedWindowAggExec: wdw=[dense_rank() ORDER BY [total DESC NULLS FIRST, k ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "dense_rank() ORDER BY [total DESC NULLS FIRST, k ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----ProjectionExec: expr=[k@0 as k, sum(s) PARTITION BY [k] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@2 as total] +04)------WindowAggExec: wdw=[sum(s) PARTITION BY [k] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "sum(s) PARTITION BY [k] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] +05)--------ProjectionExec: expr=[lower(per_partition_constant.g)@0 as k, sum(per_partition_constant.v)@1 as s] +06)----------AggregateExec: mode=Single, gby=[lower(g@0) as lower(per_partition_constant.g)], aggr=[sum(per_partition_constant.v)] +07)------------FilterExec: g@0 = a +08)--------------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +reset datafusion.explain.physical_plan_only; + +statement ok +DROP TABLE per_partition_constant; + +# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# restore that value rather than resetting. +statement ok +set datafusion.execution.target_partitions = 4; From 6e6a8ae45c794ac79a0026dd6e9367ee1f045e68 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sat, 7 Nov 2026 13:03:55 +0800 Subject: [PATCH 02/19] fix: keep the sort above a projection that cannot prove the ordering --- .../src/equivalence/properties/walk.rs | 8 ++- .../enforce_sorting/sort_pushdown.rs | 24 +++++++-- datafusion/sqllogictest/test_files/window.slt | 51 +++++++++++++++++++ 3 files changed, 78 insertions(+), 5 deletions(-) diff --git a/datafusion/physical-expr/src/equivalence/properties/walk.rs b/datafusion/physical-expr/src/equivalence/properties/walk.rs index 826e6cb222418..1ed6c260710e7 100644 --- a/datafusion/physical-expr/src/equivalence/properties/walk.rs +++ b/datafusion/physical-expr/src/equivalence/properties/walk.rs @@ -43,8 +43,8 @@ pub(super) struct SortKeyWalk<'a> { /// found by [`EquivalenceProperties::implied_ordering`] as the walk goes. implied: Vec, /// For each ordering (the normal orderings of `props`, then `implied`), - /// the index of its head. Empty until the first pin, so that a walk with - /// nothing pinned allocates nothing; until then every head is at index 0. + /// the index of its head. Empty until a head moves, so that a walk that + /// moves no head allocates nothing; until then every head is at index 0. head_idxs: Vec, } @@ -296,6 +296,10 @@ impl<'a> SortKeyWalk<'a> { if class.is_some_and(|class| class.constant.is_some()) { return true; } + // A column is only fixed if it is pinned or constant: + if expr.is::() { + return false; + } // A volatile expression can change from row to row, even when all of // its inputs are fixed: if expr.is_volatile_node() { diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs index f2a01598b982e..3493c26a56428 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs @@ -1318,10 +1318,28 @@ fn handle_projection_pushdown( // (expressed in the projection's output schema) must be remapped into the // child schema before being pushed down. Forwarding `score@1` unchanged // could refer to a different input column and sort on the wrong key. - Ok( + let Some(requirements) = remap_requirement_through_projection(projection_exec, parent_required) - .map(|requirements| vec![Some(requirements)]), - ) + else { + return Ok(None); + }; + // If the input already has this ordering, no sort will be added below the + // projection. The parent still needs the projection to prove it, which it + // may not be able to do: the proof can depend on a column that it drops, + // such as a leading per-partition constant. Keep the sort above the + // projection in that case. + if projection_exec + .input() + .equivalence_properties() + .ordering_satisfy_requirement(requirements.first().clone())? + && !projection_exec + .properties() + .equivalence_properties() + .ordering_satisfy_requirement(parent_required.first().clone())? + { + return Ok(None); + } + Ok(Some(vec![Some(requirements)])) } #[cfg(test)] diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 54dee349d8be8..74791f544f8f1 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -7878,6 +7878,57 @@ physical_plan statement ok reset datafusion.explain.physical_plan_only; +# Same, with `k` a per-partition constant that is not derived from a filter +# (one row per partition), at more than one target partition. +statement ok +set datafusion.execution.target_partitions = 4; + +query IIII +SELECT k, c, total, DENSE_RANK() OVER (ORDER BY total DESC, k) AS r +FROM ( + SELECT k, c, SUM(c) OVER (PARTITION BY k) AS total + FROM (SELECT SUM(v) AS k, COUNT(*) AS c FROM per_partition_constant) +); +---- +3 2 2 1 + +statement ok +set datafusion.execution.target_partitions = 1; + +# When a projection drops the leading constant, it can no longer prove the +# ordering that its input has, so the sort stays above the projection. +query II +SELECT total, DENSE_RANK() OVER (ORDER BY total DESC) AS r +FROM ( + SELECT SUM(s) OVER (PARTITION BY k) AS total + FROM (SELECT lower(g) AS k, SUM(v) AS s FROM per_partition_constant WHERE g = 'a' GROUP BY lower(g)) +); +---- +1 1 + +# The same when the proof depends on any other dropped column, without a +# constant: `[k, a, c]` with `a = CAST(c AS DATE)` has `[k, c]`, but not after +# `a` is dropped. +statement ok +CREATE TABLE ordered_through_cast(k VARCHAR, a DATE, c TIMESTAMP) AS VALUES + ('a', DATE '2026-01-02', TIMESTAMP '2026-01-02 05:00:00'), + ('a', DATE '2026-01-01', TIMESTAMP '2026-01-01 09:00:00'), + ('a', DATE '2026-01-02', TIMESTAMP '2026-01-02 01:00:00'), + ('b', DATE '2026-01-03', TIMESTAMP '2026-01-03 01:00:00'); + +query TPI +SELECT k, c, ROW_NUMBER() OVER (ORDER BY k, c) AS r +FROM (SELECT k, a, c FROM ordered_through_cast WHERE a = CAST(c AS DATE) ORDER BY k, a, c LIMIT 10) +ORDER BY r; +---- +a 2026-01-01T09:00:00 1 +a 2026-01-02T01:00:00 2 +a 2026-01-02T05:00:00 3 +b 2026-01-03T01:00:00 4 + +statement ok +DROP TABLE ordered_through_cast; + statement ok DROP TABLE per_partition_constant; From c777864f7c89f8c54a7a50bc480a785df17185c7 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sat, 7 Nov 2026 14:13:36 +0800 Subject: [PATCH 03/19] fix: discover implied orderings when an ordering is added --- .../src/equivalence/properties/mod.rs | 15 +++++++ .../src/equivalence/properties/walk.rs | 34 +++++++++++++++ datafusion/sqllogictest/test_files/window.slt | 41 +++++++++++++++++++ 3 files changed, 90 insertions(+) diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 1206315bbb274..7a2f861668d06 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -392,12 +392,27 @@ impl EquivalenceProperties { .filter_map(|o| self.normalize_sort_exprs(o)) .collect(); if !normal_orderings.is_empty() { + // What the new orderings imply through known equalities. `add_constants` + // and `add_equal_conditions` discover this too, and whether an ordering + // is satisfied must not depend on which of them ran last. If it cannot + // be derived, nothing is implied. + let implied: Vec<_> = normal_orderings + .iter() + .filter_map(|ordering| { + self.implied_ordering(&ordering[0], ordering[1..].iter()) + .ok() + .flatten() + }) + .collect(); self.oeq_class.extend(orderings); // Normalize given orderings to update the cache: self.oeq_cache.normal_cls.extend(normal_orderings); // TODO: If no ordering is found to be redundant during extension, we // can use a shortcut algorithm to update the leading map. self.oeq_cache.update_map(); + // Each implied ordering is shorter than the one implying it, so this + // terminates: + self.add_orderings(implied); } } diff --git a/datafusion/physical-expr/src/equivalence/properties/walk.rs b/datafusion/physical-expr/src/equivalence/properties/walk.rs index 1ed6c260710e7..1cb7b7a370e03 100644 --- a/datafusion/physical-expr/src/equivalence/properties/walk.rs +++ b/datafusion/physical-expr/src/equivalence/properties/walk.rs @@ -352,6 +352,40 @@ mod tests { }; use datafusion_expr::{ColumnarValue, Operator, Volatility, create_udf}; + /// With `CAST(c AS DATE) = a`, the ordering `[a, c]` implies `[c]`, whether + /// the ordering or the equality comes first. + #[test] + fn implied_ordering_does_not_depend_on_call_order() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Date32, true), + Field::new("c", DataType::Timestamp(TimeUnit::Nanosecond, None), true), + ])); + let [a, c] = ["a", "c"].map(|name| col(name, &schema).unwrap()); + let cast_c = Arc::new(CastExpr::new_with_target_field( + Arc::clone(&c), + Arc::new(Field::new("c", DataType::Date32, true)), + None, + )) as Arc; + let ordering = + [&a, &c].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))); + + let mut equality_first = EquivalenceProperties::new(Arc::clone(&schema)); + equality_first.add_equal_conditions(Arc::clone(&cast_c), Arc::clone(&a))?; + equality_first.add_ordering(ordering.clone()); + + let mut ordering_first = EquivalenceProperties::new(Arc::clone(&schema)); + ordering_first.add_ordering(ordering); + ordering_first.add_equal_conditions(cast_c, Arc::clone(&a))?; + + for properties in [equality_first, ordering_first] { + assert!( + properties + .ordering_satisfy([PhysicalSortExpr::new_default(Arc::clone(&c))])? + ); + } + Ok(()) + } + /// A normal ordering keeps a constant whose value differs between /// partitions (`AcrossPartitions::Heterogeneous`), and such a constant /// can lead it. Within a partition the constant has one value, so it is diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 74791f544f8f1..44b1c51996130 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -7926,6 +7926,47 @@ a 2026-01-02T01:00:00 2 a 2026-01-02T05:00:00 3 b 2026-01-03T01:00:00 4 +# An ordering that a sort adds implies other orderings through known equalities, +# as it does when the equality comes after the sort: `[a, c]` with +# `a = CAST(c AS DATE)` has `[c]`. +query DPI +SELECT a, c, ROW_NUMBER() OVER (ORDER BY c) AS r +FROM (SELECT a, c FROM ordered_through_cast WHERE a = CAST(c AS DATE) ORDER BY a, c LIMIT 10) +ORDER BY r; +---- +2026-01-01 2026-01-01T09:00:00 1 +2026-01-02 2026-01-02T01:00:00 2 +2026-01-02 2026-01-02T05:00:00 3 +2026-01-03 2026-01-03T01:00:00 4 + +# The same with a per-partition constant leading the sort. +query TDPI +SELECT k, a, c, ROW_NUMBER() OVER (ORDER BY c) AS r +FROM (SELECT lower(k) AS k, a, c FROM ordered_through_cast WHERE k = 'a' AND a = CAST(c AS DATE) ORDER BY k, a, c LIMIT 10) +ORDER BY r; +---- +a 2026-01-01 2026-01-01T09:00:00 1 +a 2026-01-02 2026-01-02T01:00:00 2 +a 2026-01-02 2026-01-02T05:00:00 3 + +# And when a merge of partitions is added above the sort, which keeps only its +# merge key. +statement ok +set datafusion.execution.target_partitions = 4; + +query DPI +SELECT a, c, ROW_NUMBER() OVER (ORDER BY c) AS r +FROM (SELECT a, c FROM ordered_through_cast WHERE a = CAST(c AS DATE) ORDER BY a, c LIMIT 10) +ORDER BY r; +---- +2026-01-01 2026-01-01T09:00:00 1 +2026-01-02 2026-01-02T01:00:00 2 +2026-01-02 2026-01-02T05:00:00 3 +2026-01-03 2026-01-03T01:00:00 4 + +statement ok +set datafusion.execution.target_partitions = 1; + statement ok DROP TABLE ordered_through_cast; From 6425ba0c982719002ff78a0afd6dae639de77737 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sat, 7 Nov 2026 20:44:46 +0800 Subject: [PATCH 04/19] refactor: return errors from EquivalenceProperties::add_orderings --- datafusion/core/tests/execution/coop.rs | 3 +- .../tests/fuzz_cases/equivalence/ordering.rs | 5 +- .../tests/fuzz_cases/equivalence/utils.rs | 6 +- .../sort_preserving_repartition_fuzz.rs | 2 +- .../physical_optimizer/ensure_requirements.rs | 3 +- datafusion/datasource-parquet/src/source.rs | 10 +-- .../functions-table/src/generate_series.rs | 2 +- .../benches/equivalence_properties.rs | 24 ++++--- .../physical-expr/src/equivalence/mod.rs | 2 +- .../physical-expr/src/equivalence/ordering.rs | 4 +- .../src/equivalence/properties/dependency.rs | 44 ++++++------- .../src/equivalence/properties/joins.rs | 14 ++--- .../src/equivalence/properties/mod.rs | 63 +++++++++---------- .../src/equivalence/properties/union.rs | 2 +- .../src/equivalence/properties/walk.rs | 43 +++++++++---- datafusion/physical-expr/src/projection.rs | 6 +- .../physical-expr/src/window/standard.rs | 4 +- .../enforce_sorting/sort_pushdown.rs | 2 +- datafusion/physical-plan/src/memory.rs | 7 ++- .../physical-plan/src/repartition/mod.rs | 19 +++--- .../src/sorts/partitioned_topk.rs | 3 +- datafusion/physical-plan/src/sorts/sort.rs | 8 ++- .../src/sorts/sort_preserving_merge.rs | 22 ++++--- datafusion/physical-plan/src/union.rs | 2 +- datafusion/physical-plan/src/windows/mod.rs | 26 ++++---- .../library-user-guide/upgrading/56.0.0.md | 20 ++++++ 26 files changed, 203 insertions(+), 143 deletions(-) diff --git a/datafusion/core/tests/execution/coop.rs b/datafusion/core/tests/execution/coop.rs index eda649770ff5c..5349b8adf3997 100644 --- a/datafusion/core/tests/execution/coop.rs +++ b/datafusion/core/tests/execution/coop.rs @@ -148,7 +148,8 @@ fn make_lazy_exec_with_range( exec.add_ordering(vec![PhysicalSortExpr::new( Arc::new(Column::new(column_name, 0)), SortOptions::new(false, true), - )]); + )]) + .unwrap(); exec } diff --git a/datafusion/core/tests/fuzz_cases/equivalence/ordering.rs b/datafusion/core/tests/fuzz_cases/equivalence/ordering.rs index 576d7e557dad3..1684b02e4dad1 100644 --- a/datafusion/core/tests/fuzz_cases/equivalence/ordering.rs +++ b/datafusion/core/tests/fuzz_cases/equivalence/ordering.rs @@ -347,7 +347,7 @@ fn test_ordering_satisfy_all_null_ordered_column() -> Result<()> { eq_properties.add_ordering([ PhysicalSortExpr::new(Arc::clone(col_a), options), PhysicalSortExpr::new(Arc::clone(col_b), options), - ]); + ])?; let batch = RecordBatch::try_new( Arc::clone(&schema), vec![ @@ -392,7 +392,8 @@ fn test_ordering_satisfy_null_constant_operand() -> Result<()> { as Arc; for options in ALL_SORT_OPTIONS { let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - eq_properties.add_ordering([PhysicalSortExpr::new(col("a", &schema)?, options)]); + eq_properties + .add_ordering([PhysicalSortExpr::new(col("a", &schema)?, options)])?; let batch = RecordBatch::try_new( Arc::clone(&schema), vec![sorted_f64_array(options, 2)?], diff --git a/datafusion/core/tests/fuzz_cases/equivalence/utils.rs b/datafusion/core/tests/fuzz_cases/equivalence/utils.rs index 4e536ce9be7bb..bff76e5351b2f 100644 --- a/datafusion/core/tests/fuzz_cases/equivalence/utils.rs +++ b/datafusion/core/tests/fuzz_cases/equivalence/utils.rs @@ -119,7 +119,7 @@ pub fn create_random_schema( options: random_sort_options(&mut rng), }); - eq_properties.add_ordering(ordering); + eq_properties.add_ordering(ordering)?; } Ok((test_schema, eq_properties)) @@ -408,7 +408,7 @@ pub fn create_test_params() -> Result<(SchemaRef, EquivalenceProperties)> { ], ]; let orderings = convert_to_orderings(&orderings); - eq_properties.add_orderings(orderings); + eq_properties.add_orderings(orderings)?; Ok((test_schema, eq_properties)) } @@ -669,7 +669,7 @@ pub fn create_random_table_with_casts_and_windows( ALL_SORT_OPTIONS.map(|options| vec![sort(&k, opts_k), sort(&w, options)]), ); orderings.shuffle(&mut rng); - eq_properties.add_orderings(orderings); + eq_properties.add_orderings(orderings)?; let keys = vec![ (Arc::clone(&k), Some(opts_k)), diff --git a/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs b/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs index 312128750ded4..690237e0b7ff6 100644 --- a/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs @@ -90,7 +90,7 @@ mod sp_repartition_fuzz_tests { options: options_asc, }); - eq_properties.add_ordering(ordering); + eq_properties.add_ordering(ordering)?; } Ok((test_schema, eq_properties)) diff --git a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs index 86adb307db75b..49c42f8cecc8c 100644 --- a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs @@ -87,7 +87,8 @@ impl MockMultiPartitionExec { nulls_first: false, }, )]) { - eq.add_orderings(vec![ordering.into_iter().collect::>()]); + eq.add_orderings(vec![ordering.into_iter().collect::>()]) + .unwrap(); } let properties = PlanProperties::new( eq, diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index ac7e10f94e6f5..1e98f0dd2c47d 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -1014,7 +1014,7 @@ impl FileSource for ParquetSource { .collect::>() }) .collect::>(); - new.add_orderings(reversed_orderings); + new.add_orderings(reversed_orderings)?; new }; let reversed_satisfies = @@ -1712,7 +1712,8 @@ mod tests { let source = ParquetSource::new(Arc::clone(&schema)); // Source declares `[a DESC]`. let mut eq = EquivalenceProperties::new(Arc::clone(&schema)); - eq.add_ordering(vec![sort_expr_on(&schema, "a", true)]); + eq.add_ordering(vec![sort_expr_on(&schema, "a", true)]) + .unwrap(); // Request `[a ASC]`. let order = vec![sort_expr_on(&schema, "a", false)]; @@ -1785,7 +1786,7 @@ mod tests { }; let declared = request_expr.reverse(); let mut eq = EquivalenceProperties::new(Arc::clone(&full_schema)); - eq.add_ordering(vec![declared]); + eq.add_ordering(vec![declared]).unwrap(); let order = vec![request_expr]; let result = source.try_pushdown_sort(&order, &eq).unwrap(); @@ -1847,7 +1848,8 @@ mod tests { nulls_first: false, }, }, - ]); + ]) + .unwrap(); // Request `[a DESC, b ASC NULLS LAST, c DESC]` — three columns; // source's two-column declaration is a strict prefix. diff --git a/datafusion/functions-table/src/generate_series.rs b/datafusion/functions-table/src/generate_series.rs index 7efd6d9e18f41..fb53d537a10a6 100644 --- a/datafusion/functions-table/src/generate_series.rs +++ b/datafusion/functions-table/src/generate_series.rs @@ -730,7 +730,7 @@ impl TableProvider for GenerateSeriesTable { .with_projection(projection.map(|p| p.to_vec())); if let Some(ordering) = self.output_ordering(exec.schema().as_ref()) { - exec.add_ordering([ordering]); + exec.add_ordering([ordering])?; } Ok(Arc::new(exec)) diff --git a/datafusion/physical-expr/benches/equivalence_properties.rs b/datafusion/physical-expr/benches/equivalence_properties.rs index 0b14efcd3e24f..645de1d8cdbdb 100644 --- a/datafusion/physical-expr/benches/equivalence_properties.rs +++ b/datafusion/physical-expr/benches/equivalence_properties.rs @@ -64,11 +64,13 @@ fn asc(i: usize) -> PhysicalSortExpr { fn properties(n_classes: usize) -> EquivalenceProperties { let schema = schema(8 + 2 * n_classes); let mut props = EquivalenceProperties::new(schema); - props.add_orderings([ - vec![asc(0), asc(1), asc(2), asc(3)], - vec![asc(4), asc(5)], - vec![asc(6)], - ]); + props + .add_orderings([ + vec![asc(0), asc(1), asc(2), asc(3)], + vec![asc(4), asc(5)], + vec![asc(6)], + ]) + .unwrap(); for i in 0..n_classes { props .add_equal_conditions(col(8 + 2 * i), col(9 + 2 * i)) @@ -225,7 +227,7 @@ fn bench_shapes(c: &mut Criterion) { ]); } } - props.add_orderings(orderings); + props.add_orderings(orderings).unwrap(); let keys = vec![ PhysicalSortExpr::new(col(0), asc_nulls_last), PhysicalSortExpr::new(col(1), asc_nulls_last), @@ -307,10 +309,12 @@ fn bench_project_wide_struct_fields(c: &mut Criterion) { let mut group = c.benchmark_group("equivalence_properties"); for (suffix, sort_col) in [("", col_a), ("/ordered_by_s", col_s)] { let mut props = EquivalenceProperties::new(Arc::clone(&schema)); - props.add_orderings([vec![PhysicalSortExpr::new( - sort_col, - SortOptions::default(), - )]]); + props + .add_orderings([vec![PhysicalSortExpr::new( + sort_col, + SortOptions::default(), + )]]) + .unwrap(); group.bench_function( format!("project/wide_struct_fields/{READ}_of_{FIELDS}{suffix}"), |b| b.iter(|| props.project(&mapping, Arc::clone(&output))), diff --git a/datafusion/physical-expr/src/equivalence/mod.rs b/datafusion/physical-expr/src/equivalence/mod.rs index ff57146efcce9..5d228ac2f9714 100644 --- a/datafusion/physical-expr/src/equivalence/mod.rs +++ b/datafusion/physical-expr/src/equivalence/mod.rs @@ -137,7 +137,7 @@ mod tests { ], ]; let orderings = convert_to_orderings(&orderings); - eq_properties.add_orderings(orderings); + eq_properties.add_orderings(orderings)?; Ok((test_schema, eq_properties)) } diff --git a/datafusion/physical-expr/src/equivalence/ordering.rs b/datafusion/physical-expr/src/equivalence/ordering.rs index 7de47767f62fd..1d938b29cf0c5 100644 --- a/datafusion/physical-expr/src/equivalence/ordering.rs +++ b/datafusion/physical-expr/src/equivalence/ordering.rs @@ -662,7 +662,7 @@ mod tests { ); let mut eq_properties = EquivalenceProperties::new(Arc::clone(&test_schema)); let orderings = convert_to_orderings(&orderings); - eq_properties.add_orderings(orderings); + eq_properties.add_orderings(orderings)?; let classes = eq_group .into_iter() .map(|eq_class| EquivalenceClass::new(eq_class.into_iter().cloned())); @@ -706,7 +706,7 @@ mod tests { let orderings = convert_to_orderings(&orderings); // Column [a ASC], [e ASC], [d ASC, f ASC] are all valid orderings for the schema. - eq_properties.add_orderings(orderings); + eq_properties.add_orderings(orderings)?; // First entry in the tuple is required ordering, second entry is the expected flag // that indicates whether this required ordering is satisfied. diff --git a/datafusion/physical-expr/src/equivalence/properties/dependency.rs b/datafusion/physical-expr/src/equivalence/properties/dependency.rs index b626e644a2a9b..836aa4247249d 100644 --- a/datafusion/physical-expr/src/equivalence/properties/dependency.rs +++ b/datafusion/physical-expr/src/equivalence/properties/dependency.rs @@ -460,7 +460,7 @@ mod tests { parse_sort_expr("b", &input_schema), parse_sort_expr("c", &input_schema), parse_sort_expr("d", &input_schema), - ]); + ])?; // add equivalent ordering [a, c, b, d] input_properties.add_ordering([ @@ -468,7 +468,7 @@ mod tests { parse_sort_expr("c", &input_schema), parse_sort_expr("b", &input_schema), // NB b and c are swapped parse_sort_expr("d", &input_schema), - ]); + ])?; // simply project all the columns in order let proj_exprs = vec![ @@ -504,13 +504,13 @@ mod tests { eq_properties.add_orderings([ vec![PhysicalSortExpr::new_default(Arc::clone(&col_b_expr))], vec![PhysicalSortExpr::new_default(Arc::clone(&col_c_expr))], - ]); + ])?; let mut expected_eqs = EquivalenceProperties::new(Arc::new(schema)); expected_eqs.add_orderings([ vec![PhysicalSortExpr::new_default(col_b_expr)], vec![PhysicalSortExpr::new_default(col_c_expr)], - ]); + ])?; assert!(eq_properties.oeq_class().eq(expected_eqs.oeq_class())); Ok(()) @@ -532,7 +532,7 @@ mod tests { eq_properties.add_ordering([ PhysicalSortExpr::new(Arc::new(Column::new("b", 1)), sort_options_not), PhysicalSortExpr::new(Arc::new(Column::new("a", 0)), sort_options), - ]); + ])?; let (result, idxs) = eq_properties.find_longest_permutation(&required_columns)?; assert_eq!(idxs, vec![0, 1]); assert_eq!( @@ -561,7 +561,7 @@ mod tests { PhysicalSortExpr::new(Arc::new(Column::new("b", 1)), sort_options_not), PhysicalSortExpr::new(Arc::new(Column::new("a", 0)), sort_options), ], - ]); + ])?; let (result, idxs) = eq_properties.find_longest_permutation(&required_columns)?; assert_eq!(idxs, vec![0, 1]); assert_eq!( @@ -588,7 +588,7 @@ mod tests { PhysicalSortExpr::new(Arc::new(Column::new("b", 1)), sort_options_not), PhysicalSortExpr::new(Arc::new(Column::new("c", 2)), sort_options), PhysicalSortExpr::new(Arc::new(Column::new("a", 0)), sort_options), - ]); + ])?; let (_, idxs) = eq_properties.find_longest_permutation(&required_columns)?; assert_eq!(idxs, vec![0]); @@ -621,7 +621,7 @@ mod tests { for known in all_options { let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); eq_properties - .add_ordering([PhysicalSortExpr::new(Arc::clone(&sqrt_a), known)]); + .add_ordering([PhysicalSortExpr::new(Arc::clone(&sqrt_a), known)])?; for required in all_options { let required_ordering = [PhysicalSortExpr::new(Arc::clone(&sqrt_a), required)]; @@ -638,7 +638,7 @@ mod tests { let required_ordering = [PhysicalSortExpr::new_default(Arc::clone(&sqrt_a))]; let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); assert!(!eq_properties.ordering_satisfy(required_ordering.clone())?); - eq_properties.add_ordering([PhysicalSortExpr::new_default(a)]); + eq_properties.add_ordering([PhysicalSortExpr::new_default(a)])?; assert!(!eq_properties.ordering_satisfy(required_ordering)?); Ok(()) } @@ -667,7 +667,7 @@ mod tests { eq_properties.add_orderings([ vec![PhysicalSortExpr::new(Arc::clone(&col_b), option_asc)], vec![PhysicalSortExpr::new(Arc::clone(&col_d), option_asc)], - ]); + ])?; let test_cases = vec![ // d + b can wrap @@ -736,7 +736,7 @@ mod tests { eq_properties.add_ordering([ PhysicalSortExpr::new(Arc::clone(col_d), option_asc), PhysicalSortExpr::new(Arc::clone(col_h), option_desc), - ]); + ])?; let test_cases = vec![ // TEST CASE 1 (vec![col_a], vec![(col_a, option_asc)]), @@ -1078,7 +1078,7 @@ mod tests { ] .into(); - eq_properties.add_ordering(initial_ordering.clone()); + eq_properties.add_ordering(initial_ordering.clone())?; // Add equality condition c = concat(a, b) eq_properties.add_equal_conditions(Arc::clone(&col_c), a_concat_b)?; @@ -1120,7 +1120,7 @@ mod tests { ] .into(); - eq_properties.add_ordering(initial_ordering.clone()); + eq_properties.add_ordering(initial_ordering.clone())?; // Add equality condition c = a * b eq_properties.add_equal_conditions(col_c, a_times_b)?; @@ -1210,7 +1210,7 @@ mod tests { PhysicalSortExpr::new(Arc::clone(&col_a), asc), PhysicalSortExpr::new(Arc::clone(&col_b), desc), PhysicalSortExpr::new(Arc::clone(&col_c), asc), - ]); + ])?; // New ordering: [a ASC] let new_order = vec![PhysicalSortExpr::new(Arc::clone(&col_a), asc)]; @@ -1248,7 +1248,7 @@ mod tests { eq_properties.add_ordering([ PhysicalSortExpr::new_default(Arc::clone(&col_a)), PhysicalSortExpr::new_default(Arc::clone(&col_c)), - ]); + ])?; // New ordering: [b ASC] let new_order = vec![PhysicalSortExpr::new_default(Arc::clone(&col_b))]; @@ -1289,7 +1289,7 @@ mod tests { eq_properties.add_ordering([ PhysicalSortExpr::new(Arc::clone(&col_a), asc), PhysicalSortExpr::new(Arc::clone(&col_b), desc), - ]); + ])?; // New ordering: [a DESC] let new_order = vec![PhysicalSortExpr::new(Arc::clone(&col_a), desc)]; @@ -1329,7 +1329,7 @@ mod tests { PhysicalSortExpr::new_default(Arc::clone(&col_a)), ], vec![PhysicalSortExpr::new_default(Arc::clone(&col_e))], - ]); + ])?; // New ordering: [b ASC, c ASC] let new_order = vec![ @@ -1478,7 +1478,7 @@ mod tests { expr: col(col_name, schema).unwrap(), options: SortOptions::default(), }); - eq_properties.add_ordering(base_ordering); + eq_properties.add_ordering(base_ordering)?; // Add constraints eq_properties = @@ -1564,7 +1564,7 @@ mod tests { nulls_first: false, }, ); - input_properties.add_orderings([vec![sort_a]]); + input_properties.add_orderings([vec![sort_a]])?; // Project through the named_struct let out_schema = output_schema(&projection_mapping, &input_schema)?; @@ -1728,7 +1728,7 @@ mod tests { input.add_orderings([vec![PhysicalSortExpr::new( Arc::clone(&a), SortOptions::default(), - )]]); + )]])?; // Several sources over each column, so all but the first of each are // served from the cache: `a` through the sort key, `b` as a plain @@ -1777,7 +1777,7 @@ mod tests { input.add_orderings([vec![PhysicalSortExpr::new( Arc::clone(&s), SortOptions::default(), - )]]); + )]])?; // Every probe refers to the sort key, so both substituting the // ordering and deriving the projected one examine each of them. @@ -1838,7 +1838,7 @@ mod tests { input.add_orderings([vec![PhysicalSortExpr::new( Arc::clone(&untyped), SortOptions::default(), - )]]); + )]])?; // The probe is ordered only through a derivation from `a + b`, which // fails before the probe is asked for its properties. The output type diff --git a/datafusion/physical-expr/src/equivalence/properties/joins.rs b/datafusion/physical-expr/src/equivalence/properties/joins.rs index 2ba46f5d9b920..3c4175ee0f24b 100644 --- a/datafusion/physical-expr/src/equivalence/properties/joins.rs +++ b/datafusion/physical-expr/src/equivalence/properties/joins.rs @@ -57,9 +57,9 @@ pub fn join_equivalence_properties( // then we should add `a ASC, b ASC` to the ordering equivalences // of the join output. let out_oeq_class = left_oeq_class.join_suffix(&right_oeq_class); - result.add_orderings(out_oeq_class); + result.add_orderings(out_oeq_class)?; } else { - result.add_orderings(left_oeq_class); + result.add_orderings(left_oeq_class)?; } } [false, true] => { @@ -82,9 +82,9 @@ pub fn join_equivalence_properties( // then we should add `b ASC, a ASC` to the ordering equivalences // of the join output. let out_oeq_class = right_oeq_class.join_suffix(&left_oeq_class); - result.add_orderings(out_oeq_class); + result.add_orderings(out_oeq_class)?; } else { - result.add_orderings(right_oeq_class); + result.add_orderings(right_oeq_class)?; } } [false, false] => {} @@ -185,8 +185,8 @@ mod tests { let left_orderings = convert_to_orderings(&left_orderings); let right_orderings = convert_to_orderings(&right_orderings); let expected = convert_to_orderings(&expected); - left_eq_properties.add_orderings(left_orderings); - right_eq_properties.add_orderings(right_orderings); + left_eq_properties.add_orderings(left_orderings)?; + right_eq_properties.add_orderings(right_orderings)?; let join_eq = join_equivalence_properties( left_eq_properties, right_eq_properties, @@ -261,7 +261,7 @@ mod tests { &join_type, left_columns_len, )?; - join_eq_properties.add_orderings(right_oeq_class); + join_eq_properties.add_orderings(right_oeq_class)?; let result = join_eq_properties.oeq_class().clone(); // [x ASC, y ASC], [z ASC, w ASC] diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 7a2f861668d06..4a7cb0b227ef3 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -114,7 +114,7 @@ use itertools::{Either, Itertools}; /// eq_properties.add_ordering([ /// PhysicalSortExpr::new_default(col_a).asc(), /// PhysicalSortExpr::new_default(col_c).desc(), -/// ]); +/// ]).unwrap(); /// /// assert_eq!( /// eq_properties.to_string(), @@ -345,7 +345,7 @@ impl EquivalenceProperties { pub fn extend(mut self, other: Self) -> Result { self.constraints.extend(other.constraints); self.add_equivalence_group(other.eq_group)?; - self.add_orderings(other.oeq_class); + self.add_orderings(other.oeq_class)?; self.add_groupings(other.geq_class); Ok(self) } @@ -379,11 +379,12 @@ impl EquivalenceProperties { } } - /// Adds new orderings into the existing ordering equivalence class. + /// Adds new orderings into the existing ordering equivalence class, along + /// with the orderings they imply through known equalities. pub fn add_orderings( &mut self, orderings: impl IntoIterator>, - ) { + ) -> Result<()> { let orderings: Vec<_> = orderings.into_iter().filter_map(LexOrdering::new).collect(); let normal_orderings: Vec<_> = orderings @@ -391,34 +392,33 @@ impl EquivalenceProperties { .cloned() .filter_map(|o| self.normalize_sort_exprs(o)) .collect(); - if !normal_orderings.is_empty() { - // What the new orderings imply through known equalities. `add_constants` - // and `add_equal_conditions` discover this too, and whether an ordering - // is satisfied must not depend on which of them ran last. If it cannot - // be derived, nothing is implied. - let implied: Vec<_> = normal_orderings - .iter() - .filter_map(|ordering| { - self.implied_ordering(&ordering[0], ordering[1..].iter()) - .ok() - .flatten() - }) - .collect(); - self.oeq_class.extend(orderings); - // Normalize given orderings to update the cache: - self.oeq_cache.normal_cls.extend(normal_orderings); - // TODO: If no ordering is found to be redundant during extension, we - // can use a shortcut algorithm to update the leading map. - self.oeq_cache.update_map(); - // Each implied ordering is shorter than the one implying it, so this - // terminates: - self.add_orderings(implied); + if normal_orderings.is_empty() { + return Ok(()); + } + // `add_constants` and `add_equal_conditions` discover implied orderings + // too, and whether an ordering is satisfied must not depend on which of + // them ran last. + let mut implied = vec![]; + for ordering in &normal_orderings { + implied.extend(self.implied_ordering(&ordering[0], ordering[1..].iter())?); } + self.oeq_class.extend(orderings); + // Normalize given orderings to update the cache: + self.oeq_cache.normal_cls.extend(normal_orderings); + // TODO: If no ordering is found to be redundant during extension, we + // can use a shortcut algorithm to update the leading map. + self.oeq_cache.update_map(); + // Each implied ordering is shorter than the one implying it, so this + // terminates: + self.add_orderings(implied) } /// Adds a single ordering to the existing ordering equivalence class. - pub fn add_ordering(&mut self, ordering: impl IntoIterator) { - self.add_orderings(std::iter::once(ordering)); + pub fn add_ordering( + &mut self, + ordering: impl IntoIterator, + ) -> Result<()> { + self.add_orderings(std::iter::once(ordering)) } /// Adds complete grouping tuples whose values are contiguous within @@ -586,10 +586,7 @@ impl EquivalenceProperties { new_orderings .extend(self.implied_ordering(&ordering[0], ordering[1..].iter())?); } - if !new_orderings.is_empty() { - self.add_orderings(new_orderings); - } - Ok(()) + self.add_orderings(new_orderings) } /// Returns the ordering that the normal ordering `[leading, rest..]` @@ -677,7 +674,7 @@ impl EquivalenceProperties { if normal_ordering.len() != self.common_sort_prefix_length(&normal_ordering)? { // If the ordering is unsatisfied, replace existing orderings: self.clear_orderings(); - self.add_ordering(ordering_tee); + self.add_ordering(ordering_tee)?; return Ok(true); } Ok(false) diff --git a/datafusion/physical-expr/src/equivalence/properties/union.rs b/datafusion/physical-expr/src/equivalence/properties/union.rs index a3d5c1c454a47..111456cbc3a38 100644 --- a/datafusion/physical-expr/src/equivalence/properties/union.rs +++ b/datafusion/physical-expr/src/equivalence/properties/union.rs @@ -63,7 +63,7 @@ fn calculate_union_binary( let mut eq_properties = EquivalenceProperties::new(lhs.schema); eq_properties.add_constants(constants)?; - eq_properties.add_orderings(orderings); + eq_properties.add_orderings(orderings)?; Ok(eq_properties) } diff --git a/datafusion/physical-expr/src/equivalence/properties/walk.rs b/datafusion/physical-expr/src/equivalence/properties/walk.rs index 1cb7b7a370e03..bb8a016a220ad 100644 --- a/datafusion/physical-expr/src/equivalence/properties/walk.rs +++ b/datafusion/physical-expr/src/equivalence/properties/walk.rs @@ -338,7 +338,7 @@ impl<'a> SortKeyWalk<'a> { mod tests { use std::sync::Arc; - use crate::expressions::{BinaryExpr, CastExpr, col, lit}; + use crate::expressions::{BinaryExpr, CastExpr, Column, col, lit}; use crate::scalar_subquery::ScalarSubqueryExpr; use crate::{AcrossPartitions, ConstExpr, EquivalenceProperties}; use crate::{PhysicalExpr, PhysicalSortExpr, ScalarFunctionExpr}; @@ -371,10 +371,10 @@ mod tests { let mut equality_first = EquivalenceProperties::new(Arc::clone(&schema)); equality_first.add_equal_conditions(Arc::clone(&cast_c), Arc::clone(&a))?; - equality_first.add_ordering(ordering.clone()); + equality_first.add_ordering(ordering.clone())?; let mut ordering_first = EquivalenceProperties::new(Arc::clone(&schema)); - ordering_first.add_ordering(ordering); + ordering_first.add_ordering(ordering)?; ordering_first.add_equal_conditions(cast_c, Arc::clone(&a))?; for properties in [equality_first, ordering_first] { @@ -386,6 +386,27 @@ mod tests { Ok(()) } + /// `add_ordering` reports what it cannot derive instead of dropping it: + /// here `c` is not in the schema, so its type is unknown. + #[test] + fn add_ordering_returns_implied_ordering_errors() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Date32, true)])); + let a = col("a", &schema)?; + let c = Arc::new(Column::new("c", 1)) as Arc; + let cast_c = Arc::new(CastExpr::new_with_target_field( + Arc::clone(&c), + Arc::new(Field::new("c", DataType::Date32, true)), + None, + )) as Arc; + + let mut properties = EquivalenceProperties::new(Arc::clone(&schema)); + properties.add_equal_conditions(cast_c, Arc::clone(&a))?; + let ordering = + [&a, &c].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))); + assert!(properties.add_ordering(ordering).is_err()); + Ok(()) + } + /// A normal ordering keeps a constant whose value differs between /// partitions (`AcrossPartitions::Heterogeneous`), and such a constant /// can lead it. Within a partition the constant has one value, so it is @@ -409,7 +430,7 @@ mod tests { // Added after the constant, so that `k` stays in the normal ordering: properties.add_ordering( [&k, &b].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), - ); + )?; for required in [vec![&b], vec![&b, &k], vec![&k, &b]] { let required = required @@ -441,7 +462,7 @@ mod tests { let mut properties = EquivalenceProperties::new(Arc::clone(&schema)); properties.add_ordering( [&k, &a, &c].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), - ); + )?; properties.add_equal_conditions(cast_c, Arc::clone(&a))?; let required = @@ -478,7 +499,7 @@ mod tests { properties.add_equal_conditions(Arc::clone(&x), cast(&y, DataType::Int64))?; properties.add_ordering( [&k, &x, &y, &z].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), - ); + )?; let required = [&k, &z].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))); @@ -513,7 +534,7 @@ mod tests { vec![sort(&a, true), sort(&k, false), sort(&c, false)], vec![sort(&k, false), sort(&a, true), sort(&c, false)], vec![sort(&k, false), sort(&a, false)], - ]); + ])?; assert!(properties.ordering_satisfy([sort(&k, false), sort(&a, false)])?); Ok(()) @@ -549,7 +570,7 @@ mod tests { PhysicalSortExpr::new_default(Arc::clone(&k)), PhysicalSortExpr::new(Arc::clone(&x), options), ] - })); + }))?; for options in directions { let required = [ @@ -586,7 +607,7 @@ mod tests { properties.add_ordering([ PhysicalSortExpr::new_default(plus(&b)), PhysicalSortExpr::new_default(Arc::clone(&d)), - ]); + ])?; for right in [&b, &c] { let (_, indices) = @@ -619,7 +640,7 @@ mod tests { properties.add_ordering([ PhysicalSortExpr::new_default(Arc::clone(&k)), PhysicalSortExpr::new_default(Arc::clone(&b_mod_5)), - ]); + ])?; let (_, indices) = properties.find_longest_permutation(&[lit(5i32), k, b_mod_5])?; @@ -659,7 +680,7 @@ mod tests { }; let mut properties = EquivalenceProperties::new(Arc::clone(&schema)); - properties.add_ordering([PhysicalSortExpr::new_default(Arc::clone(&a))]); + properties.add_ordering([PhysicalSortExpr::new_default(Arc::clone(&a))])?; let satisfied = |exprs: [&Arc; 2]| { properties.ordering_satisfy( exprs.map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), diff --git a/datafusion/physical-expr/src/projection.rs b/datafusion/physical-expr/src/projection.rs index 79c2dc56adc00..efe993be84650 100644 --- a/datafusion/physical-expr/src/projection.rs +++ b/datafusion/physical-expr/src/projection.rs @@ -1966,7 +1966,7 @@ pub(crate) mod tests { let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); let orderings = convert_to_orderings(&orderings); - eq_properties.add_orderings(orderings); + eq_properties.add_orderings(orderings)?; let proj_exprs = proj_exprs .into_iter() @@ -2169,7 +2169,7 @@ pub(crate) mod tests { let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); let orderings = convert_to_orderings(orderings); - eq_properties.add_orderings(orderings); + eq_properties.add_orderings(orderings)?; let expected = convert_to_orderings(expected); @@ -2296,7 +2296,7 @@ pub(crate) mod tests { } let orderings = convert_to_orderings(&orderings); - eq_properties.add_orderings(orderings); + eq_properties.add_orderings(orderings)?; let expected = convert_to_orderings(&expected); diff --git a/datafusion/physical-expr/src/window/standard.rs b/datafusion/physical-expr/src/window/standard.rs index 135b691615bd0..0b5fefaa50576 100644 --- a/datafusion/physical-expr/src/window/standard.rs +++ b/datafusion/physical-expr/src/window/standard.rs @@ -300,7 +300,7 @@ pub(crate) fn add_new_ordering_expr_with_partition_by( ) -> Result<()> { if partition_by.is_empty() { // In the absence of a PARTITION BY, ordering of `self.expr` is global: - eqp.add_ordering([expr]); + eqp.add_ordering([expr])?; } else { // If we have a PARTITION BY, standard functions can not introduce // a global ordering unless the existing ordering is compatible @@ -311,7 +311,7 @@ pub(crate) fn add_new_ordering_expr_with_partition_by( let (mut ordering, _) = eqp.find_longest_permutation(partition_by)?; if ordering.len() == partition_by.len() { ordering.push(expr); - eqp.add_ordering(ordering); + eqp.add_ordering(ordering)?; } } Ok(()) diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs index 3493c26a56428..0256960c062db 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs @@ -1086,7 +1086,7 @@ fn is_custom_sort_pushdown_valid( // Clear old ordering metadata so the proposal must prove this on its own. let mut proposed_properties = child_properties.clone(); proposed_properties.clear_orderings(); - proposed_properties.add_ordering(ordering.clone()); + proposed_properties.add_ordering(ordering.clone())?; if !proposed_properties.ordering_satisfy(sort.expr().clone())? { return Ok(false); } diff --git a/datafusion/physical-plan/src/memory.rs b/datafusion/physical-plan/src/memory.rs index 0688a87cb420e..71cf750367a8c 100644 --- a/datafusion/physical-plan/src/memory.rs +++ b/datafusion/physical-plan/src/memory.rs @@ -257,10 +257,13 @@ impl LazyMemoryExec { Ok(()) } - pub fn add_ordering(&mut self, ordering: impl IntoIterator) { + pub fn add_ordering( + &mut self, + ordering: impl IntoIterator, + ) -> Result<()> { Arc::make_mut(&mut self.cache) .eq_properties - .add_orderings(std::iter::once(ordering)); + .add_orderings(std::iter::once(ordering)) } /// Get the batch generators diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 13274d26e1fb4..394d61909044a 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -2178,7 +2178,7 @@ impl RepartitionExec { partitioning: Partitioning, ) -> Result { let preserve_order = false; - let cache = Self::compute_properties(&input, partitioning, preserve_order); + let cache = Self::compute_properties(&input, partitioning, preserve_order)?; Ok(RepartitionExec { input, state: Default::default(), @@ -2200,19 +2200,19 @@ impl RepartitionExec { fn eq_properties_helper( input: &Arc, preserve_order: bool, - ) -> EquivalenceProperties { + ) -> Result { let mut eq_properties = input.equivalence_properties().clone(); // Fusing several input partitions keeps only the merge key ordering, // and per-partition constants no longer hold. if input.output_partitioning().partition_count() > 1 { eq_properties.clear_orderings(); if preserve_order && let Some(ordering) = input.output_ordering() { - eq_properties.add_ordering(ordering.clone()); + eq_properties.add_ordering(ordering.clone())?; } eq_properties.clear_groupings(); eq_properties.clear_per_partition_constants(); } - eq_properties + Ok(eq_properties) } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -2220,15 +2220,15 @@ impl RepartitionExec { input: &Arc, partitioning: Partitioning, preserve_order: bool, - ) -> PlanProperties { - PlanProperties::new( - Self::eq_properties_helper(input, preserve_order), + ) -> Result { + Ok(PlanProperties::new( + Self::eq_properties_helper(input, preserve_order)?, partitioning, input.pipeline_behavior(), input.boundedness(), ) .with_scheduling_type(SchedulingType::Cooperative) - .with_evaluation_type(EvaluationType::Eager) + .with_evaluation_type(EvaluationType::Eager)) } /// Specify if this repartitioning operation should preserve the order of @@ -2245,7 +2245,8 @@ impl RepartitionExec { // if there is only one input partition, merging is not required // to maintain order self.input.output_partitioning().partition_count() > 1; - let eq_properties = Self::eq_properties_helper(&self.input, self.preserve_order); + let eq_properties = + Self::eq_properties_helper(&self.input, self.preserve_order).unwrap(); Arc::make_mut(&mut self.cache).set_eq_properties(eq_properties); self } diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs index 3950894b322f8..3d0cee02a68d0 100644 --- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs +++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs @@ -366,7 +366,8 @@ impl PartitionedTopKExec { // value only ascends *within* a partition, so the ordering has to // be prefixed by the partition keys. `reorder` above already // established that prefix. - eq_properties.add_ordering(partition_exprs.into_iter().chain([ranking_sort])); + eq_properties + .add_ordering(partition_exprs.into_iter().chain([ranking_sort]))?; } Ok(PlanProperties::new( diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 96e673e4043bc..bf90665dcb726 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -2826,9 +2826,11 @@ mod tests { impl SortedUnboundedExec { fn compute_properties(schema: SchemaRef) -> PlanProperties { let mut eq_properties = EquivalenceProperties::new(schema); - eq_properties.add_ordering([PhysicalSortExpr::new_default(Arc::new( - Column::new("c1", 0), - ))]); + eq_properties + .add_ordering([PhysicalSortExpr::new_default(Arc::new(Column::new( + "c1", 0, + )))]) + .unwrap(); PlanProperties::new( eq_properties, Partitioning::UnknownPartitioning(1), diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 641e932c280c1..86175daa8f440 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -94,7 +94,7 @@ pub struct SortPreservingMergeExec { impl SortPreservingMergeExec { /// Create a new sort execution plan pub fn new(expr: LexOrdering, input: Arc) -> Self { - let cache = Self::compute_properties(&input, expr.clone()); + let cache = Self::compute_properties(&input, expr.clone()).unwrap(); Self { input, expr, @@ -148,7 +148,7 @@ impl SortPreservingMergeExec { fn compute_properties( input: &Arc, ordering: LexOrdering, - ) -> PlanProperties { + ) -> Result { let input_partitions = input.output_partitioning().partition_count(); let (drive, scheduling) = if input_partitions > 1 { (EvaluationType::Eager, SchedulingType::Cooperative) @@ -167,15 +167,15 @@ impl SortPreservingMergeExec { eq_properties.clear_groupings(); } eq_properties.clear_per_partition_constants(); - eq_properties.add_ordering(ordering); - PlanProperties::new( + eq_properties.add_ordering(ordering)?; + Ok(PlanProperties::new( eq_properties, // Equivalence Properties Partitioning::UnknownPartitioning(1), // Output Partitioning input.pipeline_behavior(), // Pipeline Behavior input.boundedness(), // Boundedness ) .with_evaluation_type(drive) - .with_scheduling_type(scheduling) + .with_scheduling_type(scheduling)) } } @@ -1890,11 +1890,13 @@ mod tests { .map(|(i, f)| Arc::new(Column::new(f.name(), i)) as Arc) .collect::>(); let mut eq_properties = EquivalenceProperties::new(schema); - eq_properties.add_ordering( - columns - .iter() - .map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), - ); + eq_properties + .add_ordering( + columns + .iter() + .map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), + ) + .unwrap(); PlanProperties::new( eq_properties, Partitioning::Hash(columns, 3), diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 042f3217e22c5..dd4f7a76edf54 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -1662,7 +1662,7 @@ mod tests { let child2 = Arc::new(TestMemoryExec::update_cache(&child2)); let mut union_expected_eq = EquivalenceProperties::new(Arc::clone(&schema)); - union_expected_eq.add_orderings(union_expected_orderings); + union_expected_eq.add_orderings(union_expected_orderings)?; let union: Arc = UnionExec::try_new(vec![child1, child2])?; let union_eq_properties = union.properties().equivalence_properties(); diff --git a/datafusion/physical-plan/src/windows/mod.rs b/datafusion/physical-plan/src/windows/mod.rs index 1e5eea0b01285..6954dec824869 100644 --- a/datafusion/physical-plan/src/windows/mod.rs +++ b/datafusion/physical-plan/src/windows/mod.rs @@ -454,7 +454,7 @@ pub(crate) fn window_equivalence_properties( existing }) }); - window_eq_properties.add_orderings(new_lexs); + window_eq_properties.add_orderings(new_lexs)?; } } else { // The window frame is ever expanding, so set monotonicity comes @@ -483,7 +483,7 @@ pub(crate) fn window_equivalence_properties( window_eq_properties.add_ordering([PhysicalSortExpr::new( Arc::new(window_col), SortOptions::new(increasing, true), - )]); + )])?; } else { // Reverse set-monotonic cases for all orderings: for mut lex in all_satisfied_lexs.into_iter() { @@ -491,7 +491,7 @@ pub(crate) fn window_equivalence_properties( Arc::new(window_col.clone()), SortOptions::new(increasing, true), )); - window_eq_properties.add_ordering(lex); + window_eq_properties.add_ordering(lex)?; } } } @@ -553,15 +553,19 @@ pub(crate) fn window_equivalence_properties( set_monotonicity.eq(&SetMonotonicity::Increasing); let window_col = Column::new(expr.name(), i + input_schema_len); if increasing && (asc || no_partitioning) { - window_eq_properties.add_ordering([PhysicalSortExpr::new( - Arc::new(window_col), - SortOptions::new(false, false), - )]); + window_eq_properties.add_ordering([ + PhysicalSortExpr::new( + Arc::new(window_col), + SortOptions::new(false, false), + ), + ])?; } else if !increasing && (!asc || no_partitioning) { - window_eq_properties.add_ordering([PhysicalSortExpr::new( - Arc::new(window_col), - SortOptions::new(true, false), - )]); + window_eq_properties.add_ordering([ + PhysicalSortExpr::new( + Arc::new(window_col), + SortOptions::new(true, false), + ), + ])?; } } } diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 849f3237d82eb..de286462c9e0e 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -1567,6 +1567,26 @@ A custom rule list that left out `SanityCheckPlan` is now checked too. A plan whose requirements do not hold, such as a `SortMergeJoinExec` over unsorted input, now fails to plan instead of possibly returning wrong results. +### `EquivalenceProperties::add_ordering` and `add_orderings` return `Result` + +`add_ordering` and `add_orderings` now also add the orderings that a new ordering +implies through known equalities. Deriving them can fail, so both now return +`Result<()>`, as does `LazyMemoryExec::add_ordering`. Code that ignores the result +gets an `unused_must_use` warning. + +`SortPreservingMergeExec::new` and `RepartitionExec::with_preserve_order` keep +their signatures and panic on such a failure, as `SortExec::new` already does. + +**Migration guide:** + +```rust,ignore +// Before +eq_properties.add_ordering(ordering); + +// After +eq_properties.add_ordering(ordering)?; +``` + ### `AggregateUDFImpl::is_nullable` is deprecated `AggregateUDFImpl::is_nullable` is deprecated; set the nullability in the field From 9c7ee55ad9fccfa77f86f858bbc5a7f0a5d9c79e Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sat, 7 Nov 2026 20:47:18 +0800 Subject: [PATCH 05/19] fix: look past leading constants when discovering implied orderings --- .../src/equivalence/properties/mod.rs | 21 ++++- .../src/equivalence/properties/walk.rs | 87 +++++++++++++++++++ datafusion/sqllogictest/test_files/window.slt | 31 +++++++ 3 files changed, 136 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 4a7cb0b227ef3..522a5d2171890 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -400,7 +400,7 @@ impl EquivalenceProperties { // them ran last. let mut implied = vec![]; for ordering in &normal_orderings { - implied.extend(self.implied_ordering(&ordering[0], ordering[1..].iter())?); + implied.extend(self.implied_by(ordering)?); } self.oeq_class.extend(orderings); // Normalize given orderings to update the cache: @@ -583,12 +583,27 @@ impl EquivalenceProperties { let mut new_orderings = vec![]; for idx in ordering_idxs { let ordering = &self.oeq_cache.normal_cls[*idx]; - new_orderings - .extend(self.implied_ordering(&ordering[0], ordering[1..].iter())?); + new_orderings.extend(self.implied_by(ordering)?); } self.add_orderings(new_orderings) } + /// Returns what the normal ordering `ordering` implies, looking past the + /// constants that a normal ordering can keep (those that differ between + /// partitions): within a partition they have one value. + fn implied_by( + &self, + ordering: &LexOrdering, + ) -> Result>> { + let mut rest = ordering.iter().filter(|sort_expr| { + self.eq_group.is_expr_constant(&sort_expr.expr).is_none() + }); + match rest.next() { + Some(leading) => self.implied_ordering(leading, rest), + None => Ok(None), + } + } + /// Returns the ordering that the normal ordering `[leading, rest..]` /// implies through an equality on its leading expression, if any. /// diff --git a/datafusion/physical-expr/src/equivalence/properties/walk.rs b/datafusion/physical-expr/src/equivalence/properties/walk.rs index bb8a016a220ad..5e3f5f0b5f15f 100644 --- a/datafusion/physical-expr/src/equivalence/properties/walk.rs +++ b/datafusion/physical-expr/src/equivalence/properties/walk.rs @@ -386,6 +386,93 @@ mod tests { Ok(()) } + /// A normal ordering can keep a leading per-partition constant, and the + /// equality behind it still counts: with `k` such a constant and + /// `CAST(c AS DATE) = a`, the ordering `[k, a, c]` implies `[c]`, whichever + /// order the constant, the equality and the ordering are added in. + #[test] + fn implied_ordering_behind_a_leading_constant() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, true), + Field::new("a", DataType::Date32, true), + Field::new("c", DataType::Timestamp(TimeUnit::Nanosecond, None), true), + ])); + let [k, a, c] = ["k", "a", "c"].map(|name| col(name, &schema).unwrap()); + let cast_c = Arc::new(CastExpr::new_with_target_field( + Arc::clone(&c), + Arc::new(Field::new("c", DataType::Date32, true)), + None, + )) as Arc; + let ordering = + [&k, &a, &c].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))); + let constant = || ConstExpr::new(Arc::clone(&k), AcrossPartitions::Heterogeneous); + + let mut constant_equality_ordering = + EquivalenceProperties::new(Arc::clone(&schema)); + constant_equality_ordering.add_constants([constant()])?; + constant_equality_ordering + .add_equal_conditions(Arc::clone(&cast_c), Arc::clone(&a))?; + constant_equality_ordering.add_ordering(ordering.clone())?; + + let mut constant_ordering_equality = + EquivalenceProperties::new(Arc::clone(&schema)); + constant_ordering_equality.add_constants([constant()])?; + constant_ordering_equality.add_ordering(ordering.clone())?; + constant_ordering_equality + .add_equal_conditions(Arc::clone(&cast_c), Arc::clone(&a))?; + + let mut equality_ordering_constant = + EquivalenceProperties::new(Arc::clone(&schema)); + equality_ordering_constant.add_equal_conditions(cast_c, Arc::clone(&a))?; + equality_ordering_constant.add_ordering(ordering)?; + equality_ordering_constant.add_constants([constant()])?; + + for properties in [ + constant_equality_ordering, + constant_ordering_equality, + equality_ordering_constant, + ] { + assert!( + properties + .ordering_satisfy([PhysicalSortExpr::new_default(Arc::clone(&c))])? + ); + } + Ok(()) + } + + /// An implied ordering can imply another: with `CAST(c AS DATE) = a` and + /// `CAST(a AS DATE64) = d`, the ordering `[d, a, c]` implies `[a, c]`, + /// which implies `[c]`. + #[test] + fn implied_orderings_chain() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("d", DataType::Date64, true), + Field::new("a", DataType::Date32, true), + Field::new("c", DataType::Timestamp(TimeUnit::Nanosecond, None), true), + ])); + let [d, a, c] = ["d", "a", "c"].map(|name| col(name, &schema).unwrap()); + let cast = |expr: &Arc, data_type| { + Arc::new(CastExpr::new_with_target_field( + Arc::clone(expr), + Arc::new(Field::new("cast", data_type, true)), + None, + )) as Arc + }; + + let mut properties = EquivalenceProperties::new(Arc::clone(&schema)); + properties.add_equal_conditions(Arc::clone(&a), cast(&c, DataType::Date32))?; + properties.add_equal_conditions(cast(&a, DataType::Date64), Arc::clone(&d))?; + properties.add_ordering( + [&d, &a, &c].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), + )?; + + assert!( + properties + .ordering_satisfy([PhysicalSortExpr::new_default(Arc::clone(&c))])? + ); + Ok(()) + } + /// `add_ordering` reports what it cannot derive instead of dropping it: /// here `c` is not in the schema, so its type is unknown. #[test] diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 44b1c51996130..58b03b2c6ff14 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -7949,6 +7949,37 @@ a 2026-01-01 2026-01-01T09:00:00 1 a 2026-01-02 2026-01-02T01:00:00 2 a 2026-01-02 2026-01-02T05:00:00 3 +# The same when a projection drops `a`: the implied `[c]` is stored, so the +# projection can prove it and no sort is added above the projection. +query TPI +SELECT k, c, ROW_NUMBER() OVER (ORDER BY c) AS r +FROM (SELECT lower(k) AS k, a, c FROM ordered_through_cast WHERE k = 'a' AND a = CAST(c AS DATE) ORDER BY k, a, c LIMIT 10) +ORDER BY r; +---- +a 2026-01-01T09:00:00 1 +a 2026-01-02T01:00:00 2 +a 2026-01-02T05:00:00 3 + +statement ok +set datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT k, c, ROW_NUMBER() OVER (ORDER BY c) AS r +FROM (SELECT lower(k) AS k, a, c FROM ordered_through_cast WHERE k = 'a' AND a = CAST(c AS DATE) ORDER BY k, a, c LIMIT 10) +ORDER BY r; +---- +physical_plan +01)ProjectionExec: expr=[k@0 as k, c@1 as c, row_number() ORDER BY [ordered_through_cast.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as r] +02)--BoundedWindowAggExec: wdw=[row_number() ORDER BY [ordered_through_cast.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() ORDER BY [ordered_through_cast.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----ProjectionExec: expr=[k@0 as k, c@2 as c] +04)------SortExec: TopK(fetch=10), expr=[a@1 ASC NULLS LAST, c@2 ASC NULLS LAST], preserve_partitioning=[false] +05)--------ProjectionExec: expr=[lower(k@0) as k, a@1 as a, c@2 as c] +06)----------FilterExec: k@0 = a AND a@1 = CAST(c@2 AS Date32) +07)------------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +reset datafusion.explain.physical_plan_only; + # And when a merge of partitions is added above the sort, which keeps only its # merge key. statement ok From 9210b517599b985873e838a0a845f100710268b1 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sat, 7 Nov 2026 20:48:23 +0800 Subject: [PATCH 06/19] fix: compare the pushed alternative in the projection sort guard --- .../enforce_sorting/sort_pushdown.rs | 60 ++++++++++++++++--- 1 file changed, 53 insertions(+), 7 deletions(-) diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs index 0256960c062db..d7d5ee83e10c8 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs @@ -1327,15 +1327,21 @@ fn handle_projection_pushdown( // projection. The parent still needs the projection to prove it, which it // may not be able to do: the proof can depend on a column that it drops, // such as a leading per-partition constant. Keep the sort above the - // projection in that case. - if projection_exec - .input() - .equivalence_properties() - .ordering_satisfy_requirement(requirements.first().clone())? + // projection in that case. Alternatives that cannot be remapped are not + // pushed down, so compare against the first one that can. + let pushed = parent_required.iter().find(|required| { + remap_lex_requirement_through_projection(projection_exec.expr(), required) + .is_some() + }); + if let Some(pushed) = pushed + && projection_exec + .input() + .equivalence_properties() + .ordering_satisfy_requirement(requirements.first().clone())? && !projection_exec .properties() .equivalence_properties() - .ordering_satisfy_requirement(parent_required.first().clone())? + .ordering_satisfy_requirement(pushed.clone())? { return Ok(None); } @@ -1353,7 +1359,7 @@ mod tests { ScalarSubqueryResults, SubqueryIndex, }; use datafusion_physical_expr::PhysicalExpr; - use datafusion_physical_expr::expressions::{BinaryExpr, col}; + use datafusion_physical_expr::expressions::{BinaryExpr, col, lit}; use datafusion_physical_plan::empty::EmptyExec; use datafusion_physical_plan::limit::GlobalLimitExec; use datafusion_physical_plan::scalar_subquery::{ @@ -1537,6 +1543,46 @@ mod tests { assert!(matches!(remapped, OrderingRequirements::Soft(_))); } + /// The guard compares the alternative that is pushed down, not the first + /// one, which may be an alternative that cannot be remapped. Here the + /// input is sorted by `c`, and the projection proves `score` (`c` renamed), + /// but not `score % 2`, so the requirement can be pushed down. + #[test] + fn projection_guard_compares_the_pushed_alternative() -> Result<()> { + let child = child_schema(); + let empty: Arc = Arc::new(EmptyExec::new(Arc::clone(&child))); + let sorted: Arc = Arc::new(SortExec::new( + [PhysicalSortExpr::new(col("c", &child)?, DESC)].into(), + empty, + )); + let projection = ProjectionExec::try_new( + vec![ + (col("a", &child)?, "a".to_string()), + (col("c", &child)?, "score".to_string()), + (col("b", &child)?, "value".to_string()), + ], + sorted, + )?; + let out = projection.schema(); + + let computed: Arc = Arc::new(BinaryExpr::new( + col("score", &out)?, + Operator::Modulo, + lit(2), + )); + let mut required = + OrderingRequirements::new(lex([PhysicalSortRequirement::new( + computed, + Some(ASC), + )])); + required.add_alternative(lex([req("score", &out, DESC)])); + + let pushed = handle_projection_pushdown(&projection, &required)?; + let expected = OrderingRequirements::new(lex([req("c", &child, DESC)])); + assert_eq!(pushed, Some(vec![Some(expected)])); + Ok(()) + } + #[test] fn remap_preserves_all_hard_alternatives() { let projection = reordering_projection(); From 17c73e6d3127884736a9197023a4dc6a35603608 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sat, 7 Nov 2026 20:49:01 +0800 Subject: [PATCH 07/19] test: check the plans of the implied ordering window cases --- datafusion/sqllogictest/test_files/window.slt | 71 +++++++++++++++---- 1 file changed, 56 insertions(+), 15 deletions(-) diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 58b03b2c6ff14..c81a86324b725 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -7928,7 +7928,7 @@ b 2026-01-03T01:00:00 4 # An ordering that a sort adds implies other orderings through known equalities, # as it does when the equality comes after the sort: `[a, c]` with -# `a = CAST(c AS DATE)` has `[c]`. +# `a = CAST(c AS DATE)` has `[c]`, so the window needs no sort of its own. query DPI SELECT a, c, ROW_NUMBER() OVER (ORDER BY c) AS r FROM (SELECT a, c FROM ordered_through_cast WHERE a = CAST(c AS DATE) ORDER BY a, c LIMIT 10) @@ -7939,6 +7939,21 @@ ORDER BY r; 2026-01-02 2026-01-02T05:00:00 3 2026-01-03 2026-01-03T01:00:00 4 +statement ok +set datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT a, c, ROW_NUMBER() OVER (ORDER BY c) AS r +FROM (SELECT a, c FROM ordered_through_cast WHERE a = CAST(c AS DATE) ORDER BY a, c LIMIT 10) +ORDER BY r; +---- +physical_plan +01)ProjectionExec: expr=[a@0 as a, c@1 as c, row_number() ORDER BY [ordered_through_cast.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as r] +02)--BoundedWindowAggExec: wdw=[row_number() ORDER BY [ordered_through_cast.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() ORDER BY [ordered_through_cast.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: TopK(fetch=10), expr=[a@0 ASC NULLS LAST, c@1 ASC NULLS LAST], preserve_partitioning=[false] +04)------FilterExec: a@0 = CAST(c@1 AS Date32) +05)--------DataSourceExec: partitions=1, partition_sizes=[1] + # The same with a per-partition constant leading the sort. query TDPI SELECT k, a, c, ROW_NUMBER() OVER (ORDER BY c) AS r @@ -7949,6 +7964,19 @@ a 2026-01-01 2026-01-01T09:00:00 1 a 2026-01-02 2026-01-02T01:00:00 2 a 2026-01-02 2026-01-02T05:00:00 3 +query TT +EXPLAIN SELECT k, a, c, ROW_NUMBER() OVER (ORDER BY c) AS r +FROM (SELECT lower(k) AS k, a, c FROM ordered_through_cast WHERE k = 'a' AND a = CAST(c AS DATE) ORDER BY k, a, c LIMIT 10) +ORDER BY r; +---- +physical_plan +01)ProjectionExec: expr=[k@0 as k, a@1 as a, c@2 as c, row_number() ORDER BY [ordered_through_cast.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as r] +02)--BoundedWindowAggExec: wdw=[row_number() ORDER BY [ordered_through_cast.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() ORDER BY [ordered_through_cast.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: TopK(fetch=10), expr=[a@1 ASC NULLS LAST, c@2 ASC NULLS LAST], preserve_partitioning=[false] +04)------ProjectionExec: expr=[lower(k@0) as k, a@1 as a, c@2 as c] +05)--------FilterExec: k@0 = a AND a@1 = CAST(c@2 AS Date32) +06)----------DataSourceExec: partitions=1, partition_sizes=[1] + # The same when a projection drops `a`: the implied `[c]` is stored, so the # projection can prove it and no sort is added above the projection. query TPI @@ -7960,9 +7988,6 @@ a 2026-01-01T09:00:00 1 a 2026-01-02T01:00:00 2 a 2026-01-02T05:00:00 3 -statement ok -set datafusion.explain.physical_plan_only = true; - query TT EXPLAIN SELECT k, c, ROW_NUMBER() OVER (ORDER BY c) AS r FROM (SELECT lower(k) AS k, a, c FROM ordered_through_cast WHERE k = 'a' AND a = CAST(c AS DATE) ORDER BY k, a, c LIMIT 10) @@ -7977,23 +8002,39 @@ physical_plan 06)----------FilterExec: k@0 = a AND a@1 = CAST(c@2 AS Date32) 07)------------DataSourceExec: partitions=1, partition_sizes=[1] -statement ok -reset datafusion.explain.physical_plan_only; - # And when a merge of partitions is added above the sort, which keeps only its -# merge key. +# merge key. `lower(k)` is computed, so the input is repartitioned and needs the +# merge. statement ok set datafusion.execution.target_partitions = 4; -query DPI -SELECT a, c, ROW_NUMBER() OVER (ORDER BY c) AS r -FROM (SELECT a, c FROM ordered_through_cast WHERE a = CAST(c AS DATE) ORDER BY a, c LIMIT 10) +query TDPI +SELECT l, a, c, ROW_NUMBER() OVER (ORDER BY c) AS r +FROM (SELECT lower(k) AS l, a, c FROM ordered_through_cast WHERE a = CAST(c AS DATE) ORDER BY a, c LIMIT 10) ORDER BY r; ---- -2026-01-01 2026-01-01T09:00:00 1 -2026-01-02 2026-01-02T01:00:00 2 -2026-01-02 2026-01-02T05:00:00 3 -2026-01-03 2026-01-03T01:00:00 4 +a 2026-01-01 2026-01-01T09:00:00 1 +a 2026-01-02 2026-01-02T01:00:00 2 +a 2026-01-02 2026-01-02T05:00:00 3 +b 2026-01-03 2026-01-03T01:00:00 4 + +query TT +EXPLAIN SELECT l, a, c, ROW_NUMBER() OVER (ORDER BY c) AS r +FROM (SELECT lower(k) AS l, a, c FROM ordered_through_cast WHERE a = CAST(c AS DATE) ORDER BY a, c LIMIT 10) +ORDER BY r; +---- +physical_plan +01)ProjectionExec: expr=[l@0 as l, a@1 as a, c@2 as c, row_number() ORDER BY [ordered_through_cast.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as r] +02)--BoundedWindowAggExec: wdw=[row_number() ORDER BY [ordered_through_cast.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() ORDER BY [ordered_through_cast.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortPreservingMergeExec: [a@1 ASC NULLS LAST, c@2 ASC NULLS LAST], fetch=10 +04)------SortExec: TopK(fetch=10), expr=[a@1 ASC NULLS LAST, c@2 ASC NULLS LAST], preserve_partitioning=[true] +05)--------ProjectionExec: expr=[lower(k@0) as l, a@1 as a, c@2 as c] +06)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +07)------------FilterExec: a@1 = CAST(c@2 AS Date32) +08)--------------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +reset datafusion.explain.physical_plan_only; statement ok set datafusion.execution.target_partitions = 1; From 79d9e8f922dbd096fa01e1ff6f34049e3acdafa4 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sun, 8 Nov 2026 09:43:51 +0800 Subject: [PATCH 08/19] fix: imply the ordering that leaves out a key after any prefix --- .../src/equivalence/properties/mod.rs | 31 ++++--- .../src/equivalence/properties/walk.rs | 82 +++++++++++++++++++ 2 files changed, 101 insertions(+), 12 deletions(-) diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 522a5d2171890..377c7719fd4c1 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -588,20 +588,27 @@ impl EquivalenceProperties { self.add_orderings(new_orderings) } - /// Returns what the normal ordering `ordering` implies, looking past the + /// Returns the orderings that the normal ordering `ordering` implies. If a + /// key is a function of the keys after it, as `a` is in `[k, a, c]` with + /// `a = CAST(c AS DATE)`, rows tied on the keys before it are sorted by the + /// keys after it, so the ordering without that key holds: `[k, c]`. The /// constants that a normal ordering can keep (those that differ between - /// partitions): within a partition they have one value. - fn implied_by( - &self, - ordering: &LexOrdering, - ) -> Result>> { - let mut rest = ordering.iter().filter(|sort_expr| { - self.eq_group.is_expr_constant(&sort_expr.expr).is_none() - }); - match rest.next() { - Some(leading) => self.implied_ordering(leading, rest), - None => Ok(None), + /// partitions) are left out, because within a partition they have one + /// value. + fn implied_by(&self, ordering: &LexOrdering) -> Result>> { + let keys: Vec<_> = ordering + .iter() + .filter(|sort_expr| self.eq_group.is_expr_constant(&sort_expr.expr).is_none()) + .collect(); + let mut implied = vec![]; + for (idx, key) in keys.iter().enumerate() { + if let Some(rest) = + self.implied_ordering(key, keys[idx + 1..].iter().copied())? + { + implied.push(keys[..idx].iter().copied().cloned().chain(rest).collect()); + } } + Ok(implied) } /// Returns the ordering that the normal ordering `[leading, rest..]` diff --git a/datafusion/physical-expr/src/equivalence/properties/walk.rs b/datafusion/physical-expr/src/equivalence/properties/walk.rs index 5e3f5f0b5f15f..0af56cdcb84ea 100644 --- a/datafusion/physical-expr/src/equivalence/properties/walk.rs +++ b/datafusion/physical-expr/src/equivalence/properties/walk.rs @@ -473,6 +473,88 @@ mod tests { Ok(()) } + /// A key that is a function of the keys after it implies the ordering + /// without it, after any prefix: with `CAST(c AS DATE) = a`, the ordering + /// `[k, a, c]` implies `[k, c]`, whichever is added first. The implied + /// ordering is stored, so it holds after `a` is gone. + #[test] + fn implied_ordering_after_a_prefix() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Utf8, true), + Field::new("a", DataType::Date32, true), + Field::new("c", DataType::Timestamp(TimeUnit::Nanosecond, None), true), + ])); + let [k, a, c] = ["k", "a", "c"].map(|name| col(name, &schema).unwrap()); + let cast_c = Arc::new(CastExpr::new_with_target_field( + Arc::clone(&c), + Arc::new(Field::new("c", DataType::Date32, true)), + None, + )) as Arc; + let ordering = + [&k, &a, &c].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))); + let implied = + [&k, &c].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))); + + let mut equality_first = EquivalenceProperties::new(Arc::clone(&schema)); + equality_first.add_equal_conditions(Arc::clone(&cast_c), Arc::clone(&a))?; + equality_first.add_ordering(ordering.clone())?; + + let mut ordering_first = EquivalenceProperties::new(Arc::clone(&schema)); + ordering_first.add_ordering(ordering)?; + ordering_first.add_equal_conditions(cast_c, Arc::clone(&a))?; + + for properties in [equality_first, ordering_first] { + assert!( + properties + .oeq_class() + .iter() + .any(|ordering| ordering.iter().eq(implied.iter())) + ); + } + Ok(()) + } + + /// A constant between a key and its arguments doesn't hide the equality: + /// with `k` a per-partition constant and `CAST(c AS DATE) = a`, the + /// ordering `[a, k, c]` implies `[c]`, because within a partition `k` has + /// one value. The constant is added first, so the normal ordering keeps it. + #[test] + fn implied_ordering_past_a_constant_between_keys() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, true), + Field::new("a", DataType::Date32, true), + Field::new("c", DataType::Timestamp(TimeUnit::Nanosecond, None), true), + ])); + let [k, a, c] = ["k", "a", "c"].map(|name| col(name, &schema).unwrap()); + let cast_c = Arc::new(CastExpr::new_with_target_field( + Arc::clone(&c), + Arc::new(Field::new("c", DataType::Date32, true)), + None, + )) as Arc; + let ordering = + [&a, &k, &c].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))); + let constant = || ConstExpr::new(Arc::clone(&k), AcrossPartitions::Heterogeneous); + + let mut equality_then_ordering = EquivalenceProperties::new(Arc::clone(&schema)); + equality_then_ordering.add_constants([constant()])?; + equality_then_ordering + .add_equal_conditions(Arc::clone(&cast_c), Arc::clone(&a))?; + equality_then_ordering.add_ordering(ordering.clone())?; + + let mut ordering_then_equality = EquivalenceProperties::new(Arc::clone(&schema)); + ordering_then_equality.add_constants([constant()])?; + ordering_then_equality.add_ordering(ordering)?; + ordering_then_equality.add_equal_conditions(cast_c, Arc::clone(&a))?; + + for properties in [equality_then_ordering, ordering_then_equality] { + assert!( + properties + .ordering_satisfy([PhysicalSortExpr::new_default(Arc::clone(&c))])? + ); + } + Ok(()) + } + /// `add_ordering` reports what it cannot derive instead of dropping it: /// here `c` is not in the schema, so its type is unknown. #[test] From 96a73995b3ecb1f35225b66b1256667ab49285b4 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sun, 8 Nov 2026 09:43:51 +0800 Subject: [PATCH 09/19] test: check the plans of the projection guard and implied ordering cases --- datafusion/sqllogictest/test_files/window.slt | 69 +++++++++++++++++-- 1 file changed, 64 insertions(+), 5 deletions(-) diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index c81a86324b725..49f244e94f831 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -7896,7 +7896,10 @@ statement ok set datafusion.execution.target_partitions = 1; # When a projection drops the leading constant, it can no longer prove the -# ordering that its input has, so the sort stays above the projection. +# ordering that its input has. The input satisfies `[total]`, because within a +# partition `k` has one value, so a sort pushed below the projection would be +# removed, and `SanityCheckPlan` would reject the window. The sort stays above +# the projection. query II SELECT total, DENSE_RANK() OVER (ORDER BY total DESC) AS r FROM ( @@ -7906,9 +7909,34 @@ FROM ( ---- 1 1 -# The same when the proof depends on any other dropped column, without a -# constant: `[k, a, c]` with `a = CAST(c AS DATE)` has `[k, c]`, but not after -# `a` is dropped. +statement ok +set datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT total, DENSE_RANK() OVER (ORDER BY total DESC) AS r +FROM ( + SELECT SUM(s) OVER (PARTITION BY k) AS total + FROM (SELECT lower(g) AS k, SUM(v) AS s FROM per_partition_constant WHERE g = 'a' GROUP BY lower(g)) +); +---- +physical_plan +01)ProjectionExec: expr=[total@0 as total, dense_rank() ORDER BY [total DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@1 as r] +02)--BoundedWindowAggExec: wdw=[dense_rank() ORDER BY [total DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "dense_rank() ORDER BY [total DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: expr=[total@0 DESC], preserve_partitioning=[false] +04)------ProjectionExec: expr=[sum(s) PARTITION BY [k] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@2 as total] +05)--------WindowAggExec: wdw=[sum(s) PARTITION BY [k] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "sum(s) PARTITION BY [k] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] +06)----------ProjectionExec: expr=[lower(per_partition_constant.g)@0 as k, sum(per_partition_constant.v)@1 as s] +07)------------AggregateExec: mode=Single, gby=[lower(g@0) as lower(per_partition_constant.g)], aggr=[sum(per_partition_constant.v)] +08)--------------FilterExec: g@0 = a +09)----------------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +reset datafusion.explain.physical_plan_only; + +# A key that is a function of the keys after it can be left out, after any +# prefix: `[k, a, c]` with `a = CAST(c AS DATE)` implies `[k, c]`. That ordering +# is stored, so a projection that drops `a` keeps it, and neither the window nor +# an `ORDER BY k, c` above the projection needs a sort. statement ok CREATE TABLE ordered_through_cast(k VARCHAR, a DATE, c TIMESTAMP) AS VALUES ('a', DATE '2026-01-02', TIMESTAMP '2026-01-02 05:00:00'), @@ -7926,6 +7954,36 @@ a 2026-01-02T01:00:00 2 a 2026-01-02T05:00:00 3 b 2026-01-03T01:00:00 4 +statement ok +set datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT k, c, ROW_NUMBER() OVER (ORDER BY k, c) AS r +FROM (SELECT k, a, c FROM ordered_through_cast WHERE a = CAST(c AS DATE) ORDER BY k, a, c LIMIT 10) +ORDER BY r; +---- +physical_plan +01)ProjectionExec: expr=[k@0 as k, c@1 as c, row_number() ORDER BY [ordered_through_cast.k ASC NULLS LAST, ordered_through_cast.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as r] +02)--BoundedWindowAggExec: wdw=[row_number() ORDER BY [ordered_through_cast.k ASC NULLS LAST, ordered_through_cast.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() ORDER BY [ordered_through_cast.k ASC NULLS LAST, ordered_through_cast.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----ProjectionExec: expr=[k@0 as k, c@2 as c] +04)------SortExec: TopK(fetch=10), expr=[k@0 ASC NULLS LAST, a@1 ASC NULLS LAST, c@2 ASC NULLS LAST], preserve_partitioning=[false] +05)--------FilterExec: a@1 = CAST(c@2 AS Date32) +06)----------DataSourceExec: partitions=1, partition_sizes=[1] + +query TT +EXPLAIN SELECT k, c +FROM (SELECT k, a, c FROM ordered_through_cast WHERE a = CAST(c AS DATE) ORDER BY k, a, c LIMIT 10) +ORDER BY k, c; +---- +physical_plan +01)ProjectionExec: expr=[k@0 as k, c@2 as c] +02)--SortExec: TopK(fetch=10), expr=[k@0 ASC NULLS LAST, a@1 ASC NULLS LAST, c@2 ASC NULLS LAST], preserve_partitioning=[false] +03)----FilterExec: a@1 = CAST(c@2 AS Date32) +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +reset datafusion.explain.physical_plan_only; + # An ordering that a sort adds implies other orderings through known equalities, # as it does when the equality comes after the sort: `[a, c]` with # `a = CAST(c AS DATE)` has `[c]`, so the window needs no sort of its own. @@ -7954,7 +8012,8 @@ physical_plan 04)------FilterExec: a@0 = CAST(c@1 AS Date32) 05)--------DataSourceExec: partitions=1, partition_sizes=[1] -# The same with a per-partition constant leading the sort. +# The same with a constant leading the sort. `lower(k)` under `k = 'a'` has the +# same value in every partition, so the TopK leaves it out of its sort keys. query TDPI SELECT k, a, c, ROW_NUMBER() OVER (ORDER BY c) AS r FROM (SELECT lower(k) AS k, a, c FROM ordered_through_cast WHERE k = 'a' AND a = CAST(c AS DATE) ORDER BY k, a, c LIMIT 10) From 434a1272bb6902381433d1ae45c6f27a91b33ed7 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Mon, 9 Nov 2026 20:56:49 +0800 Subject: [PATCH 10/19] refactor: skip an implied ordering that cannot be derived instead of returning an error --- datafusion/core/tests/execution/coop.rs | 3 +- .../tests/fuzz_cases/equivalence/ordering.rs | 5 +- .../tests/fuzz_cases/equivalence/utils.rs | 6 +-- .../sort_preserving_repartition_fuzz.rs | 2 +- .../physical_optimizer/ensure_requirements.rs | 3 +- datafusion/datasource-parquet/src/source.rs | 10 ++-- .../functions-table/src/generate_series.rs | 2 +- .../benches/equivalence_properties.rs | 24 ++++------ .../physical-expr/src/equivalence/mod.rs | 2 +- .../physical-expr/src/equivalence/ordering.rs | 4 +- .../src/equivalence/properties/dependency.rs | 44 +++++++++--------- .../src/equivalence/properties/joins.rs | 14 +++--- .../src/equivalence/properties/mod.rs | 28 +++++------ .../src/equivalence/properties/union.rs | 2 +- .../src/equivalence/properties/walk.rs | 46 ++++++++++--------- datafusion/physical-expr/src/projection.rs | 6 +-- .../physical-expr/src/window/standard.rs | 4 +- .../enforce_sorting/sort_pushdown.rs | 2 +- datafusion/physical-plan/src/memory.rs | 7 +-- .../physical-plan/src/repartition/mod.rs | 19 ++++---- .../src/sorts/partitioned_topk.rs | 3 +- datafusion/physical-plan/src/sorts/sort.rs | 8 ++-- .../src/sorts/sort_preserving_merge.rs | 22 ++++----- datafusion/physical-plan/src/union.rs | 2 +- datafusion/physical-plan/src/windows/mod.rs | 26 +++++------ .../library-user-guide/upgrading/56.0.0.md | 20 -------- 26 files changed, 137 insertions(+), 177 deletions(-) diff --git a/datafusion/core/tests/execution/coop.rs b/datafusion/core/tests/execution/coop.rs index 5349b8adf3997..eda649770ff5c 100644 --- a/datafusion/core/tests/execution/coop.rs +++ b/datafusion/core/tests/execution/coop.rs @@ -148,8 +148,7 @@ fn make_lazy_exec_with_range( exec.add_ordering(vec![PhysicalSortExpr::new( Arc::new(Column::new(column_name, 0)), SortOptions::new(false, true), - )]) - .unwrap(); + )]); exec } diff --git a/datafusion/core/tests/fuzz_cases/equivalence/ordering.rs b/datafusion/core/tests/fuzz_cases/equivalence/ordering.rs index 1684b02e4dad1..576d7e557dad3 100644 --- a/datafusion/core/tests/fuzz_cases/equivalence/ordering.rs +++ b/datafusion/core/tests/fuzz_cases/equivalence/ordering.rs @@ -347,7 +347,7 @@ fn test_ordering_satisfy_all_null_ordered_column() -> Result<()> { eq_properties.add_ordering([ PhysicalSortExpr::new(Arc::clone(col_a), options), PhysicalSortExpr::new(Arc::clone(col_b), options), - ])?; + ]); let batch = RecordBatch::try_new( Arc::clone(&schema), vec![ @@ -392,8 +392,7 @@ fn test_ordering_satisfy_null_constant_operand() -> Result<()> { as Arc; for options in ALL_SORT_OPTIONS { let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - eq_properties - .add_ordering([PhysicalSortExpr::new(col("a", &schema)?, options)])?; + eq_properties.add_ordering([PhysicalSortExpr::new(col("a", &schema)?, options)]); let batch = RecordBatch::try_new( Arc::clone(&schema), vec![sorted_f64_array(options, 2)?], diff --git a/datafusion/core/tests/fuzz_cases/equivalence/utils.rs b/datafusion/core/tests/fuzz_cases/equivalence/utils.rs index bff76e5351b2f..4e536ce9be7bb 100644 --- a/datafusion/core/tests/fuzz_cases/equivalence/utils.rs +++ b/datafusion/core/tests/fuzz_cases/equivalence/utils.rs @@ -119,7 +119,7 @@ pub fn create_random_schema( options: random_sort_options(&mut rng), }); - eq_properties.add_ordering(ordering)?; + eq_properties.add_ordering(ordering); } Ok((test_schema, eq_properties)) @@ -408,7 +408,7 @@ pub fn create_test_params() -> Result<(SchemaRef, EquivalenceProperties)> { ], ]; let orderings = convert_to_orderings(&orderings); - eq_properties.add_orderings(orderings)?; + eq_properties.add_orderings(orderings); Ok((test_schema, eq_properties)) } @@ -669,7 +669,7 @@ pub fn create_random_table_with_casts_and_windows( ALL_SORT_OPTIONS.map(|options| vec![sort(&k, opts_k), sort(&w, options)]), ); orderings.shuffle(&mut rng); - eq_properties.add_orderings(orderings)?; + eq_properties.add_orderings(orderings); let keys = vec![ (Arc::clone(&k), Some(opts_k)), diff --git a/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs b/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs index 690237e0b7ff6..312128750ded4 100644 --- a/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs @@ -90,7 +90,7 @@ mod sp_repartition_fuzz_tests { options: options_asc, }); - eq_properties.add_ordering(ordering)?; + eq_properties.add_ordering(ordering); } Ok((test_schema, eq_properties)) diff --git a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs index 49c42f8cecc8c..86adb307db75b 100644 --- a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs @@ -87,8 +87,7 @@ impl MockMultiPartitionExec { nulls_first: false, }, )]) { - eq.add_orderings(vec![ordering.into_iter().collect::>()]) - .unwrap(); + eq.add_orderings(vec![ordering.into_iter().collect::>()]); } let properties = PlanProperties::new( eq, diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 1e98f0dd2c47d..ac7e10f94e6f5 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -1014,7 +1014,7 @@ impl FileSource for ParquetSource { .collect::>() }) .collect::>(); - new.add_orderings(reversed_orderings)?; + new.add_orderings(reversed_orderings); new }; let reversed_satisfies = @@ -1712,8 +1712,7 @@ mod tests { let source = ParquetSource::new(Arc::clone(&schema)); // Source declares `[a DESC]`. let mut eq = EquivalenceProperties::new(Arc::clone(&schema)); - eq.add_ordering(vec![sort_expr_on(&schema, "a", true)]) - .unwrap(); + eq.add_ordering(vec![sort_expr_on(&schema, "a", true)]); // Request `[a ASC]`. let order = vec![sort_expr_on(&schema, "a", false)]; @@ -1786,7 +1785,7 @@ mod tests { }; let declared = request_expr.reverse(); let mut eq = EquivalenceProperties::new(Arc::clone(&full_schema)); - eq.add_ordering(vec![declared]).unwrap(); + eq.add_ordering(vec![declared]); let order = vec![request_expr]; let result = source.try_pushdown_sort(&order, &eq).unwrap(); @@ -1848,8 +1847,7 @@ mod tests { nulls_first: false, }, }, - ]) - .unwrap(); + ]); // Request `[a DESC, b ASC NULLS LAST, c DESC]` — three columns; // source's two-column declaration is a strict prefix. diff --git a/datafusion/functions-table/src/generate_series.rs b/datafusion/functions-table/src/generate_series.rs index fb53d537a10a6..7efd6d9e18f41 100644 --- a/datafusion/functions-table/src/generate_series.rs +++ b/datafusion/functions-table/src/generate_series.rs @@ -730,7 +730,7 @@ impl TableProvider for GenerateSeriesTable { .with_projection(projection.map(|p| p.to_vec())); if let Some(ordering) = self.output_ordering(exec.schema().as_ref()) { - exec.add_ordering([ordering])?; + exec.add_ordering([ordering]); } Ok(Arc::new(exec)) diff --git a/datafusion/physical-expr/benches/equivalence_properties.rs b/datafusion/physical-expr/benches/equivalence_properties.rs index 645de1d8cdbdb..0b14efcd3e24f 100644 --- a/datafusion/physical-expr/benches/equivalence_properties.rs +++ b/datafusion/physical-expr/benches/equivalence_properties.rs @@ -64,13 +64,11 @@ fn asc(i: usize) -> PhysicalSortExpr { fn properties(n_classes: usize) -> EquivalenceProperties { let schema = schema(8 + 2 * n_classes); let mut props = EquivalenceProperties::new(schema); - props - .add_orderings([ - vec![asc(0), asc(1), asc(2), asc(3)], - vec![asc(4), asc(5)], - vec![asc(6)], - ]) - .unwrap(); + props.add_orderings([ + vec![asc(0), asc(1), asc(2), asc(3)], + vec![asc(4), asc(5)], + vec![asc(6)], + ]); for i in 0..n_classes { props .add_equal_conditions(col(8 + 2 * i), col(9 + 2 * i)) @@ -227,7 +225,7 @@ fn bench_shapes(c: &mut Criterion) { ]); } } - props.add_orderings(orderings).unwrap(); + props.add_orderings(orderings); let keys = vec![ PhysicalSortExpr::new(col(0), asc_nulls_last), PhysicalSortExpr::new(col(1), asc_nulls_last), @@ -309,12 +307,10 @@ fn bench_project_wide_struct_fields(c: &mut Criterion) { let mut group = c.benchmark_group("equivalence_properties"); for (suffix, sort_col) in [("", col_a), ("/ordered_by_s", col_s)] { let mut props = EquivalenceProperties::new(Arc::clone(&schema)); - props - .add_orderings([vec![PhysicalSortExpr::new( - sort_col, - SortOptions::default(), - )]]) - .unwrap(); + props.add_orderings([vec![PhysicalSortExpr::new( + sort_col, + SortOptions::default(), + )]]); group.bench_function( format!("project/wide_struct_fields/{READ}_of_{FIELDS}{suffix}"), |b| b.iter(|| props.project(&mapping, Arc::clone(&output))), diff --git a/datafusion/physical-expr/src/equivalence/mod.rs b/datafusion/physical-expr/src/equivalence/mod.rs index 5d228ac2f9714..ff57146efcce9 100644 --- a/datafusion/physical-expr/src/equivalence/mod.rs +++ b/datafusion/physical-expr/src/equivalence/mod.rs @@ -137,7 +137,7 @@ mod tests { ], ]; let orderings = convert_to_orderings(&orderings); - eq_properties.add_orderings(orderings)?; + eq_properties.add_orderings(orderings); Ok((test_schema, eq_properties)) } diff --git a/datafusion/physical-expr/src/equivalence/ordering.rs b/datafusion/physical-expr/src/equivalence/ordering.rs index 1d938b29cf0c5..7de47767f62fd 100644 --- a/datafusion/physical-expr/src/equivalence/ordering.rs +++ b/datafusion/physical-expr/src/equivalence/ordering.rs @@ -662,7 +662,7 @@ mod tests { ); let mut eq_properties = EquivalenceProperties::new(Arc::clone(&test_schema)); let orderings = convert_to_orderings(&orderings); - eq_properties.add_orderings(orderings)?; + eq_properties.add_orderings(orderings); let classes = eq_group .into_iter() .map(|eq_class| EquivalenceClass::new(eq_class.into_iter().cloned())); @@ -706,7 +706,7 @@ mod tests { let orderings = convert_to_orderings(&orderings); // Column [a ASC], [e ASC], [d ASC, f ASC] are all valid orderings for the schema. - eq_properties.add_orderings(orderings)?; + eq_properties.add_orderings(orderings); // First entry in the tuple is required ordering, second entry is the expected flag // that indicates whether this required ordering is satisfied. diff --git a/datafusion/physical-expr/src/equivalence/properties/dependency.rs b/datafusion/physical-expr/src/equivalence/properties/dependency.rs index 836aa4247249d..b626e644a2a9b 100644 --- a/datafusion/physical-expr/src/equivalence/properties/dependency.rs +++ b/datafusion/physical-expr/src/equivalence/properties/dependency.rs @@ -460,7 +460,7 @@ mod tests { parse_sort_expr("b", &input_schema), parse_sort_expr("c", &input_schema), parse_sort_expr("d", &input_schema), - ])?; + ]); // add equivalent ordering [a, c, b, d] input_properties.add_ordering([ @@ -468,7 +468,7 @@ mod tests { parse_sort_expr("c", &input_schema), parse_sort_expr("b", &input_schema), // NB b and c are swapped parse_sort_expr("d", &input_schema), - ])?; + ]); // simply project all the columns in order let proj_exprs = vec![ @@ -504,13 +504,13 @@ mod tests { eq_properties.add_orderings([ vec![PhysicalSortExpr::new_default(Arc::clone(&col_b_expr))], vec![PhysicalSortExpr::new_default(Arc::clone(&col_c_expr))], - ])?; + ]); let mut expected_eqs = EquivalenceProperties::new(Arc::new(schema)); expected_eqs.add_orderings([ vec![PhysicalSortExpr::new_default(col_b_expr)], vec![PhysicalSortExpr::new_default(col_c_expr)], - ])?; + ]); assert!(eq_properties.oeq_class().eq(expected_eqs.oeq_class())); Ok(()) @@ -532,7 +532,7 @@ mod tests { eq_properties.add_ordering([ PhysicalSortExpr::new(Arc::new(Column::new("b", 1)), sort_options_not), PhysicalSortExpr::new(Arc::new(Column::new("a", 0)), sort_options), - ])?; + ]); let (result, idxs) = eq_properties.find_longest_permutation(&required_columns)?; assert_eq!(idxs, vec![0, 1]); assert_eq!( @@ -561,7 +561,7 @@ mod tests { PhysicalSortExpr::new(Arc::new(Column::new("b", 1)), sort_options_not), PhysicalSortExpr::new(Arc::new(Column::new("a", 0)), sort_options), ], - ])?; + ]); let (result, idxs) = eq_properties.find_longest_permutation(&required_columns)?; assert_eq!(idxs, vec![0, 1]); assert_eq!( @@ -588,7 +588,7 @@ mod tests { PhysicalSortExpr::new(Arc::new(Column::new("b", 1)), sort_options_not), PhysicalSortExpr::new(Arc::new(Column::new("c", 2)), sort_options), PhysicalSortExpr::new(Arc::new(Column::new("a", 0)), sort_options), - ])?; + ]); let (_, idxs) = eq_properties.find_longest_permutation(&required_columns)?; assert_eq!(idxs, vec![0]); @@ -621,7 +621,7 @@ mod tests { for known in all_options { let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); eq_properties - .add_ordering([PhysicalSortExpr::new(Arc::clone(&sqrt_a), known)])?; + .add_ordering([PhysicalSortExpr::new(Arc::clone(&sqrt_a), known)]); for required in all_options { let required_ordering = [PhysicalSortExpr::new(Arc::clone(&sqrt_a), required)]; @@ -638,7 +638,7 @@ mod tests { let required_ordering = [PhysicalSortExpr::new_default(Arc::clone(&sqrt_a))]; let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); assert!(!eq_properties.ordering_satisfy(required_ordering.clone())?); - eq_properties.add_ordering([PhysicalSortExpr::new_default(a)])?; + eq_properties.add_ordering([PhysicalSortExpr::new_default(a)]); assert!(!eq_properties.ordering_satisfy(required_ordering)?); Ok(()) } @@ -667,7 +667,7 @@ mod tests { eq_properties.add_orderings([ vec![PhysicalSortExpr::new(Arc::clone(&col_b), option_asc)], vec![PhysicalSortExpr::new(Arc::clone(&col_d), option_asc)], - ])?; + ]); let test_cases = vec![ // d + b can wrap @@ -736,7 +736,7 @@ mod tests { eq_properties.add_ordering([ PhysicalSortExpr::new(Arc::clone(col_d), option_asc), PhysicalSortExpr::new(Arc::clone(col_h), option_desc), - ])?; + ]); let test_cases = vec![ // TEST CASE 1 (vec![col_a], vec![(col_a, option_asc)]), @@ -1078,7 +1078,7 @@ mod tests { ] .into(); - eq_properties.add_ordering(initial_ordering.clone())?; + eq_properties.add_ordering(initial_ordering.clone()); // Add equality condition c = concat(a, b) eq_properties.add_equal_conditions(Arc::clone(&col_c), a_concat_b)?; @@ -1120,7 +1120,7 @@ mod tests { ] .into(); - eq_properties.add_ordering(initial_ordering.clone())?; + eq_properties.add_ordering(initial_ordering.clone()); // Add equality condition c = a * b eq_properties.add_equal_conditions(col_c, a_times_b)?; @@ -1210,7 +1210,7 @@ mod tests { PhysicalSortExpr::new(Arc::clone(&col_a), asc), PhysicalSortExpr::new(Arc::clone(&col_b), desc), PhysicalSortExpr::new(Arc::clone(&col_c), asc), - ])?; + ]); // New ordering: [a ASC] let new_order = vec![PhysicalSortExpr::new(Arc::clone(&col_a), asc)]; @@ -1248,7 +1248,7 @@ mod tests { eq_properties.add_ordering([ PhysicalSortExpr::new_default(Arc::clone(&col_a)), PhysicalSortExpr::new_default(Arc::clone(&col_c)), - ])?; + ]); // New ordering: [b ASC] let new_order = vec![PhysicalSortExpr::new_default(Arc::clone(&col_b))]; @@ -1289,7 +1289,7 @@ mod tests { eq_properties.add_ordering([ PhysicalSortExpr::new(Arc::clone(&col_a), asc), PhysicalSortExpr::new(Arc::clone(&col_b), desc), - ])?; + ]); // New ordering: [a DESC] let new_order = vec![PhysicalSortExpr::new(Arc::clone(&col_a), desc)]; @@ -1329,7 +1329,7 @@ mod tests { PhysicalSortExpr::new_default(Arc::clone(&col_a)), ], vec![PhysicalSortExpr::new_default(Arc::clone(&col_e))], - ])?; + ]); // New ordering: [b ASC, c ASC] let new_order = vec![ @@ -1478,7 +1478,7 @@ mod tests { expr: col(col_name, schema).unwrap(), options: SortOptions::default(), }); - eq_properties.add_ordering(base_ordering)?; + eq_properties.add_ordering(base_ordering); // Add constraints eq_properties = @@ -1564,7 +1564,7 @@ mod tests { nulls_first: false, }, ); - input_properties.add_orderings([vec![sort_a]])?; + input_properties.add_orderings([vec![sort_a]]); // Project through the named_struct let out_schema = output_schema(&projection_mapping, &input_schema)?; @@ -1728,7 +1728,7 @@ mod tests { input.add_orderings([vec![PhysicalSortExpr::new( Arc::clone(&a), SortOptions::default(), - )]])?; + )]]); // Several sources over each column, so all but the first of each are // served from the cache: `a` through the sort key, `b` as a plain @@ -1777,7 +1777,7 @@ mod tests { input.add_orderings([vec![PhysicalSortExpr::new( Arc::clone(&s), SortOptions::default(), - )]])?; + )]]); // Every probe refers to the sort key, so both substituting the // ordering and deriving the projected one examine each of them. @@ -1838,7 +1838,7 @@ mod tests { input.add_orderings([vec![PhysicalSortExpr::new( Arc::clone(&untyped), SortOptions::default(), - )]])?; + )]]); // The probe is ordered only through a derivation from `a + b`, which // fails before the probe is asked for its properties. The output type diff --git a/datafusion/physical-expr/src/equivalence/properties/joins.rs b/datafusion/physical-expr/src/equivalence/properties/joins.rs index 3c4175ee0f24b..2ba46f5d9b920 100644 --- a/datafusion/physical-expr/src/equivalence/properties/joins.rs +++ b/datafusion/physical-expr/src/equivalence/properties/joins.rs @@ -57,9 +57,9 @@ pub fn join_equivalence_properties( // then we should add `a ASC, b ASC` to the ordering equivalences // of the join output. let out_oeq_class = left_oeq_class.join_suffix(&right_oeq_class); - result.add_orderings(out_oeq_class)?; + result.add_orderings(out_oeq_class); } else { - result.add_orderings(left_oeq_class)?; + result.add_orderings(left_oeq_class); } } [false, true] => { @@ -82,9 +82,9 @@ pub fn join_equivalence_properties( // then we should add `b ASC, a ASC` to the ordering equivalences // of the join output. let out_oeq_class = right_oeq_class.join_suffix(&left_oeq_class); - result.add_orderings(out_oeq_class)?; + result.add_orderings(out_oeq_class); } else { - result.add_orderings(right_oeq_class)?; + result.add_orderings(right_oeq_class); } } [false, false] => {} @@ -185,8 +185,8 @@ mod tests { let left_orderings = convert_to_orderings(&left_orderings); let right_orderings = convert_to_orderings(&right_orderings); let expected = convert_to_orderings(&expected); - left_eq_properties.add_orderings(left_orderings)?; - right_eq_properties.add_orderings(right_orderings)?; + left_eq_properties.add_orderings(left_orderings); + right_eq_properties.add_orderings(right_orderings); let join_eq = join_equivalence_properties( left_eq_properties, right_eq_properties, @@ -261,7 +261,7 @@ mod tests { &join_type, left_columns_len, )?; - join_eq_properties.add_orderings(right_oeq_class)?; + join_eq_properties.add_orderings(right_oeq_class); let result = join_eq_properties.oeq_class().clone(); // [x ASC, y ASC], [z ASC, w ASC] diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 377c7719fd4c1..846695b7ea020 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -114,7 +114,7 @@ use itertools::{Either, Itertools}; /// eq_properties.add_ordering([ /// PhysicalSortExpr::new_default(col_a).asc(), /// PhysicalSortExpr::new_default(col_c).desc(), -/// ]).unwrap(); +/// ]); /// /// assert_eq!( /// eq_properties.to_string(), @@ -345,7 +345,7 @@ impl EquivalenceProperties { pub fn extend(mut self, other: Self) -> Result { self.constraints.extend(other.constraints); self.add_equivalence_group(other.eq_group)?; - self.add_orderings(other.oeq_class)?; + self.add_orderings(other.oeq_class); self.add_groupings(other.geq_class); Ok(self) } @@ -384,7 +384,7 @@ impl EquivalenceProperties { pub fn add_orderings( &mut self, orderings: impl IntoIterator>, - ) -> Result<()> { + ) { let orderings: Vec<_> = orderings.into_iter().filter_map(LexOrdering::new).collect(); let normal_orderings: Vec<_> = orderings @@ -393,14 +393,16 @@ impl EquivalenceProperties { .filter_map(|o| self.normalize_sort_exprs(o)) .collect(); if normal_orderings.is_empty() { - return Ok(()); + return; } // `add_constants` and `add_equal_conditions` discover implied orderings // too, and whether an ordering is satisfied must not depend on which of - // them ran last. + // them ran last. Deriving one fails only on an expression that does not + // match the schema; that ordering is then not implied, which can cost a + // sort but not a wrong result. let mut implied = vec![]; for ordering in &normal_orderings { - implied.extend(self.implied_by(ordering)?); + implied.extend(self.implied_by(ordering).unwrap_or_default()); } self.oeq_class.extend(orderings); // Normalize given orderings to update the cache: @@ -410,15 +412,12 @@ impl EquivalenceProperties { self.oeq_cache.update_map(); // Each implied ordering is shorter than the one implying it, so this // terminates: - self.add_orderings(implied) + self.add_orderings(implied); } /// Adds a single ordering to the existing ordering equivalence class. - pub fn add_ordering( - &mut self, - ordering: impl IntoIterator, - ) -> Result<()> { - self.add_orderings(std::iter::once(ordering)) + pub fn add_ordering(&mut self, ordering: impl IntoIterator) { + self.add_orderings(std::iter::once(ordering)); } /// Adds complete grouping tuples whose values are contiguous within @@ -585,7 +584,8 @@ impl EquivalenceProperties { let ordering = &self.oeq_cache.normal_cls[*idx]; new_orderings.extend(self.implied_by(ordering)?); } - self.add_orderings(new_orderings) + self.add_orderings(new_orderings); + Ok(()) } /// Returns the orderings that the normal ordering `ordering` implies. If a @@ -696,7 +696,7 @@ impl EquivalenceProperties { if normal_ordering.len() != self.common_sort_prefix_length(&normal_ordering)? { // If the ordering is unsatisfied, replace existing orderings: self.clear_orderings(); - self.add_ordering(ordering_tee)?; + self.add_ordering(ordering_tee); return Ok(true); } Ok(false) diff --git a/datafusion/physical-expr/src/equivalence/properties/union.rs b/datafusion/physical-expr/src/equivalence/properties/union.rs index 111456cbc3a38..a3d5c1c454a47 100644 --- a/datafusion/physical-expr/src/equivalence/properties/union.rs +++ b/datafusion/physical-expr/src/equivalence/properties/union.rs @@ -63,7 +63,7 @@ fn calculate_union_binary( let mut eq_properties = EquivalenceProperties::new(lhs.schema); eq_properties.add_constants(constants)?; - eq_properties.add_orderings(orderings)?; + eq_properties.add_orderings(orderings); Ok(eq_properties) } diff --git a/datafusion/physical-expr/src/equivalence/properties/walk.rs b/datafusion/physical-expr/src/equivalence/properties/walk.rs index 0af56cdcb84ea..05e9f5f6b584a 100644 --- a/datafusion/physical-expr/src/equivalence/properties/walk.rs +++ b/datafusion/physical-expr/src/equivalence/properties/walk.rs @@ -371,10 +371,10 @@ mod tests { let mut equality_first = EquivalenceProperties::new(Arc::clone(&schema)); equality_first.add_equal_conditions(Arc::clone(&cast_c), Arc::clone(&a))?; - equality_first.add_ordering(ordering.clone())?; + equality_first.add_ordering(ordering.clone()); let mut ordering_first = EquivalenceProperties::new(Arc::clone(&schema)); - ordering_first.add_ordering(ordering)?; + ordering_first.add_ordering(ordering); ordering_first.add_equal_conditions(cast_c, Arc::clone(&a))?; for properties in [equality_first, ordering_first] { @@ -412,19 +412,19 @@ mod tests { constant_equality_ordering.add_constants([constant()])?; constant_equality_ordering .add_equal_conditions(Arc::clone(&cast_c), Arc::clone(&a))?; - constant_equality_ordering.add_ordering(ordering.clone())?; + constant_equality_ordering.add_ordering(ordering.clone()); let mut constant_ordering_equality = EquivalenceProperties::new(Arc::clone(&schema)); constant_ordering_equality.add_constants([constant()])?; - constant_ordering_equality.add_ordering(ordering.clone())?; + constant_ordering_equality.add_ordering(ordering.clone()); constant_ordering_equality .add_equal_conditions(Arc::clone(&cast_c), Arc::clone(&a))?; let mut equality_ordering_constant = EquivalenceProperties::new(Arc::clone(&schema)); equality_ordering_constant.add_equal_conditions(cast_c, Arc::clone(&a))?; - equality_ordering_constant.add_ordering(ordering)?; + equality_ordering_constant.add_ordering(ordering); equality_ordering_constant.add_constants([constant()])?; for properties in [ @@ -464,7 +464,7 @@ mod tests { properties.add_equal_conditions(cast(&a, DataType::Date64), Arc::clone(&d))?; properties.add_ordering( [&d, &a, &c].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), - )?; + ); assert!( properties @@ -497,10 +497,10 @@ mod tests { let mut equality_first = EquivalenceProperties::new(Arc::clone(&schema)); equality_first.add_equal_conditions(Arc::clone(&cast_c), Arc::clone(&a))?; - equality_first.add_ordering(ordering.clone())?; + equality_first.add_ordering(ordering.clone()); let mut ordering_first = EquivalenceProperties::new(Arc::clone(&schema)); - ordering_first.add_ordering(ordering)?; + ordering_first.add_ordering(ordering); ordering_first.add_equal_conditions(cast_c, Arc::clone(&a))?; for properties in [equality_first, ordering_first] { @@ -539,11 +539,11 @@ mod tests { equality_then_ordering.add_constants([constant()])?; equality_then_ordering .add_equal_conditions(Arc::clone(&cast_c), Arc::clone(&a))?; - equality_then_ordering.add_ordering(ordering.clone())?; + equality_then_ordering.add_ordering(ordering.clone()); let mut ordering_then_equality = EquivalenceProperties::new(Arc::clone(&schema)); ordering_then_equality.add_constants([constant()])?; - ordering_then_equality.add_ordering(ordering)?; + ordering_then_equality.add_ordering(ordering); ordering_then_equality.add_equal_conditions(cast_c, Arc::clone(&a))?; for properties in [equality_then_ordering, ordering_then_equality] { @@ -555,10 +555,11 @@ mod tests { Ok(()) } - /// `add_ordering` reports what it cannot derive instead of dropping it: - /// here `c` is not in the schema, so its type is unknown. + /// `add_ordering` keeps the given ordering and leaves out an implied + /// ordering it cannot derive: here `c` is not in the schema, so its type + /// is unknown. #[test] - fn add_ordering_returns_implied_ordering_errors() -> Result<()> { + fn add_ordering_skips_underivable_implied_ordering() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Date32, true)])); let a = col("a", &schema)?; let c = Arc::new(Column::new("c", 1)) as Arc; @@ -572,7 +573,8 @@ mod tests { properties.add_equal_conditions(cast_c, Arc::clone(&a))?; let ordering = [&a, &c].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))); - assert!(properties.add_ordering(ordering).is_err()); + properties.add_ordering(ordering); + assert_eq!(properties.oeq_class().len(), 1); Ok(()) } @@ -599,7 +601,7 @@ mod tests { // Added after the constant, so that `k` stays in the normal ordering: properties.add_ordering( [&k, &b].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), - )?; + ); for required in [vec![&b], vec![&b, &k], vec![&k, &b]] { let required = required @@ -631,7 +633,7 @@ mod tests { let mut properties = EquivalenceProperties::new(Arc::clone(&schema)); properties.add_ordering( [&k, &a, &c].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), - )?; + ); properties.add_equal_conditions(cast_c, Arc::clone(&a))?; let required = @@ -668,7 +670,7 @@ mod tests { properties.add_equal_conditions(Arc::clone(&x), cast(&y, DataType::Int64))?; properties.add_ordering( [&k, &x, &y, &z].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), - )?; + ); let required = [&k, &z].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))); @@ -703,7 +705,7 @@ mod tests { vec![sort(&a, true), sort(&k, false), sort(&c, false)], vec![sort(&k, false), sort(&a, true), sort(&c, false)], vec![sort(&k, false), sort(&a, false)], - ])?; + ]); assert!(properties.ordering_satisfy([sort(&k, false), sort(&a, false)])?); Ok(()) @@ -739,7 +741,7 @@ mod tests { PhysicalSortExpr::new_default(Arc::clone(&k)), PhysicalSortExpr::new(Arc::clone(&x), options), ] - }))?; + })); for options in directions { let required = [ @@ -776,7 +778,7 @@ mod tests { properties.add_ordering([ PhysicalSortExpr::new_default(plus(&b)), PhysicalSortExpr::new_default(Arc::clone(&d)), - ])?; + ]); for right in [&b, &c] { let (_, indices) = @@ -809,7 +811,7 @@ mod tests { properties.add_ordering([ PhysicalSortExpr::new_default(Arc::clone(&k)), PhysicalSortExpr::new_default(Arc::clone(&b_mod_5)), - ])?; + ]); let (_, indices) = properties.find_longest_permutation(&[lit(5i32), k, b_mod_5])?; @@ -849,7 +851,7 @@ mod tests { }; let mut properties = EquivalenceProperties::new(Arc::clone(&schema)); - properties.add_ordering([PhysicalSortExpr::new_default(Arc::clone(&a))])?; + properties.add_ordering([PhysicalSortExpr::new_default(Arc::clone(&a))]); let satisfied = |exprs: [&Arc; 2]| { properties.ordering_satisfy( exprs.map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), diff --git a/datafusion/physical-expr/src/projection.rs b/datafusion/physical-expr/src/projection.rs index efe993be84650..79c2dc56adc00 100644 --- a/datafusion/physical-expr/src/projection.rs +++ b/datafusion/physical-expr/src/projection.rs @@ -1966,7 +1966,7 @@ pub(crate) mod tests { let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); let orderings = convert_to_orderings(&orderings); - eq_properties.add_orderings(orderings)?; + eq_properties.add_orderings(orderings); let proj_exprs = proj_exprs .into_iter() @@ -2169,7 +2169,7 @@ pub(crate) mod tests { let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); let orderings = convert_to_orderings(orderings); - eq_properties.add_orderings(orderings)?; + eq_properties.add_orderings(orderings); let expected = convert_to_orderings(expected); @@ -2296,7 +2296,7 @@ pub(crate) mod tests { } let orderings = convert_to_orderings(&orderings); - eq_properties.add_orderings(orderings)?; + eq_properties.add_orderings(orderings); let expected = convert_to_orderings(&expected); diff --git a/datafusion/physical-expr/src/window/standard.rs b/datafusion/physical-expr/src/window/standard.rs index 0b5fefaa50576..135b691615bd0 100644 --- a/datafusion/physical-expr/src/window/standard.rs +++ b/datafusion/physical-expr/src/window/standard.rs @@ -300,7 +300,7 @@ pub(crate) fn add_new_ordering_expr_with_partition_by( ) -> Result<()> { if partition_by.is_empty() { // In the absence of a PARTITION BY, ordering of `self.expr` is global: - eqp.add_ordering([expr])?; + eqp.add_ordering([expr]); } else { // If we have a PARTITION BY, standard functions can not introduce // a global ordering unless the existing ordering is compatible @@ -311,7 +311,7 @@ pub(crate) fn add_new_ordering_expr_with_partition_by( let (mut ordering, _) = eqp.find_longest_permutation(partition_by)?; if ordering.len() == partition_by.len() { ordering.push(expr); - eqp.add_ordering(ordering)?; + eqp.add_ordering(ordering); } } Ok(()) diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs index d7d5ee83e10c8..980a4b15dc334 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs @@ -1086,7 +1086,7 @@ fn is_custom_sort_pushdown_valid( // Clear old ordering metadata so the proposal must prove this on its own. let mut proposed_properties = child_properties.clone(); proposed_properties.clear_orderings(); - proposed_properties.add_ordering(ordering.clone())?; + proposed_properties.add_ordering(ordering.clone()); if !proposed_properties.ordering_satisfy(sort.expr().clone())? { return Ok(false); } diff --git a/datafusion/physical-plan/src/memory.rs b/datafusion/physical-plan/src/memory.rs index 71cf750367a8c..0688a87cb420e 100644 --- a/datafusion/physical-plan/src/memory.rs +++ b/datafusion/physical-plan/src/memory.rs @@ -257,13 +257,10 @@ impl LazyMemoryExec { Ok(()) } - pub fn add_ordering( - &mut self, - ordering: impl IntoIterator, - ) -> Result<()> { + pub fn add_ordering(&mut self, ordering: impl IntoIterator) { Arc::make_mut(&mut self.cache) .eq_properties - .add_orderings(std::iter::once(ordering)) + .add_orderings(std::iter::once(ordering)); } /// Get the batch generators diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 394d61909044a..13274d26e1fb4 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -2178,7 +2178,7 @@ impl RepartitionExec { partitioning: Partitioning, ) -> Result { let preserve_order = false; - let cache = Self::compute_properties(&input, partitioning, preserve_order)?; + let cache = Self::compute_properties(&input, partitioning, preserve_order); Ok(RepartitionExec { input, state: Default::default(), @@ -2200,19 +2200,19 @@ impl RepartitionExec { fn eq_properties_helper( input: &Arc, preserve_order: bool, - ) -> Result { + ) -> EquivalenceProperties { let mut eq_properties = input.equivalence_properties().clone(); // Fusing several input partitions keeps only the merge key ordering, // and per-partition constants no longer hold. if input.output_partitioning().partition_count() > 1 { eq_properties.clear_orderings(); if preserve_order && let Some(ordering) = input.output_ordering() { - eq_properties.add_ordering(ordering.clone())?; + eq_properties.add_ordering(ordering.clone()); } eq_properties.clear_groupings(); eq_properties.clear_per_partition_constants(); } - Ok(eq_properties) + eq_properties } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -2220,15 +2220,15 @@ impl RepartitionExec { input: &Arc, partitioning: Partitioning, preserve_order: bool, - ) -> Result { - Ok(PlanProperties::new( - Self::eq_properties_helper(input, preserve_order)?, + ) -> PlanProperties { + PlanProperties::new( + Self::eq_properties_helper(input, preserve_order), partitioning, input.pipeline_behavior(), input.boundedness(), ) .with_scheduling_type(SchedulingType::Cooperative) - .with_evaluation_type(EvaluationType::Eager)) + .with_evaluation_type(EvaluationType::Eager) } /// Specify if this repartitioning operation should preserve the order of @@ -2245,8 +2245,7 @@ impl RepartitionExec { // if there is only one input partition, merging is not required // to maintain order self.input.output_partitioning().partition_count() > 1; - let eq_properties = - Self::eq_properties_helper(&self.input, self.preserve_order).unwrap(); + let eq_properties = Self::eq_properties_helper(&self.input, self.preserve_order); Arc::make_mut(&mut self.cache).set_eq_properties(eq_properties); self } diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs index 3d0cee02a68d0..3950894b322f8 100644 --- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs +++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs @@ -366,8 +366,7 @@ impl PartitionedTopKExec { // value only ascends *within* a partition, so the ordering has to // be prefixed by the partition keys. `reorder` above already // established that prefix. - eq_properties - .add_ordering(partition_exprs.into_iter().chain([ranking_sort]))?; + eq_properties.add_ordering(partition_exprs.into_iter().chain([ranking_sort])); } Ok(PlanProperties::new( diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index bf90665dcb726..96e673e4043bc 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -2826,11 +2826,9 @@ mod tests { impl SortedUnboundedExec { fn compute_properties(schema: SchemaRef) -> PlanProperties { let mut eq_properties = EquivalenceProperties::new(schema); - eq_properties - .add_ordering([PhysicalSortExpr::new_default(Arc::new(Column::new( - "c1", 0, - )))]) - .unwrap(); + eq_properties.add_ordering([PhysicalSortExpr::new_default(Arc::new( + Column::new("c1", 0), + ))]); PlanProperties::new( eq_properties, Partitioning::UnknownPartitioning(1), diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 86175daa8f440..641e932c280c1 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -94,7 +94,7 @@ pub struct SortPreservingMergeExec { impl SortPreservingMergeExec { /// Create a new sort execution plan pub fn new(expr: LexOrdering, input: Arc) -> Self { - let cache = Self::compute_properties(&input, expr.clone()).unwrap(); + let cache = Self::compute_properties(&input, expr.clone()); Self { input, expr, @@ -148,7 +148,7 @@ impl SortPreservingMergeExec { fn compute_properties( input: &Arc, ordering: LexOrdering, - ) -> Result { + ) -> PlanProperties { let input_partitions = input.output_partitioning().partition_count(); let (drive, scheduling) = if input_partitions > 1 { (EvaluationType::Eager, SchedulingType::Cooperative) @@ -167,15 +167,15 @@ impl SortPreservingMergeExec { eq_properties.clear_groupings(); } eq_properties.clear_per_partition_constants(); - eq_properties.add_ordering(ordering)?; - Ok(PlanProperties::new( + eq_properties.add_ordering(ordering); + PlanProperties::new( eq_properties, // Equivalence Properties Partitioning::UnknownPartitioning(1), // Output Partitioning input.pipeline_behavior(), // Pipeline Behavior input.boundedness(), // Boundedness ) .with_evaluation_type(drive) - .with_scheduling_type(scheduling)) + .with_scheduling_type(scheduling) } } @@ -1890,13 +1890,11 @@ mod tests { .map(|(i, f)| Arc::new(Column::new(f.name(), i)) as Arc) .collect::>(); let mut eq_properties = EquivalenceProperties::new(schema); - eq_properties - .add_ordering( - columns - .iter() - .map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), - ) - .unwrap(); + eq_properties.add_ordering( + columns + .iter() + .map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), + ); PlanProperties::new( eq_properties, Partitioning::Hash(columns, 3), diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index dd4f7a76edf54..042f3217e22c5 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -1662,7 +1662,7 @@ mod tests { let child2 = Arc::new(TestMemoryExec::update_cache(&child2)); let mut union_expected_eq = EquivalenceProperties::new(Arc::clone(&schema)); - union_expected_eq.add_orderings(union_expected_orderings)?; + union_expected_eq.add_orderings(union_expected_orderings); let union: Arc = UnionExec::try_new(vec![child1, child2])?; let union_eq_properties = union.properties().equivalence_properties(); diff --git a/datafusion/physical-plan/src/windows/mod.rs b/datafusion/physical-plan/src/windows/mod.rs index 6954dec824869..1e5eea0b01285 100644 --- a/datafusion/physical-plan/src/windows/mod.rs +++ b/datafusion/physical-plan/src/windows/mod.rs @@ -454,7 +454,7 @@ pub(crate) fn window_equivalence_properties( existing }) }); - window_eq_properties.add_orderings(new_lexs)?; + window_eq_properties.add_orderings(new_lexs); } } else { // The window frame is ever expanding, so set monotonicity comes @@ -483,7 +483,7 @@ pub(crate) fn window_equivalence_properties( window_eq_properties.add_ordering([PhysicalSortExpr::new( Arc::new(window_col), SortOptions::new(increasing, true), - )])?; + )]); } else { // Reverse set-monotonic cases for all orderings: for mut lex in all_satisfied_lexs.into_iter() { @@ -491,7 +491,7 @@ pub(crate) fn window_equivalence_properties( Arc::new(window_col.clone()), SortOptions::new(increasing, true), )); - window_eq_properties.add_ordering(lex)?; + window_eq_properties.add_ordering(lex); } } } @@ -553,19 +553,15 @@ pub(crate) fn window_equivalence_properties( set_monotonicity.eq(&SetMonotonicity::Increasing); let window_col = Column::new(expr.name(), i + input_schema_len); if increasing && (asc || no_partitioning) { - window_eq_properties.add_ordering([ - PhysicalSortExpr::new( - Arc::new(window_col), - SortOptions::new(false, false), - ), - ])?; + window_eq_properties.add_ordering([PhysicalSortExpr::new( + Arc::new(window_col), + SortOptions::new(false, false), + )]); } else if !increasing && (!asc || no_partitioning) { - window_eq_properties.add_ordering([ - PhysicalSortExpr::new( - Arc::new(window_col), - SortOptions::new(true, false), - ), - ])?; + window_eq_properties.add_ordering([PhysicalSortExpr::new( + Arc::new(window_col), + SortOptions::new(true, false), + )]); } } } diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index de286462c9e0e..849f3237d82eb 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -1567,26 +1567,6 @@ A custom rule list that left out `SanityCheckPlan` is now checked too. A plan whose requirements do not hold, such as a `SortMergeJoinExec` over unsorted input, now fails to plan instead of possibly returning wrong results. -### `EquivalenceProperties::add_ordering` and `add_orderings` return `Result` - -`add_ordering` and `add_orderings` now also add the orderings that a new ordering -implies through known equalities. Deriving them can fail, so both now return -`Result<()>`, as does `LazyMemoryExec::add_ordering`. Code that ignores the result -gets an `unused_must_use` warning. - -`SortPreservingMergeExec::new` and `RepartitionExec::with_preserve_order` keep -their signatures and panic on such a failure, as `SortExec::new` already does. - -**Migration guide:** - -```rust,ignore -// Before -eq_properties.add_ordering(ordering); - -// After -eq_properties.add_ordering(ordering)?; -``` - ### `AggregateUDFImpl::is_nullable` is deprecated `AggregateUDFImpl::is_nullable` is deprecated; set the nullability in the field From 5f706697028b6e8a6df26348950e7cb777ee4559 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Tue, 10 Nov 2026 23:20:05 +0800 Subject: [PATCH 11/19] refactor: return LexOrdering from implied_by and document it with examples --- .../src/equivalence/properties/mod.rs | 57 ++++++++++++++++--- 1 file changed, 48 insertions(+), 9 deletions(-) diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 846695b7ea020..bc7fdde0eb2a1 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -588,14 +588,52 @@ impl EquivalenceProperties { Ok(()) } - /// Returns the orderings that the normal ordering `ordering` implies. If a - /// key is a function of the keys after it, as `a` is in `[k, a, c]` with - /// `a = CAST(c AS DATE)`, rows tied on the keys before it are sorted by the - /// keys after it, so the ordering without that key holds: `[k, c]`. The - /// constants that a normal ordering can keep (those that differ between - /// partitions) are left out, because within a partition they have one - /// value. - fn implied_by(&self, ordering: &LexOrdering) -> Result>> { + /// Returns the orderings that the normal ordering `ordering` implies + /// through known equalities: for each key that is a function of the keys + /// after it (see [`Self::implied_ordering`]), `ordering` without that key. + /// + /// # Example: a key that is a function of the keys after it + /// + /// * Ordering: `[k ASC, a ASC, c ASC]` + /// * Equality: `a = CAST(c AS DATE)` + /// * Returns: `[k ASC, c ASC]` + /// + /// Among the rows that tie on `k`, the rows are sorted by `a`, and the + /// rows that also tie on `a` are sorted by `c`. `a` is the date of `c`, so + /// a row with an earlier `a` also has an earlier `c`, and the rows that + /// tie on `k` are sorted by `c` alone: + /// + /// ```text + /// ┌───┬────────────┬──────────────────┐ + /// │ k │ a │ c │ + /// ├───┼────────────┼──────────────────┤ + /// │ x │ 2026-01-01 │ 2026-01-01 09:00 │ + /// │ x │ 2026-01-02 │ 2026-01-02 01:00 │ + /// │ x │ 2026-01-02 │ 2026-01-02 05:00 │ + /// │ y │ 2026-01-01 │ 2026-01-01 03:00 │ + /// └───┴────────────┴──────────────────┘ + /// ``` + /// + /// # Example: a constant that differs between partitions + /// + /// * Ordering: `[k ASC, a ASC, c ASC]`, where `k` is a constant whose value + /// differs between partitions + /// * Equality: `a = CAST(c AS DATE)` + /// * Returns: `[c ASC]` + /// + /// A normal ordering can keep such a constant (see + /// `EquivalenceGroup::normalize_sort_exprs`). Within a partition `k` has + /// one value, so it is left out of the keys and of the result. + /// + /// # Example: directions that differ + /// + /// * Ordering: `[a ASC, c DESC]` + /// * Equality: `a = CAST(c AS DATE)` + /// * Returns: nothing + /// + /// `c` descends within a date but ascends across dates, so the rows are + /// not sorted by `c` in either direction. + fn implied_by(&self, ordering: &LexOrdering) -> Result> { let keys: Vec<_> = ordering .iter() .filter(|sort_expr| self.eq_group.is_expr_constant(&sort_expr.expr).is_none()) @@ -605,7 +643,8 @@ impl EquivalenceProperties { if let Some(rest) = self.implied_ordering(key, keys[idx + 1..].iter().copied())? { - implied.push(keys[..idx].iter().copied().cloned().chain(rest).collect()); + let prefix = keys[..idx].iter().copied().cloned(); + implied.extend(LexOrdering::new(prefix.chain(rest))); } } Ok(implied) From a8ec6f44af3e05dad8a9f89b53cfea295a44f9ee Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Tue, 10 Nov 2026 23:20:06 +0800 Subject: [PATCH 12/19] docs: reword the column check in is_fixed and the projection guard test --- .../src/equivalence/properties/walk.rs | 4 +++- .../enforce_sorting/sort_pushdown.rs | 20 +++++++++++++++---- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/datafusion/physical-expr/src/equivalence/properties/walk.rs b/datafusion/physical-expr/src/equivalence/properties/walk.rs index 05e9f5f6b584a..64973451fb63c 100644 --- a/datafusion/physical-expr/src/equivalence/properties/walk.rs +++ b/datafusion/physical-expr/src/equivalence/properties/walk.rs @@ -296,7 +296,9 @@ impl<'a> SortKeyWalk<'a> { if class.is_some_and(|class| class.constant.is_some()) { return true; } - // A column is only fixed if it is pinned or constant: + // A column that is neither pinned nor constant can have a different + // value in each row, so it is not fixed. The leaf check below gives + // the same answer; returning here skips it: if expr.is::() { return false; } diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs index 980a4b15dc334..484bfa7066a0e 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs @@ -1543,10 +1543,22 @@ mod tests { assert!(matches!(remapped, OrderingRequirements::Soft(_))); } - /// The guard compares the alternative that is pushed down, not the first - /// one, which may be an alternative that cannot be remapped. Here the - /// input is sorted by `c`, and the projection proves `score` (`c` renamed), - /// but not `score % 2`, so the requirement can be pushed down. + /// The guard keeps a sort above a projection when the input already has + /// the ordering that would be pushed down, but the projection cannot show + /// that ordering to its parent. When the parent accepts several + /// orderings, the guard must check the one that is pushed down, not the + /// first one. + /// + /// * Input: sorted by `[c DESC]`; the projection outputs `c` as `score` + /// * Parent requires: `[score % 2 ASC]` or `[score DESC]` + /// * Expected: `[c DESC]` is pushed down + /// + /// `score % 2` is not a column, so `[score % 2 ASC]` cannot be rewritten + /// in terms of the input and is dropped; `[score DESC]` becomes + /// `[c DESC]`. The projection does provide `[score DESC]`, so the guard + /// does not apply. Checking the first alternative, `[score % 2 ASC]`, + /// which the projection does not provide, would keep an unneeded sort + /// above the projection. #[test] fn projection_guard_compares_the_pushed_alternative() -> Result<()> { let child = child_schema(); From 28a60df75026ea4ba57e3827841a05f537340503 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Tue, 10 Nov 2026 23:20:06 +0800 Subject: [PATCH 13/19] test: check the implied ordering for every pair of sort options --- .../src/equivalence/properties/walk.rs | 51 ++++++++++++------- 1 file changed, 33 insertions(+), 18 deletions(-) diff --git a/datafusion/physical-expr/src/equivalence/properties/walk.rs b/datafusion/physical-expr/src/equivalence/properties/walk.rs index 64973451fb63c..3d15f9ea880c6 100644 --- a/datafusion/physical-expr/src/equivalence/properties/walk.rs +++ b/datafusion/physical-expr/src/equivalence/properties/walk.rs @@ -354,8 +354,10 @@ mod tests { }; use datafusion_expr::{ColumnarValue, Operator, Volatility, create_udf}; - /// With `CAST(c AS DATE) = a`, the ordering `[a, c]` implies `[c]`, whether - /// the ordering or the equality comes first. + /// With `CAST(c AS DATE) = a`, the ordering `[a, c]` implies `[c]` when + /// `a` and `c` have the same sort options, whether the ordering or the + /// equality comes first. Otherwise it implies nothing: with + /// `[a DESC, c ASC]`, `c` descends across dates and ascends within one. #[test] fn implied_ordering_does_not_depend_on_call_order() -> Result<()> { let schema = Arc::new(Schema::new(vec![ @@ -368,22 +370,35 @@ mod tests { Arc::new(Field::new("c", DataType::Date32, true)), None, )) as Arc; - let ordering = - [&a, &c].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))); - - let mut equality_first = EquivalenceProperties::new(Arc::clone(&schema)); - equality_first.add_equal_conditions(Arc::clone(&cast_c), Arc::clone(&a))?; - equality_first.add_ordering(ordering.clone()); - - let mut ordering_first = EquivalenceProperties::new(Arc::clone(&schema)); - ordering_first.add_ordering(ordering); - ordering_first.add_equal_conditions(cast_c, Arc::clone(&a))?; - - for properties in [equality_first, ordering_first] { - assert!( - properties - .ordering_satisfy([PhysicalSortExpr::new_default(Arc::clone(&c))])? - ); + let all_options = [(false, false), (false, true), (true, false), (true, true)] + .map(|(descending, nulls_first)| SortOptions::new(descending, nulls_first)); + + for a_options in all_options { + for c_options in all_options { + let ordering = [ + PhysicalSortExpr::new(Arc::clone(&a), a_options), + PhysicalSortExpr::new(Arc::clone(&c), c_options), + ]; + + let mut equality_first = EquivalenceProperties::new(Arc::clone(&schema)); + equality_first + .add_equal_conditions(Arc::clone(&cast_c), Arc::clone(&a))?; + equality_first.add_ordering(ordering.clone()); + + let mut ordering_first = EquivalenceProperties::new(Arc::clone(&schema)); + ordering_first.add_ordering(ordering); + ordering_first + .add_equal_conditions(Arc::clone(&cast_c), Arc::clone(&a))?; + + let required = [PhysicalSortExpr::new(Arc::clone(&c), c_options)]; + for properties in [equality_first, ordering_first] { + assert_eq!( + properties.ordering_satisfy(required.clone())?, + a_options == c_options, + "ordering [a {a_options:?}, c {c_options:?}]" + ); + } + } } Ok(()) } From 2ceaba5e15cfc284173c3a1c84c66a04088a7d61 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Wed, 11 Nov 2026 21:15:27 +0800 Subject: [PATCH 14/19] fix: clear per-partition constants before RepartitionExec adds the merge ordering --- .../physical-plan/src/repartition/mod.rs | 105 +++++++++++++++++- 1 file changed, 103 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 13274d26e1fb4..95f26bd2b8020 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -2206,11 +2206,13 @@ impl RepartitionExec { // and per-partition constants no longer hold. if input.output_partitioning().partition_count() > 1 { eq_properties.clear_orderings(); + eq_properties.clear_groupings(); + // Clear the constants first: an ordering added while they hold + // implies orderings that hold only within each input partition. + eq_properties.clear_per_partition_constants(); if preserve_order && let Some(ordering) = input.output_ordering() { eq_properties.add_ordering(ordering.clone()); } - eq_properties.clear_groupings(); - eq_properties.clear_per_partition_constants(); } eq_properties } @@ -4942,6 +4944,105 @@ mod test { Ok(()) } + /// Merging partitions keeps only the merge key ordering. Here each input + /// partition has its own constant `k`, `a = CAST(c AS DATE)`, and is sorted + /// by `[k, a, c]`, so within a partition the rows are also sorted by `[c]`. + /// Partition `k = 0` has later dates than partition `k = 1`, so the merged + /// rows are not sorted by `[c]`. + #[tokio::test] + async fn test_preserve_order_drops_per_partition_implied_ordering() -> Result<()> { + use crate::filter::FilterExec; + use crate::sorts::sort::SortExec; + use arrow::array::{Date32Array, TimestampSecondArray}; + use arrow::datatypes::TimeUnit; + use datafusion_expr::Operator; + use datafusion_physical_expr::expressions::{BinaryExpr, CastExpr, lit}; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Date32, true), + Field::new("c", DataType::Timestamp(TimeUnit::Second, None), true), + ])); + let partition = |k: i32, days: [i32; 2]| -> Result> { + // One o'clock on each day: + let times = days.map(|day| i64::from(day) * 86_400 + 3_600); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Date32Array::from(days.to_vec())), + Arc::new(TimestampSecondArray::from(times.to_vec())), + ], + )?; + let source = + TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?; + Ok(Arc::new(ProjectionExec::try_new( + vec![ + (lit(ScalarValue::Int32(Some(k))), "k".to_string()), + (col("a", &schema)?, "a".to_string()), + (col("c", &schema)?, "c".to_string()), + ], + source, + )?)) + }; + let union = + UnionExec::try_new(vec![partition(0, [10, 11])?, partition(1, [1, 2])?])?; + let union_schema = union.schema(); + let cast_c = Arc::new(CastExpr::new_with_target_field( + col("c", &union_schema)?, + Arc::new(Field::new("c", DataType::Date32, true)), + None, + )); + let filter = Arc::new(FilterExec::try_new( + Arc::new(BinaryExpr::new( + col("a", &union_schema)?, + Operator::Eq, + cast_c, + )), + union, + )?); + let ordering = |names: &[&str]| { + names + .iter() + .map(|name| Ok(PhysicalSortExpr::new_default(col(name, &union_schema)?))) + .collect::>>() + }; + let sort = Arc::new( + SortExec::new( + LexOrdering::new(ordering(&["k", "a", "c"])?).unwrap(), + filter, + ) + .with_preserve_partitioning(true), + ); + assert!( + sort.properties() + .equivalence_properties() + .ordering_satisfy(ordering(&["c"])?)? + ); + + let exec = Arc::new( + RepartitionExec::try_new(sort, Partitioning::RoundRobinBatch(1))? + .with_preserve_order(), + ); + let properties = exec.properties().equivalence_properties(); + assert!(properties.ordering_satisfy(ordering(&["k", "a", "c"])?)?); + assert!(!properties.ordering_satisfy(ordering(&["c"])?)?); + + let batches = crate::collect(exec, Arc::new(TaskContext::default())).await?; + assert_batches_eq!( + [ + "+---+------------+---------------------+", + "| k | a | c |", + "+---+------------+---------------------+", + "| 0 | 1970-01-11 | 1970-01-11T01:00:00 |", + "| 0 | 1970-01-12 | 1970-01-12T01:00:00 |", + "| 1 | 1970-01-02 | 1970-01-02T01:00:00 |", + "| 1 | 1970-01-03 | 1970-01-03T01:00:00 |", + "+---+------------+---------------------+", + ], + &batches + ); + Ok(()) + } + #[tokio::test] async fn test_preserve_order_with_spilling() -> Result<()> { use datafusion_execution::runtime_env::RuntimeEnvBuilder; From 7f87bff2beb246c947174c6c11d3d44a1cb0717d Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Wed, 11 Nov 2026 21:15:28 +0800 Subject: [PATCH 15/19] fix: skip an implied ordering that cannot be derived when discovering orderings too --- .../src/equivalence/properties/mod.rs | 4 +++- .../src/equivalence/properties/walk.rs | 24 ++++++++++++------- 2 files changed, 19 insertions(+), 9 deletions(-) diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index bc7fdde0eb2a1..ce0d33139e233 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -582,7 +582,9 @@ impl EquivalenceProperties { let mut new_orderings = vec![]; for idx in ordering_idxs { let ordering = &self.oeq_cache.normal_cls[*idx]; - new_orderings.extend(self.implied_by(ordering)?); + // As in `add_orderings`, an ordering that cannot be derived is + // not implied, so the result does not depend on the call order: + new_orderings.extend(self.implied_by(ordering).unwrap_or_default()); } self.add_orderings(new_orderings); Ok(()) diff --git a/datafusion/physical-expr/src/equivalence/properties/walk.rs b/datafusion/physical-expr/src/equivalence/properties/walk.rs index 3d15f9ea880c6..36c69f5e3eff9 100644 --- a/datafusion/physical-expr/src/equivalence/properties/walk.rs +++ b/datafusion/physical-expr/src/equivalence/properties/walk.rs @@ -572,11 +572,11 @@ mod tests { Ok(()) } - /// `add_ordering` keeps the given ordering and leaves out an implied - /// ordering it cannot derive: here `c` is not in the schema, so its type - /// is unknown. + /// An implied ordering that cannot be derived is left out, and the given + /// ordering is kept, whether the ordering or the equality comes first: + /// here `c` is not in the schema, so its type is unknown. #[test] - fn add_ordering_skips_underivable_implied_ordering() -> Result<()> { + fn underivable_implied_ordering_is_skipped_in_either_order() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Date32, true)])); let a = col("a", &schema)?; let c = Arc::new(Column::new("c", 1)) as Arc; @@ -586,12 +586,20 @@ mod tests { None, )) as Arc; - let mut properties = EquivalenceProperties::new(Arc::clone(&schema)); - properties.add_equal_conditions(cast_c, Arc::clone(&a))?; let ordering = [&a, &c].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))); - properties.add_ordering(ordering); - assert_eq!(properties.oeq_class().len(), 1); + + let mut equality_first = EquivalenceProperties::new(Arc::clone(&schema)); + equality_first.add_equal_conditions(Arc::clone(&cast_c), Arc::clone(&a))?; + equality_first.add_ordering(ordering.clone()); + + let mut ordering_first = EquivalenceProperties::new(Arc::clone(&schema)); + ordering_first.add_ordering(ordering); + ordering_first.add_equal_conditions(cast_c, Arc::clone(&a))?; + + for properties in [equality_first, ordering_first] { + assert_eq!(properties.oeq_class().len(), 1); + } Ok(()) } From 50456087ddd5ac5292fd8a9ae1fa424bb0b9244e Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Wed, 11 Nov 2026 21:15:28 +0800 Subject: [PATCH 16/19] perf: expand each implied ordering once in add_orderings --- .../physical-expr/src/equivalence/properties/mod.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index ce0d33139e233..c39f603bbe59b 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -400,9 +400,14 @@ impl EquivalenceProperties { // them ran last. Deriving one fails only on an expression that does not // match the schema; that ordering is then not implied, which can cost a // sort but not a wrong result. - let mut implied = vec![]; + let mut implied: Vec = vec![]; for ordering in &normal_orderings { - implied.extend(self.implied_by(ordering).unwrap_or_default()); + // Two orderings can imply the same one; expand it only once: + for ordering in self.implied_by(ordering).unwrap_or_default() { + if !implied.contains(&ordering) { + implied.push(ordering); + } + } } self.oeq_class.extend(orderings); // Normalize given orderings to update the cache: From 27ed6c621dcbdbb6d564f57397ad1561b90e750d Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Wed, 11 Nov 2026 21:15:28 +0800 Subject: [PATCH 17/19] refactor: find ordering heads with one helper in SortKeyWalk --- .../src/equivalence/properties/walk.rs | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/datafusion/physical-expr/src/equivalence/properties/walk.rs b/datafusion/physical-expr/src/equivalence/properties/walk.rs index 36c69f5e3eff9..0eb27ad0020bc 100644 --- a/datafusion/physical-expr/src/equivalence/properties/walk.rs +++ b/datafusion/physical-expr/src/equivalence/properties/walk.rs @@ -67,12 +67,7 @@ impl<'a> SortKeyWalk<'a> { { walk.head_idxs = normal .iter() - .map(|ordering| { - ordering - .iter() - .position(|sort_expr| !walk.is_fixed(&sort_expr.expr)) - .unwrap_or(ordering.len()) - }) + .map(|ordering| walk.skip_fixed(ordering, 0)) .collect(); } walk @@ -173,11 +168,7 @@ impl<'a> SortKeyWalk<'a> { let mut found = vec![]; for idx in first..count { let ordering = self.ordering(idx); - let mut head_idx = self.head_idxs[idx]; - while head_idx < ordering.len() && self.is_fixed(&ordering[head_idx].expr) - { - head_idx += 1; - } + let head_idx = self.skip_fixed(ordering, self.head_idxs[idx]); if let Some(head) = ordering.get(head_idx) { let rest = ordering[head_idx + 1..] .iter() @@ -284,6 +275,16 @@ impl<'a> SortKeyWalk<'a> { Ok(Transformed::yes(node)) } + /// Returns the index of the head of `ordering` (see [`SortKeyWalk`]), + /// searching from `start`, or `ordering.len()` if it has no head there. + fn skip_fixed(&self, ordering: &LexOrdering, start: usize) -> usize { + let mut idx = start; + while idx < ordering.len() && self.is_fixed(&ordering[idx].expr) { + idx += 1; + } + idx + } + /// Whether `expr` is fixed (see [`SortKeyWalk`]). Callers pass `expr` in /// normal form. Pins are stored as given, so a pinned expression that is /// not in normal form is found only through its equivalence class, whose From f2e397e9027984cd43c98628098d850115b4d7de Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Wed, 11 Nov 2026 21:15:28 +0800 Subject: [PATCH 18/19] refactor: take the pushed alternative and its remap from one pass in the projection guard --- .../enforce_sorting/sort_pushdown.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs index 484bfa7066a0e..bc88a529f38a6 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs @@ -1329,19 +1329,19 @@ fn handle_projection_pushdown( // such as a leading per-partition constant. Keep the sort above the // projection in that case. Alternatives that cannot be remapped are not // pushed down, so compare against the first one that can. - let pushed = parent_required.iter().find(|required| { + let pushed = parent_required.iter().find_map(|required| { remap_lex_requirement_through_projection(projection_exec.expr(), required) - .is_some() + .map(|remapped| (required, remapped)) }); - if let Some(pushed) = pushed + if let Some((required, remapped)) = pushed && projection_exec .input() .equivalence_properties() - .ordering_satisfy_requirement(requirements.first().clone())? + .ordering_satisfy_requirement(remapped)? && !projection_exec .properties() .equivalence_properties() - .ordering_satisfy_requirement(pushed.clone())? + .ordering_satisfy_requirement(required.clone())? { return Ok(None); } From 878968001f32b09cd01fee27904e52c813a39f30 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Wed, 11 Nov 2026 21:15:28 +0800 Subject: [PATCH 19/19] test: check the plan of the per-partition constant case at 4 target partitions --- datafusion/sqllogictest/test_files/window.slt | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 49f244e94f831..2a9e3084c1084 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -7892,6 +7892,28 @@ FROM ( ---- 3 2 2 1 +statement ok +set datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT k, c, total, DENSE_RANK() OVER (ORDER BY total DESC, k) AS r +FROM ( + SELECT k, c, SUM(c) OVER (PARTITION BY k) AS total + FROM (SELECT SUM(v) AS k, COUNT(*) AS c FROM per_partition_constant) +); +---- +physical_plan +01)ProjectionExec: expr=[k@0 as k, c@1 as c, total@2 as total, dense_rank() ORDER BY [total DESC NULLS FIRST, k ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as r] +02)--BoundedWindowAggExec: wdw=[dense_rank() ORDER BY [total DESC NULLS FIRST, k ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "dense_rank() ORDER BY [total DESC NULLS FIRST, k ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----ProjectionExec: expr=[k@0 as k, c@1 as c, sum(c) PARTITION BY [k] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@2 as total] +04)------WindowAggExec: wdw=[sum(c) PARTITION BY [k] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "sum(c) PARTITION BY [k] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] +05)--------ProjectionExec: expr=[sum(per_partition_constant.v)@0 as k, count(Int64(1))@1 as c] +06)----------AggregateExec: mode=Single, gby=[], aggr=[sum(per_partition_constant.v), count(Int64(1))] +07)------------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +reset datafusion.explain.physical_plan_only; + statement ok set datafusion.execution.target_partitions = 1;