viirya commented on code in PR #16217:
URL: https://github.com/apache/datafusion/pull/16217#discussion_r2124375915
##########
datafusion/physical-expr/src/equivalence/properties/mod.rs:
##########
@@ -190,382 +241,363 @@ impl EquivalenceProperties {
&self.oeq_class
}
- /// Return the inner OrderingEquivalenceClass, consuming self
- pub fn into_oeq_class(self) -> OrderingEquivalenceClass {
- self.oeq_class
- }
-
/// Returns a reference to the equivalence group within.
pub fn eq_group(&self) -> &EquivalenceGroup {
&self.eq_group
}
- /// Returns a reference to the constant expressions
- pub fn constants(&self) -> &[ConstExpr] {
- &self.constants
- }
-
+ /// Returns a reference to the constraints within.
pub fn constraints(&self) -> &Constraints {
&self.constraints
}
- /// Returns the output ordering of the properties.
- pub fn output_ordering(&self) -> Option<LexOrdering> {
- let constants = self.constants();
- let mut output_ordering =
self.oeq_class().output_ordering().unwrap_or_default();
- // Prune out constant expressions
- output_ordering
- .retain(|sort_expr| !const_exprs_contains(constants,
&sort_expr.expr));
- (!output_ordering.is_empty()).then_some(output_ordering)
+ /// Returns all the known constants expressions.
+ pub fn constants(&self) -> Vec<ConstExpr> {
+ self.eq_group
+ .iter()
+ .filter_map(|c| {
+ c.constant.as_ref().and_then(|across| {
+ c.canonical_expr()
+ .map(|expr| ConstExpr::new(Arc::clone(expr),
across.clone()))
+ })
+ })
+ .collect()
}
- /// Returns the normalized version of the ordering equivalence class
within.
- /// Normalization removes constants and duplicates as well as standardizing
- /// expressions according to the equivalence group within.
- pub fn normalized_oeq_class(&self) -> OrderingEquivalenceClass {
- OrderingEquivalenceClass::new(
- self.oeq_class
- .iter()
- .map(|ordering| self.normalize_sort_exprs(ordering))
- .collect(),
- )
+ /// Returns the output ordering of the properties.
+ pub fn output_ordering(&self) -> Option<LexOrdering> {
+ let concat = self.oeq_class.iter().flat_map(|o| o.iter().cloned());
+ self.normalize_sort_exprs(concat)
}
/// Extends this `EquivalenceProperties` with the `other` object.
- pub fn extend(mut self, other: Self) -> Self {
- self.eq_group.extend(other.eq_group);
- self.oeq_class.extend(other.oeq_class);
- self.with_constants(other.constants)
+ pub fn extend(mut self, other: Self) -> Result<Self> {
+ self.constraints.extend(other.constraints);
+ self.add_equivalence_group(other.eq_group)?;
+ self.add_orderings(other.oeq_class);
+ Ok(self)
}
/// Clears (empties) the ordering equivalence class within this object.
/// Call this method when existing orderings are invalidated.
pub fn clear_orderings(&mut self) {
self.oeq_class.clear();
+ self.oeq_cache.clear();
}
/// Removes constant expressions that may change across partitions.
- /// This method should be used when data from different partitions are
merged.
+ /// This method should be used when merging data from different partitions.
pub fn clear_per_partition_constants(&mut self) {
- self.constants.retain(|item| {
- matches!(item.across_partitions(), AcrossPartitions::Uniform(_))
- })
- }
-
- /// Extends this `EquivalenceProperties` by adding the orderings inside the
- /// ordering equivalence class `other`.
- pub fn add_ordering_equivalence_class(&mut self, other:
OrderingEquivalenceClass) {
- self.oeq_class.extend(other);
+ if self.eq_group.clear_per_partition_constants() {
+ // Renormalize orderings if the equivalence group changes:
+ let normal_orderings = self
+ .oeq_class
+ .iter()
+ .cloned()
+ .map(|o| self.eq_group.normalize_sort_exprs(o));
+ self.oeq_cache = OrderingEquivalenceCache::new(normal_orderings);
+ }
}
/// Adds new orderings into the existing ordering equivalence class.
- pub fn add_new_orderings(
+ pub fn add_orderings(
&mut self,
- orderings: impl IntoIterator<Item = LexOrdering>,
+ orderings: impl IntoIterator<Item = impl IntoIterator<Item =
PhysicalSortExpr>>,
) {
- self.oeq_class.add_new_orderings(orderings);
+ let orderings: Vec<_> =
+ orderings.into_iter().filter_map(LexOrdering::new).collect();
+ let normal_orderings: Vec<_> = orderings
+ .iter()
+ .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 redunant during extension,
we
+ // can use a shortcut algorithm to update the leading map.
+ self.oeq_cache.update_map();
+ }
}
/// Adds a single ordering to the existing ordering equivalence class.
- pub fn add_new_ordering(&mut self, ordering: LexOrdering) {
- self.add_new_orderings([ordering]);
+ pub fn add_ordering(&mut self, ordering: impl IntoIterator<Item =
PhysicalSortExpr>) {
+ self.add_orderings(std::iter::once(ordering));
}
/// Incorporates the given equivalence group to into the existing
/// equivalence group within.
- pub fn add_equivalence_group(&mut self, other_eq_group: EquivalenceGroup) {
- self.eq_group.extend(other_eq_group);
+ pub fn add_equivalence_group(
+ &mut self,
+ other_eq_group: EquivalenceGroup,
+ ) -> Result<()> {
+ if !other_eq_group.is_empty() {
+ self.eq_group.extend(other_eq_group);
+ // Renormalize orderings if the equivalence group changes:
+ let normal_cls = mem::take(&mut self.oeq_cache.normal_cls);
+ let normal_orderings = normal_cls
+ .into_iter()
+ .map(|o| self.eq_group.normalize_sort_exprs(o));
+ self.oeq_cache.normal_cls =
OrderingEquivalenceClass::new(normal_orderings);
+ self.oeq_cache.update_map();
+ // Discover any new orderings based on the new equivalence classes:
+ let leading_exprs: Vec<_> =
+ self.oeq_cache.leading_map.keys().cloned().collect();
+ for expr in leading_exprs {
+ self.discover_new_orderings(expr)?;
+ }
+ }
+ Ok(())
+ }
+
+ /// Returns the ordering equivalence class within in normal form.
+ /// Normalization standardizes expressions according to the equivalence
+ /// group within, and removes constants/duplicates.
+ pub fn normalized_oeq_class(&self) -> OrderingEquivalenceClass {
+ self.oeq_class
+ .iter()
+ .cloned()
+ .filter_map(|ordering| self.normalize_sort_exprs(ordering))
+ .collect::<Vec<_>>()
+ .into()
}
/// Adds a new equality condition into the existing equivalence group.
/// If the given equality defines a new equivalence class, adds this new
/// equivalence class to the equivalence group.
pub fn add_equal_conditions(
&mut self,
- left: &Arc<dyn PhysicalExpr>,
- right: &Arc<dyn PhysicalExpr>,
+ left: Arc<dyn PhysicalExpr>,
+ right: Arc<dyn PhysicalExpr>,
) -> Result<()> {
- // Discover new constants in light of new the equality:
- if self.is_expr_constant(left) {
- // Left expression is constant, add right as constant
- if !const_exprs_contains(&self.constants, right) {
- let const_expr = ConstExpr::from(right)
-
.with_across_partitions(self.get_expr_constant_value(left));
- self.constants.push(const_expr);
- }
- } else if self.is_expr_constant(right) {
- // Right expression is constant, add left as constant
- if !const_exprs_contains(&self.constants, left) {
- let const_expr = ConstExpr::from(left)
-
.with_across_partitions(self.get_expr_constant_value(right));
- self.constants.push(const_expr);
- }
+ // Add equal expressions to the state:
+ if self.eq_group.add_equal_conditions(Arc::clone(&left), right) {
+ // Renormalize orderings if the equivalence group changes:
+ let normal_cls = mem::take(&mut self.oeq_cache.normal_cls);
+ let normal_orderings = normal_cls
+ .into_iter()
+ .map(|o| self.eq_group.normalize_sort_exprs(o));
+ self.oeq_cache.normal_cls =
OrderingEquivalenceClass::new(normal_orderings);
+ self.oeq_cache.update_map();
+ // Discover any new orderings:
+ self.discover_new_orderings(left)?;
}
-
- // Add equal expressions to the state
- self.eq_group.add_equal_conditions(left, right);
-
- // Discover any new orderings
- self.discover_new_orderings(left)?;
Ok(())
}
/// Track/register physical expressions with constant values.
- #[deprecated(since = "43.0.0", note = "Use [`with_constants`] instead")]
- pub fn add_constants(self, constants: impl IntoIterator<Item = ConstExpr>)
-> Self {
- self.with_constants(constants)
- }
-
- /// Remove the specified constant
- pub fn remove_constant(mut self, c: &ConstExpr) -> Self {
- self.constants.retain(|existing| existing != c);
- self
- }
-
- /// Track/register physical expressions with constant values.
- pub fn with_constants(
- mut self,
+ pub fn add_constants(
+ &mut self,
constants: impl IntoIterator<Item = ConstExpr>,
- ) -> Self {
- let normalized_constants = constants
- .into_iter()
- .filter_map(|c| {
- let across_partitions = c.across_partitions();
- let expr = c.owned_expr();
- let normalized_expr = self.eq_group.normalize_expr(expr);
-
- if const_exprs_contains(&self.constants, &normalized_expr) {
- return None;
- }
-
- let const_expr = ConstExpr::from(normalized_expr)
- .with_across_partitions(across_partitions);
-
- Some(const_expr)
+ ) -> Result<()> {
+ // Add the new constant to the equivalence group:
+ for constant in constants {
+ self.eq_group.add_constant(constant);
+ }
+ // Renormalize the orderings after adding new constants by removing
+ // the constants from existing orderings:
+ let normal_cls = mem::take(&mut self.oeq_cache.normal_cls);
+ let normal_orderings = normal_cls.into_iter().map(|ordering| {
+ ordering.into_iter().filter(|sort_expr| {
+ self.eq_group.is_expr_constant(&sort_expr.expr).is_none()
})
- .collect::<Vec<_>>();
-
- // Add all new normalized constants
- self.constants.extend(normalized_constants);
-
- // Discover any new orderings based on the constants
- for ordering in self.normalized_oeq_class().iter() {
- if let Err(e) = self.discover_new_orderings(&ordering[0].expr) {
- log::debug!("error discovering new orderings: {e}");
- }
+ });
+ self.oeq_cache.normal_cls =
OrderingEquivalenceClass::new(normal_orderings);
+ self.oeq_cache.update_map();
+ // Discover any new orderings based on the constants:
+ let leading_exprs: Vec<_> =
self.oeq_cache.leading_map.keys().cloned().collect();
+ for expr in leading_exprs {
+ self.discover_new_orderings(expr)?;
}
-
- self
+ Ok(())
}
- // Discover new valid orderings in light of a new equality.
- // Accepts a single argument (`expr`) which is used to determine
- // which orderings should be updated.
- // When constants or equivalence classes are changed, there may be new
orderings
- // that can be discovered with the new equivalence properties.
- // For a discussion, see: https://github.com/apache/datafusion/issues/9812
- fn discover_new_orderings(&mut self, expr: &Arc<dyn PhysicalExpr>) ->
Result<()> {
- let normalized_expr = self.eq_group().normalize_expr(Arc::clone(expr));
+ /// Discover new valid orderings in light of a new equality. Accepts a
single
+ /// argument (`expr`) which is used to determine the orderings to update.
+ /// When constants or equivalence classes change, there may be new
orderings
+ /// that can be discovered with the new equivalence properties.
+ /// For a discussion, see:
<https://github.com/apache/datafusion/issues/9812>
+ fn discover_new_orderings(
+ &mut self,
+ normal_expr: Arc<dyn PhysicalExpr>,
+ ) -> Result<()> {
+ let Some(ordering_idxs) = self.oeq_cache.leading_map.get(&normal_expr)
else {
+ return Ok(());
+ };
let eq_class = self
.eq_group
- .iter()
- .find_map(|class| {
- class
- .contains(&normalized_expr)
- .then(|| class.clone().into_vec())
- })
- .unwrap_or_else(|| vec![Arc::clone(&normalized_expr)]);
-
- let mut new_orderings: Vec<LexOrdering> = vec![];
- for ordering in self.normalized_oeq_class().iter() {
- if !ordering[0].expr.eq(&normalized_expr) {
- continue;
- }
+ .get_equivalence_class(&normal_expr)
+ .map_or_else(|| vec![normal_expr], |class| class.clone().into());
+ let mut new_orderings = vec![];
+ for idx in ordering_idxs {
+ let ordering = &self.oeq_cache.normal_cls[*idx];
let leading_ordering_options = ordering[0].options;
- for equivalent_expr in &eq_class {
+ 'exprs: for equivalent_expr in &eq_class {
let children = equivalent_expr.children();
if children.is_empty() {
continue;
}
-
- // Check if all children match the next expressions in the
ordering
- let mut all_children_match = true;
+ // Check if all children match the next expressions in the
ordering:
let mut child_properties = vec![];
-
- // Build properties for each child based on the next
expressions
- for (i, child) in children.iter().enumerate() {
- if let Some(next) = ordering.get(i + 1) {
- if !child.as_ref().eq(next.expr.as_ref()) {
- all_children_match = false;
- break;
- }
- child_properties.push(ExprProperties {
- sort_properties:
SortProperties::Ordered(next.options),
- range: Interval::make_unbounded(
- &child.data_type(&self.schema)?,
- )?,
- preserves_lex_ordering: true,
- });
- } else {
- all_children_match = false;
- break;
+ // Build properties for each child based on the next
expression:
+ for (i, child) in children.into_iter().enumerate() {
+ let Some(next) = ordering.get(i + 1) else {
+ break 'exprs;
+ };
+ if !next.expr.eq(child) {
+ break 'exprs;
}
+ let data_type = child.data_type(&self.schema)?;
+ child_properties.push(ExprProperties {
+ sort_properties: SortProperties::Ordered(next.options),
+ range: Interval::make_unbounded(&data_type)?,
+ preserves_lex_ordering: true,
+ });
}
-
- if all_children_match {
- // Check if the expression is monotonic in all arguments
- if let Ok(expr_properties) =
- equivalent_expr.get_properties(&child_properties)
- {
- if expr_properties.preserves_lex_ordering
- &&
SortProperties::Ordered(leading_ordering_options)
- == expr_properties.sort_properties
- {
- // Assume existing ordering is [c ASC, a ASC, b
ASC]
- // When equality c = f(a,b) is given, if we know
that given ordering `[a ASC, b ASC]`,
- // ordering `[f(a,b) ASC]` is valid, then we can
deduce that ordering `[a ASC, b ASC]` is also valid.
- // Hence, ordering `[a ASC, b ASC]` can be added
to the state as a valid ordering.
- // (e.g. existing ordering where leading ordering
is removed)
-
new_orderings.push(LexOrdering::new(ordering[1..].to_vec()));
- break;
- }
- }
+ // Check if the expression is monotonic in all arguments:
+ let expr_properties =
+ equivalent_expr.get_properties(&child_properties)?;
+ if expr_properties.preserves_lex_ordering
+ && expr_properties.sort_properties
+ == SortProperties::Ordered(leading_ordering_options)
+ {
+ // Assume that `[c ASC, a ASC, b ASC]` is among existing
+ // orderings. If equality `c = f(a, b)` is given, ordering
+ // `[a ASC, b ASC]` implies the ordering `[c ASC]`. Thus,
+ // ordering `[a ASC, b ASC]` is also a valid ordering.
Review Comment:
Got it. I saw the `preserves_lex_ordering` check.
--
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: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
