Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
b42a16b
fix: skip a leading per-partition constant when finding ordering heads
alamb Nov 5, 2026
6e6a8ae
fix: keep the sort above a projection that cannot prove the ordering
jayzhan211 Nov 7, 2026
c777864
fix: discover implied orderings when an ordering is added
jayzhan211 Nov 7, 2026
6425ba0
refactor: return errors from EquivalenceProperties::add_orderings
jayzhan211 Nov 7, 2026
9c7ee55
fix: look past leading constants when discovering implied orderings
jayzhan211 Nov 7, 2026
9210b51
fix: compare the pushed alternative in the projection sort guard
jayzhan211 Nov 7, 2026
17c73e6
test: check the plans of the implied ordering window cases
jayzhan211 Nov 7, 2026
79d9e8f
fix: imply the ordering that leaves out a key after any prefix
jayzhan211 Nov 8, 2026
96a7399
test: check the plans of the projection guard and implied ordering cases
jayzhan211 Nov 8, 2026
434a127
refactor: skip an implied ordering that cannot be derived instead of …
jayzhan211 Nov 9, 2026
5f70669
refactor: return LexOrdering from implied_by and document it with exa…
jayzhan211 Nov 10, 2026
a8ec6f4
docs: reword the column check in is_fixed and the projection guard test
jayzhan211 Nov 10, 2026
28a60df
test: check the implied ordering for every pair of sort options
jayzhan211 Nov 10, 2026
2ceaba5
fix: clear per-partition constants before RepartitionExec adds the me…
jayzhan211 Nov 11, 2026
7f87bff
fix: skip an implied ordering that cannot be derived when discovering…
jayzhan211 Nov 11, 2026
5045608
perf: expand each implied ordering once in add_orderings
jayzhan211 Nov 11, 2026
27ed6c6
refactor: find ordering heads with one helper in SortKeyWalk
jayzhan211 Nov 11, 2026
f2e397e
refactor: take the pushed alternative and its remap from one pass in …
jayzhan211 Nov 11, 2026
8789680
test: check the plan of the per-partition constant case at 4 target p…
jayzhan211 Nov 11, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 93 additions & 13 deletions datafusion/physical-expr/src/equivalence/properties/mod.rs

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

discover_new_orderings now duplicates work: it computes implied_ordering for each leading ordering and passes the results to add_orderings, which immediately recomputes implied_ordering on those results and recurses again, each level re-normalizing and calling update_map.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These aren't the same computation: discover_new_orderings derives from an ordering, and add_orderings then derives from what that ordering implied. The recursion is what finds chains. With d = CAST(a), a = CAST(c) and ordering [d, a, c], [c] is found with it and missed with one level, and implied_orderings_chain covers that. It only runs when something was implied.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The walk does not derive it because nothing is pinned. This is the order dependence the new comment says must not happen.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I couldn't get an order dependence from new_with_orderings. It starts with no equalities, and add_equal_conditions and add_constants both re-derive the implied orderings of every stored ordering. With k a per-partition constant and a = CAST(c AS DATE), ordering_satisfy([c]) is true for [a, c] and for [k, a, c], whichever of the equality and the constant is added first, and the same as with add_ordering. Which sequence of calls do you have in mind?

Original file line number Diff line number Diff line change
Expand Up @@ -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(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Now that add_orderings derives implied orderings, a caller that adds an ordering before clearing per-partition constants keeps orderings that no longer hold. RepartitionExec::eq_properties_helper does this: it calls add_ordering(merge key) and then clear_per_partition_constants(). SortPreservingMergeExec::compute_properties clears first.

With k a Heterogeneous constant, CAST(c AS DATE) = a, and merge key [k, a, c], the repartition order leaves ordering_satisfy([c]) true, although [c] doesn't hold once partitions with different k are merged. The SPM order gives false. Could eq_properties_helper clear the constants before add_ordering, the way SPM does? I don't see a SQL path to this today, because UnionExec drops equalities, so it mainly affects custom plans and any future source of per-partition constants.

&mut self,
orderings: impl IntoIterator<Item = impl IntoIterator<Item = PhysicalSortExpr>>,
Expand All @@ -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<LexOrdering> = 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);
Comment on lines +418 to +420

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we use a deduplicated worklist and avoid expanding already-known candidates?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Each round's implied orderings are now deduplicated before they are expanded. I didn't skip candidates that are already in the class: remove_redundant_entries drops the repeats anyway, and skipping them would rely on every earlier call having derived their implied orderings.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

add_orderings passes implied back to itself without removing duplicates. Each level also runs remove_redundant_entries (O(n²)) and update_map over the whole class, so the work can grow with the number of key subsets.

}

/// Adds a single ordering to the existing ordering equivalence class.
Expand Down Expand Up @@ -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<Vec<LexOrdering>> {
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())?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it seems like the fact that implied_ordering returns a Result here requires many signature changes (and the upgrade guide entry) which significantly complicates the PR.

Since the only error that is actually possible is if the type can't be found from the schema, which seems like a safe thing to ignore (.ok(), despite my previous comment to the contrary 🤦 )

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could switch back to ok(). I prefer that, but I'm not strongly against returning an Error either

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I also vote to return ok() and leave a comment explaining why the error is ignored

This will make it less disruptive to backport

Then we can make a PR to main perhaps to propagate the error

{
let prefix = keys[..idx].iter().copied().cloned();

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

IMO the Simpler fix would be: keep the constants in the stored implied ordering (store [k, c], not [c]). SortKeyWalk::new already skips a leading fixed head, so [c] still holds within a partition, and clearing k then removes it correctly.

wdyt?

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This handles a leading constant, but not one between a key and its arguments, as in implied_ordering_past_a_constant_between_keys: [a, k, c] with a = CAST(c AS DATE) would store [k, c]. That holds within each partition, but not after RepartitionExec merges on [a, k, c]. Rows from two partitions on the same date come out as k = 1, 2, and the next date starts again at k = 1. Since the merge adds the ordering before clearing k, [k, c] would still be kept.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't understand the implications of this comment

Specifically, if we didn't do this change, would we get wrong plans / errors? Or is it just a missed optimization opportunity?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd keep [c] in this PR. As viirya says, the RepartitionExec call order needs fixing either way, and that fix is now in the PR. With it, SortPreservingMergeExec, RepartitionExec and CoalescePartitionsExec all drop [c] when they merge partitions.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the issue is not critical but the change is quite small

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It was a wrong-result risk, not a missed optimization: before the reorder, the merged output could claim [c] although it doesn't hold, so a sort above it could be removed and rows come out in the wrong order. We don't know a SQL query that reaches it, which is why it wasn't critical. 2ceaba5 fixes it, and test_preserve_order_drops_per_partition_implied_ordering covers it.

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.
///
Expand Down
Loading
Loading