viirya commented on code in PR #26853:
URL: https://github.com/apache/datafusion/pull/26853#discussion_r4418466393
##########
datafusion/physical-expr/src/equivalence/properties/walk.rs:
##########
@@ -0,0 +1,368 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+use std::sync::Arc;
+
+use super::EquivalenceProperties;
+use crate::equivalence::ordering::is_partial_const;
+use crate::expressions::{Column, Literal};
+use crate::{LexOrdering, PhysicalExpr, PhysicalSortExpr};
+
+use arrow::compute::SortOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_common::{HashSet, Result};
+use datafusion_expr::interval_arithmetic::Interval;
+use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
+use datafusion_physical_expr_common::sort_expr::options_compatible;
+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.
+///
+/// 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, or 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.
+///
+/// Pins live in the walk; the equivalence properties it reads do not change.
+pub(super) struct Walk<'a> {
+ props: &'a EquivalenceProperties,
+ /// Pinned expressions. Pinning a member of an equivalence class pins all
+ /// of its members.
+ pinned: HashSet<Arc<dyn PhysicalExpr>>,
+ /// Orderings that hold among the rows that tie on the pinned expressions,
+ /// found by [`EquivalenceProperties::implied_ordering`] as the walk goes.
+ implied: Vec<LexOrdering>,
+ /// 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.
+ head_idxs: Vec<usize>,
+}
+
+impl<'a> Walk<'a> {
+ pub(super) fn new(props: &'a EquivalenceProperties) -> Self {
+ Self {
+ props,
+ pinned: HashSet::new(),
+ implied: vec![],
+ head_idxs: vec![],
+ }
+ }
+
+ /// 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.
+ pub(super) fn sorted_prefix_len<'k>(
+ mut self,
+ keys: impl ExactSizeIterator<Item = (&'k Arc<dyn PhysicalExpr>,
Option<SortOptions>)>,
+ ) -> Result<usize> {
+ let schema = self.props.schema();
+ let len = keys.len();
+ for (idx, (expr, options)) in keys.enumerate() {
+ let satisfy = match
self.expr_properties(Arc::clone(expr)).sort_properties {
+ SortProperties::Ordered(existing) =>
options.is_none_or(|options| {
+ let nullable = expr.nullable(schema).unwrap_or(true);
+ // `existing` comes from the first head that `expr` (or,
+ // for a computed `expr`, its arguments) matches; another
+ // head may be `expr` itself in the required direction.
+ options_compatible(&existing, &options, nullable)
+ || self.heads().any(|head| {
+ options_compatible(&head.options, &options,
nullable)
+ && head.expr.eq(expr)
+ })
+ }),
+ // Singleton expressions satisfy any ordering.
+ SortProperties::Singleton => true,
+ SortProperties::Unordered => false,
+ };
+ if !satisfy {
+ return Ok(idx);
+ }
+ if idx + 1 < len {
+ self.pin_key(Arc::clone(expr))?;
+ }
+ }
+ Ok(len)
+ }
+
+ /// Pins a satisfied sort key, and the sub-expressions it pins down: if an
+ /// expression is strictly order-preserving, equal outputs imply equal
+ /// values of its ordered children, so within a tie group of the key those
+ /// children have one value as well. For example, if data is sorted by
+ /// `[a, b]`, the requirement `[CAST(a AS BIGINT) ASC, b ASC]` is
+ /// satisfied: `a` has one value within each group of equal
+ /// `CAST(a AS BIGINT)` values, and hence `b` is sorted within each such
+ /// group.
+ fn pin_key(&mut self, key: Arc<dyn PhysicalExpr>) -> Result<()> {
+ let mut stack = vec![key];
+ while let Some(expr) = stack.pop() {
+ let properties = self.expr_properties(Arc::clone(&expr));
+ if properties.strictly_order_preserving {
+ for child in expr.children() {
+ let child_properties =
self.expr_properties(Arc::clone(child));
+ if matches!(
+ child_properties.sort_properties,
+ SortProperties::Ordered(_)
+ ) {
+ stack.push(Arc::clone(child));
+ }
+ }
+ }
+ self.pin(expr)?;
+ }
+ 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.
+ pub(super) fn pin(&mut self, expr: Arc<dyn PhysicalExpr>) -> Result<()> {
+ match self.props.eq_group.get_equivalence_class(&expr) {
+ Some(class) => self.pinned.extend(class.iter().cloned()),
+ None => {
+ self.pinned.insert(expr);
+ }
+ }
+ let ordering_count = self.props.oeq_cache.normal_cls.len() +
self.implied.len();
+ self.head_idxs.resize(ordering_count, 0);
+ // Pins do not change below, so after the first round only the
+ // orderings found in the previous round can imply new ones.
+ let mut first = 0;
+ loop {
+ let count = self.head_idxs.len();
+ 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;
+ }
+ if let Some(head) = ordering.get(head_idx) {
+ let rest = ordering[head_idx + 1..]
+ .iter()
+ .filter(|sort_expr| !self.is_fixed(&sort_expr.expr));
+ found.extend(self.props.implied_ordering(head, rest)?);
+ }
+ self.head_idxs[idx] = head_idx;
+ }
+ for ordering in found.into_iter().filter_map(LexOrdering::new) {
+ let known =
+ (0..self.head_idxs.len()).any(|idx| *self.ordering(idx) ==
ordering);
+ if !known {
+ self.implied.push(ordering);
+ self.head_idxs.push(0);
+ }
+ }
+ // No new ordering:
+ if self.head_idxs.len() == count {
+ return Ok(());
+ }
+ first = count;
+ }
+ }
+
+ /// Returns the properties of `expr` among the rows that tie on the
+ /// pinned expressions.
+ pub(super) fn expr_properties(&self, expr: Arc<dyn PhysicalExpr>) ->
ExprProperties {
+ ExprPropertiesNode::new_unknown(expr)
+ .transform_up(|node| self.update_properties(node))
+ .data()
+ .map(|node| node.data)
+ .unwrap_or_else(|_| ExprProperties::new_unknown())
+ }
+
+ /// Calculates the properties of a given [`ExprPropertiesNode`], whose
+ /// children are already done.
+ ///
+ /// - An intermediate node takes its order and range from its children,
+ /// through [`PhysicalExpr::get_properties`]. A `Literal` leaf is a
+ /// singleton with a point range. A `Column` leaf is unordered, with an
+ /// unbounded range of its data type.
+ /// - Then, whatever its children said: a node that is fixed or partially
+ /// constant is a singleton, and a node that heads an ordering is ordered
+ /// like the first ordering it heads.
+ fn update_properties(
+ &self,
+ mut node: ExprPropertiesNode,
+ ) -> Result<Transformed<ExprPropertiesNode>> {
+ // First, try to gather the information from the children:
+ if !node.expr.children().is_empty() {
+ // We have an intermediate (non-leaf) node, account for its
children:
+ let children_props =
+ node.children.iter().map(|c| c.data.clone()).collect_vec();
+ node.data = node.expr.get_properties(&children_props)?;
+ } else if node.expr.is::<Literal>() {
+ // We have a Literal, which is one of the two possible leaf node
types:
+ node.data = node.expr.get_properties(&[])?;
+ } else if node.expr.is::<Column>() {
+ // We have a Column, which is the other possible leaf node type:
+ node.data.range =
+
Interval::make_unbounded(&node.expr.data_type(self.props.schema())?)?;
+ // A column is the identity mapping of itself, which is trivially
+ // strict:
+ node.data.strictly_order_preserving = true;
+ }
+ // Now, check what we know about orderings:
+ 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) {
Review Comment:
Nothing tests this `is_partial_const` call. `test_is_expr_partial_const`
covers `OrderingEquivalenceClass::is_expr_partial_const`, which nothing inside
DataFusion calls anymore. If I drop `|| is_partial_const(self.heads(),
&normal_expr)` here, the physical-expr and physical-optimizer tests and the
window/order/sort sqllogictests all still pass.
A walk-level test would pin it. For example, add the whole-partition window
shape `[k, x]` in all four directions of `x`, then check `[k, CAST(x AS BIGINT)
DESC NULLS FIRST]`. That is `true` with this check and `false` without it. The
any-head rule in `sorted_prefix_len` doesn't cover it, because the key is a
computed expression and `x` gets its direction from the first head.
##########
datafusion/physical-expr/src/equivalence/properties/mod.rs:
##########
@@ -1336,17 +1240,18 @@ impl EquivalenceProperties {
break;
}
// Remove indices that have an ordering from `search_indices`, and
- // treat ordered expressions as constants in subsequent iterations.
- // We can do this because the "next" key only matters in a
lexicographical
- // ordering when the keys to its left have the same values.
- //
- // Note that these expressions are not properly "constants". This
is just
- // an implementation strategy confined to this function.
- for (PhysicalSortExpr { expr, .. }, idx) in &ordered_exprs {
- let const_expr = ConstExpr::from(Arc::clone(expr));
- eq_properties.add_constants(std::iter::once(const_expr))?;
+ // pin the ordered expressions for the next iterations, if any. We
+ // can do this because the "next" key only matters in a
+ // lexicographical ordering when the keys to its left have the same
+ // values.
+ for (_, idx) in &ordered_exprs {
search_indices.shift_remove(idx);
}
+ if !search_indices.is_empty() {
+ for (PhysicalSortExpr { expr, .. }, _) in &ordered_exprs {
+ walk.pin(Arc::clone(expr))?;
Review Comment:
`exprs[idx]` is pinned as given, not in normal form. A non-normalized
compound expression is then never seen as fixed. For example, with `b = c` and
the ordering stored as `[a + b, d]`, `find_longest_permutation([a + c, d])`
finds only `a + c`. That matches `main`, because `add_constant` doesn't
normalize either, so this isn't a regression. But pinning
`self.eq_group.normalize_expr(Arc::clone(expr))` here looks like a cheap
improvement. Fine to leave for a follow-up.
Separately, the literal case from the description (`[5, k, b % 5]`, now
fully found) changes behavior here but has no test. A small unit test would
keep it from regressing.
##########
datafusion/physical-expr/src/equivalence/properties/walk.rs:
##########
@@ -0,0 +1,368 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+use std::sync::Arc;
+
+use super::EquivalenceProperties;
+use crate::equivalence::ordering::is_partial_const;
+use crate::expressions::{Column, Literal};
+use crate::{LexOrdering, PhysicalExpr, PhysicalSortExpr};
+
+use arrow::compute::SortOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_common::{HashSet, Result};
+use datafusion_expr::interval_arithmetic::Interval;
+use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
+use datafusion_physical_expr_common::sort_expr::options_compatible;
+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.
+///
+/// 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, or 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.
+///
+/// Pins live in the walk; the equivalence properties it reads do not change.
+pub(super) struct Walk<'a> {
Review Comment:
nit: `Walk` reads a bit generic next to `TreeNode` walking/visiting in the
same crate. Something like `SortKeyWalk` would say what it walks. Not blocking.
##########
datafusion/physical-expr/src/equivalence/ordering.rs:
##########
@@ -210,25 +210,19 @@ impl OrderingEquivalenceClass {
None
}
- /// Checks whether the given expression is partially constant according to
- /// this ordering equivalence class.
+ /// Checks whether `expr` has a single value according to this ordering
+ /// equivalence class: whether one ordering starts with `expr` in some
+ /// direction and another starts with it in the exact reverse, such as
+ /// `ASC NULLS LAST` and `DESC NULLS FIRST`. Rows sorted both ways all have
+ /// the same value (or are all null).
///
- /// This function determines whether `expr` appears in at least one
combination
- /// of `descending` and `nulls_first` options that indicate partial
constantness
- /// in a lexicographical ordering. Specifically, an expression is
considered
- /// a partial constant in this context if its `SortOptions` satisfies
either
- /// of the following conditions:
- /// - It is `descending` with `nulls_first` and _also_ `ascending` with
- /// `nulls_last`, OR
- /// - It is `descending` with `nulls_last` and _also_ `ascending` with
- /// `nulls_first`.
- ///
- /// The equivalence mechanism primarily uses `ConstExpr`s to represent
globally
- /// constant expressions. However, some expressions may only be partially
- /// constant within a lexicographical ordering. This function helps
identify
- /// such cases. If an expression is constant within a prefix ordering, it
is
- /// added as a constant during `ordering_satisfy_requirement()` iterations
- /// after the corresponding prefix requirement is satisfied.
+ /// This is how a value that is constant only within groups of rows is
+ /// recorded, since a `ConstExpr` is constant across all rows. For example,
+ /// `SUM(x) OVER (PARTITION BY k)` has one value per `k`, so the window
adds
+ /// `[k, w]` once for each of the four directions of `w`. Sort requirement
Review Comment:
This doc now describes how the requirement checks use partial constants
("Sort requirement checks look at the orderings past the keys they have
satisfied..."). That is the walk's job now, since it calls the free
`is_partial_const` with its own heads, not this method. Maybe keep the window
example here and move the sentence about satisfied keys to
`Walk::update_properties`, where it actually happens?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
