From 85a486d8fd48d3ac60f7abd0ee2b697d11fb90da Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 9 Nov 2026 09:49:03 -0500 Subject: [PATCH 1/3] docs: explain SortKeyWalk with worked examples Co-Authored-By: Claude Fable 5.1 --- .../src/equivalence/properties/walk.rs | 169 ++++++++++++++---- 1 file changed, 138 insertions(+), 31 deletions(-) diff --git a/datafusion/physical-expr/src/equivalence/properties/walk.rs b/datafusion/physical-expr/src/equivalence/properties/walk.rs index e0487cdd23dfa..858de0463be93 100644 --- a/datafusion/physical-expr/src/equivalence/properties/walk.rs +++ b/datafusion/physical-expr/src/equivalence/properties/walk.rs @@ -22,29 +22,127 @@ use itertools::Itertools; /// The state of one walk along a sequence of sort keys, such as the keys of a /// required ordering. /// -/// A sort key only matters among the rows that tie on the keys before it. So -/// once a key is satisfied, the walk *pins* it: for the rest of the walk, the -/// key has one value, like a constant. An expression is *fixed* if it has one -/// value among those rows: it is a literal, a constant, pinned, a leaf that -/// reports itself as a singleton (such as an uncorrelated scalar subquery), or -/// not volatile and computed only from fixed expressions. The *head* of an -/// ordering is its first expression that is not fixed; an ordering whose -/// expressions are all fixed has no head. With orderings `[a, b]` and -/// `[a, c]`, the requirement `[a, c]` is satisfied, because pinning `a` makes -/// `c` the head of the second ordering. +/// [`EquivalenceProperties`] uses a `SortKeyWalk` to decide if a set of +/// orderings satisfy a requirement. See [Checking whether a requirement is +/// satisfied] for more details. /// -/// Pins live in the walk; the equivalence properties it reads do not change. +/// # Intuition and Introduction, and Definitions +/// +/// An ordering `[x, y, z]` is a chain of claims about the rows: they are sorted +/// by x; among rows with equal x, sorted by y; among rows with equal x and y, +/// sorted by z. +/// +/// ## Pinning and Fixed Expressions +/// +/// During a walk, the code scans the ordering requirement left to right. When +/// it finds key `k` satisfied, it *pins* `k`, meaning that for the next steps +/// in the walk, `k` is effectively a constant across the remaining ordering +/// requirement. +/// +/// An expression is *fixed* if it has one value for all rows that tie on the +/// currently pinned expressions. Fixed expressions include the pinned expressions +/// themselves, as well as literals, constants, and singleton leaf expressions +/// (such as an uncorrelated scalar subquery). See [`Self::is_fixed`] for more +/// details. +/// +/// ## Heads +/// +/// For each equivalent ordering, the walk tracks its current *head* expression, +/// which is the current claim that is valid given the currently pinned +/// expressions. When an expression is pinned, it advances the head of any +/// existing claim about that key. An ordering with nothing left to promise is +/// used up. +/// +/// For example, if we have data ordered by `[a, b]` as well as `[c]`, and we +/// walk the requirement `[a, c]`, pinning `a` after the first step results in: +/// +/// | Step | Requirement Key | Pinned | Head `[a, b]` | Head `[c]` | Outcome | +/// |------|-----------------|--------|---------------|------------|-----------| +/// | 1 | a | {} | a | c | pin a | +/// | 2 | c | {a} | b | c | satisfied | +/// +/// It may be confusing that the head of `[c]` does not advance. The intuition +/// is that `[c]` still promises the data is sorted by `c` for all equal values +/// of `a`. This is true because the data was entirely sorted by `[c]` to begin +/// with, and pinning `a` does not change the order of the rows that tie on `a`, +/// as illustrated in this example: +/// +/// ```text +/// ┌───┬───┬───┐ +/// │ a │ b │ c │ +/// ├───┼───┼───┤ +/// │ 1 │ 1 │ 2 │ +/// │ 1 │ 3 │ 5 │ +/// │ 2 │ 1 │ 7 │ +/// │ 2 │ 2 │ 9 │ +/// └───┴───┴───┘ +/// ``` +/// +/// # Example: Walking a Requirement +/// +/// * Orderings: `[a ASC, b ASC, c ASC]` and `[d ASC]` +/// * Requirement: `[a ASC, b ASC]` +/// +/// | Step | Requirement Key | Pinned | Head `[a, b, c]` | Head `[d]` | Outcome | +/// |------|-----------------|--------|------------------|------------|-----------| +/// | 1 | a | {} | a | d | pin a | +/// | 2 | b | {a} | b | d | satisfied | +/// +/// This ordering is satisfied, (intuitively since the requirement is a prefix +/// of the first ordering). The second ordering plays no part. +/// +/// * Orderings: `[a ASC, b ASC]` and `[b ASC, c ASC]` +/// * Requirement: `[a ASC, b ASC, c ASC]` +/// +/// | Step | Requirement Key | Pinned | Head `[a, b]` | Head `[b, c]` | Outcome | +/// |------|-----------------|--------|---------------|---------------|-----------| +/// | 1 | a | {} | a | b | pin a | +/// | 2 | b | {a} | b | b | pin b | +/// | 3 | c | {a, b} | - | c | satisfied | +/// +/// Note that After step 2, `[a, b]` is used up as both of its expressions are +/// pinned. Among the rows that tie on `a` and `b`, the second ordering says the +/// rows are sorted by `c`, which can satisfy the third key. +/// +/// * Orderings: `[a ASC, c ASC]` and `[c ASC, b ASC]` +/// * Requirement: `[a ASC, b ASC, c ASC]` +/// +/// | Step | Requirement Key | Pinned | Head `[a, c]` | Head `[c, b]` | Outcome | +/// |------|-----------------|--------|---------------|---------------|-------------| +/// | 1 | a | {} | a | c | pin a | +/// | 2 | b | {a} | c | c | unsatisfied | +/// +/// After pinning `a`, both orderings are headed by `c`. Since nothing says the +/// rows that tie on `a` are sorted by `b`, the walk stops, and reports that the +/// requirement is not satisfied. +/// +/// * Orderings: `[a ASC, b ASC]` and `[a ASC, c ASC]` +/// * Requirement: `[a ASC, c ASC]` +/// +/// | Step | Requirement Key | Pinned | Head `[a, b]` | Head `[a, c]` | Outcome | +/// |------|-----------------|--------|---------------|---------------|-----------| +/// | 1 | a | {} | a | a | pin a | +/// | 2 | c | {a} | b | c | satisfied | +/// +/// [Checking whether a requirement is satisfied]: EquivalenceProperties#checking-whether-a-requirement-is-satisfied pub(super) struct SortKeyWalk<'a> { props: &'a EquivalenceProperties, - /// Pinned expressions. Pinning a member of an equivalence class pins all - /// of its members. + /// Pinned expressions. pinned: HashSet>, /// Orderings that hold among the rows that tie on the pinned expressions, /// 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 position of the head of each ordering. Empty until the first pin. + /// + /// `head_idxs[i]` tracks the position in the `LexOrdering` at index `i`, + /// as returned by [`Self::ordering`]. + /// + /// The ordering at index `i` is: + /// * if `i < self.props.normal_cls.len()` then `self.props.normal_cls[i]` + /// * else, `self.implied[i - self.props.normal_cls.len()]` otherwise. See + /// + /// The value is either an index into that ordering's sort expressions, or + /// its length when the ordering has been exhausted. head_idxs: Vec, } @@ -132,10 +230,13 @@ impl<'a> SortKeyWalk<'a> { Ok(()) } - /// Pins `expr` and the rest of its equivalence class (but, unlike - /// [`Self::pin_key`], not its children), then moves the head of each - /// ordering past the expressions that are now fixed, and adds the - /// orderings that the heads imply, until no new ordering is found. + /// Pins `expr`. + /// + /// See [`SortKeyWalk`] for an explanation of pinning. + /// + /// + /// This method pins `expr` and thus by extension all members of its + /// equivalence class), but not its children. See [`Self::pin_key`] for more details. pub(super) fn pin(&mut self, expr: Arc) -> Result<()> { match self.props.eq_group.get_equivalence_class(&expr) { Some(class) => self.pinned.extend(class.iter().cloned()), @@ -250,11 +351,12 @@ impl<'a> SortKeyWalk<'a> { // strict: node.data.strictly_order_preserving = true; } - // Now, check what we know about orderings. A whole-partition window - // such as `SUM(x) OVER (PARTITION BY k)` adds `[k, w]` in all four - // directions of `w` (see `OrderingEquivalenceClass::is_expr_partial_const`). - // Once `k` is pinned, `w` heads orderings in opposite directions, so it - // has one value among the rows that tie on `k`: + // Now, check what we know about orderings. A window aggregate over a + // whole window partition, such as `SUM(x) OVER (PARTITION BY k)`, adds + // `[k, w]` in all four directions of `w` (see `EquivalenceProperties` + // documentation for more details). Once `k` is pinned, `w` heads + // orderings in opposite directions, so it has one value among the rows + // that tie on `k`: let normal_expr = self.props.eq_group.normalize_expr(Arc::clone(&node.expr)); if self.is_fixed(&normal_expr) || is_partial_const(self.heads(), &normal_expr) { node.data.sort_properties = SortProperties::Singleton; @@ -264,10 +366,13 @@ impl<'a> SortKeyWalk<'a> { Ok(Transformed::yes(node)) } - /// 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 - /// members are all pinned. + /// Whether `expr` is fixed + /// + /// Note `expr` must be in [normal form]. + /// + /// See [`SortKeyWalk`] for more details + /// + /// [normal form]: EquivalenceProperties#normal-form fn is_fixed(&self, expr: &Arc) -> bool { if expr.is::() || self.pinned.contains(expr) { return true; @@ -427,8 +532,10 @@ mod tests { Ok(()) } - /// A window over a whole partition, such as `SUM(..) OVER (PARTITION BY - /// k)`, adds `[k, x]` in all four directions of `x`. Pinning `k` leaves `x` + /// A window over a whole window partition, such as `SUM(..) OVER + /// (PARTITION BY k)`, adds `[k, x]` in all four directions of `x`. + /// + /// Pinning `k` leaves `x` /// heading orderings in opposite directions, so `x` has one value, and so /// does `CAST(x AS BIGINT)`: `[k, CAST(x AS BIGINT)]` holds in every /// direction. Without the partial constant check, `x` would only be From f3fd64ad500d6f2ae6dd5b32e6bbd2cd61e78ed5 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 9 Nov 2026 09:57:13 -0500 Subject: [PATCH 2/3] cleanup --- .../src/equivalence/properties/walk.rs | 52 +++++++++---------- 1 file changed, 25 insertions(+), 27 deletions(-) diff --git a/datafusion/physical-expr/src/equivalence/properties/walk.rs b/datafusion/physical-expr/src/equivalence/properties/walk.rs index 858de0463be93..6b9f6d3e9b9c4 100644 --- a/datafusion/physical-expr/src/equivalence/properties/walk.rs +++ b/datafusion/physical-expr/src/equivalence/properties/walk.rs @@ -19,31 +19,28 @@ use datafusion_physical_expr_common::utils::ExprPropertiesNode; use itertools::Itertools; -/// The state of one walk along a sequence of sort keys, such as the keys of a -/// required ordering. +/// The state of a walk along a required ordering (a sequence of sort keys). /// /// [`EquivalenceProperties`] uses a `SortKeyWalk` to decide if a set of -/// orderings satisfy a requirement. See [Checking whether a requirement is +/// orderings satisfies a requirement. See [Checking whether a requirement is /// satisfied] for more details. /// -/// # Intuition and Introduction, and Definitions +/// # Intuition and Definitions /// -/// An ordering `[x, y, z]` is a chain of claims about the rows: they are sorted -/// by x; among rows with equal x, sorted by y; among rows with equal x and y, -/// sorted by z. +/// An ordering `[x, y, z]` can be thought of as a chain of claims about the +/// rows: they are sorted by `x`; among rows with equal `x`, sorted by `y`; +/// among rows with equal `x` and `y`, sorted by `z`. /// /// ## Pinning and Fixed Expressions /// /// During a walk, the code scans the ordering requirement left to right. When -/// it finds key `k` satisfied, it *pins* `k`, meaning that for the next steps -/// in the walk, `k` is effectively a constant across the remaining ordering -/// requirement. +/// it finds key `k` satisfied, it *pins* `k`, meaning that `k` is effectively a +/// constant across the remaining ordering requirement. /// /// An expression is *fixed* if it has one value for all rows that tie on the -/// currently pinned expressions. Fixed expressions include the pinned expressions -/// themselves, as well as literals, constants, and singleton leaf expressions -/// (such as an uncorrelated scalar subquery). See [`Self::is_fixed`] for more -/// details. +/// currently pinned expressions. Fixed expressions include the pinned +/// expressions themselves, as well as literals, constants, and singleton leaf +/// expressions (e.g. an uncorrelated scalar subquery). /// /// ## Heads /// @@ -88,8 +85,8 @@ use itertools::Itertools; /// | 1 | a | {} | a | d | pin a | /// | 2 | b | {a} | b | d | satisfied | /// -/// This ordering is satisfied, (intuitively since the requirement is a prefix -/// of the first ordering). The second ordering plays no part. +/// This requirement is satisfied (intuitively, because it is a prefix of the +/// first ordering). The second ordering plays no part. /// /// * Orderings: `[a ASC, b ASC]` and `[b ASC, c ASC]` /// * Requirement: `[a ASC, b ASC, c ASC]` @@ -100,9 +97,9 @@ use itertools::Itertools; /// | 2 | b | {a} | b | b | pin b | /// | 3 | c | {a, b} | - | c | satisfied | /// -/// Note that After step 2, `[a, b]` is used up as both of its expressions are +/// Note that after step 2, `[a, b]` is used up, as both of its expressions are /// pinned. Among the rows that tie on `a` and `b`, the second ordering says the -/// rows are sorted by `c`, which can satisfy the third key. +/// rows are sorted by `c`, which satisfies the third key. /// /// * Orderings: `[a ASC, c ASC]` and `[c ASC, b ASC]` /// * Requirement: `[a ASC, b ASC, c ASC]` @@ -113,7 +110,7 @@ use itertools::Itertools; /// | 2 | b | {a} | c | c | unsatisfied | /// /// After pinning `a`, both orderings are headed by `c`. Since nothing says the -/// rows that tie on `a` are sorted by `b`, the walk stops, and reports that the +/// rows that tie on `a` are sorted by `b`, the walk stops and reports that the /// requirement is not satisfied. /// /// * Orderings: `[a ASC, b ASC]` and `[a ASC, c ASC]` @@ -138,8 +135,9 @@ pub(super) struct SortKeyWalk<'a> { /// as returned by [`Self::ordering`]. /// /// The ordering at index `i` is: - /// * if `i < self.props.normal_cls.len()` then `self.props.normal_cls[i]` - /// * else, `self.implied[i - self.props.normal_cls.len()]` otherwise. See + /// * `self.props.oeq_cache.normal_cls[i]` if `i` is less than + /// `self.props.oeq_cache.normal_cls.len()`, + /// * `self.implied[i - self.props.oeq_cache.normal_cls.len()]` otherwise. /// /// The value is either an index into that ordering's sort expressions, or /// its length when the ordering has been exhausted. @@ -234,9 +232,9 @@ impl<'a> SortKeyWalk<'a> { /// /// See [`SortKeyWalk`] for an explanation of pinning. /// - /// - /// This method pins `expr` and thus by extension all members of its - /// equivalence class), but not its children. See [`Self::pin_key`] for more details. + /// This method pins `expr` (and thus, by extension, all members of its + /// equivalence class), but not its children. See [`Self::pin_key`] for + /// more details. pub(super) fn pin(&mut self, expr: Arc) -> Result<()> { match self.props.eq_group.get_equivalence_class(&expr) { Some(class) => self.pinned.extend(class.iter().cloned()), @@ -366,11 +364,11 @@ impl<'a> SortKeyWalk<'a> { Ok(Transformed::yes(node)) } - /// Whether `expr` is fixed + /// Whether `expr` is fixed. /// - /// Note `expr` must be in [normal form]. + /// Note that `expr` must be in [normal form]. /// - /// See [`SortKeyWalk`] for more details + /// See [`SortKeyWalk`] for more details. /// /// [normal form]: EquivalenceProperties#normal-form fn is_fixed(&self, expr: &Arc) -> bool { From 0c5a1550307b22c63b12a42229aa81052d227c4f Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 9 Nov 2026 10:04:02 -0500 Subject: [PATCH 3/3] cleanup --- .../src/equivalence/properties/walk.rs | 27 ++++++++++++++----- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/datafusion/physical-expr/src/equivalence/properties/walk.rs b/datafusion/physical-expr/src/equivalence/properties/walk.rs index 6b9f6d3e9b9c4..7f5d5651838de 100644 --- a/datafusion/physical-expr/src/equivalence/properties/walk.rs +++ b/datafusion/physical-expr/src/equivalence/properties/walk.rs @@ -154,9 +154,11 @@ impl<'a> SortKeyWalk<'a> { } } - /// Returns how many leading keys (in normal form) the existing orderings + /// Returns how many leading keys (in [normal form]) the existing orderings /// satisfy, pinning each satisfied key for the keys after it. A key /// without sort options only needs to be ordered in some direction. + /// + /// [normal form]: EquivalenceProperties#normal-form pub(super) fn sorted_prefix_len<'k>( mut self, keys: impl ExactSizeIterator, Option)>, @@ -228,13 +230,15 @@ impl<'a> SortKeyWalk<'a> { Ok(()) } - /// Pins `expr`. + /// Pins `expr`, an expression in [normal form]. /// /// See [`SortKeyWalk`] for an explanation of pinning. /// /// This method pins `expr` (and thus, by extension, all members of its /// equivalence class), but not its children. See [`Self::pin_key`] for /// more details. + /// + /// [normal form]: EquivalenceProperties#normal-form pub(super) fn pin(&mut self, expr: Arc) -> Result<()> { match self.props.eq_group.get_equivalence_class(&expr) { Some(class) => self.pinned.extend(class.iter().cloned()), @@ -395,8 +399,10 @@ impl<'a> SortKeyWalk<'a> { children.into_iter().all(|child| self.is_fixed(child)) } - /// The ordering at `idx`: the normal orderings of `props`, then + /// The ordering at `idx`: the orderings of `props` in [normal form], then /// `implied`. + /// + /// [normal form]: EquivalenceProperties#normal-form fn ordering(&self, idx: usize) -> &LexOrdering { let normal = &self.props.oeq_cache.normal_cls; normal @@ -574,10 +580,17 @@ mod tests { Ok(()) } - /// `find_longest_permutation` pins what it finds in normal form. With - /// `b = c` and ordering `[a + b, d]`, finding either `a + b` or `a + c` - /// pins the head of the ordering, whichever of `b` and `c` the class - /// normalizes to, so `d` is found next. + /// `find_longest_permutation` pins what it finds in [normal form]. + /// + /// For example, with `b = c` and ordering `[a + b, d]`, finding either + /// + /// * `a + b` + /// * `a + c` + /// + /// pins the head of the ordering (whichever of `b` or `c` the ordering + /// normalizes to), so `d` is found next. + /// + /// [normal form]: EquivalenceProperties#normal-form #[test] fn find_longest_permutation_pins_normal_form() -> Result<()> { let schema = Arc::new(Schema::new(