diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 1206315bbb274..c39f603bbe59b 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -379,7 +379,8 @@ 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>, @@ -391,14 +392,32 @@ impl EquivalenceProperties { .cloned() .filter_map(|o| self.normalize_sort_exprs(o)) .collect(); - if !normal_orderings.is_empty() { - 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(); + if normal_orderings.is_empty() { + 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. 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 = vec![]; + for ordering in &normal_orderings { + // 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: + 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. @@ -568,15 +587,76 @@ 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())?); - } - if !new_orderings.is_empty() { - self.add_orderings(new_orderings); + // 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(()) } + /// 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()) + .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())? + { + let prefix = keys[..idx].iter().copied().cloned(); + implied.extend(LexOrdering::new(prefix.chain(rest))); + } + } + Ok(implied) + } + /// 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 e0487cdd23dfa..0eb27ad0020bc 100644 --- a/datafusion/physical-expr/src/equivalence/properties/walk.rs +++ b/datafusion/physical-expr/src/equivalence/properties/walk.rs @@ -43,19 +43,34 @@ 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, } 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| walk.skip_fixed(ordering, 0)) + .collect(); } + walk } /// Returns how many leading keys (in normal form) the existing orderings @@ -153,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() @@ -264,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 @@ -276,6 +297,12 @@ impl<'a> SortKeyWalk<'a> { if class.is_some_and(|class| class.constant.is_some()) { return true; } + // 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; + } // A volatile expression can change from row to row, even when all of // its inputs are fixed: if expr.is_volatile_node() { @@ -314,7 +341,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}; @@ -328,6 +355,289 @@ mod tests { }; use datafusion_expr::{ColumnarValue, Operator, Volatility, create_udf}; + /// 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![ + 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 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(()) + } + + /// 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(()) + } + + /// 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(()) + } + + /// 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 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; + 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_eq!(properties.oeq_class().len(), 1); + } + 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 + /// 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/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs index f2a01598b982e..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 @@ -1318,10 +1318,34 @@ 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. Alternatives that cannot be remapped are not + // pushed down, so compare against the first one that can. + let pushed = parent_required.iter().find_map(|required| { + remap_lex_requirement_through_projection(projection_exec.expr(), required) + .map(|remapped| (required, remapped)) + }); + if let Some((required, remapped)) = pushed + && projection_exec + .input() + .equivalence_properties() + .ordering_satisfy_requirement(remapped)? + && !projection_exec + .properties() + .equivalence_properties() + .ordering_satisfy_requirement(required.clone())? + { + return Ok(None); + } + Ok(Some(vec![Some(requirements)])) } #[cfg(test)] @@ -1335,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::{ @@ -1519,6 +1543,58 @@ mod tests { assert!(matches!(remapped, OrderingRequirements::Soft(_))); } + /// 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(); + 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(); 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; diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 8508ca6ff66bb..2a9e3084c1084 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -7834,3 +7834,299 @@ 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; + +# 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.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; + +# When a projection drops the leading constant, it can no longer prove the +# 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 ( + 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 + +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'), + ('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 +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. +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.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 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) +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 + +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 +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 + +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] + +# And when a merge of partitions is added above the sort, which keeps only its +# merge key. `lower(k)` is computed, so the input is repartitioned and needs the +# merge. +statement ok +set datafusion.execution.target_partitions = 4; + +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; +---- +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; + +statement ok +DROP TABLE ordered_through_cast; + +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;