[GitHub] [arrow-datafusion] mustafasrepo commented on a diff in pull request #6452: Ordering Equivalence Builder

["mustafasrepo (via GitHub)" <gi...@apache.org>]

mustafasrepo commented on code in PR #6452:
URL: https://github.com/apache/arrow-datafusion/pull/6452#discussion_r1206790211


##########
datafusion/physical-expr/src/equivalence.rs:
##########
@@ -259,6 +261,80 @@ impl OrderingEquivalentClass {
     }
 }
 
+/// This is a builder object facilitating incremental construction
+/// for ordering equivalences.
+pub struct OrderingEquivalenceBuilder {
+    eq_properties: EquivalenceProperties,
+    ordering_eq_properties: OrderingEquivalenceProperties,
+    existing_ordering: Vec<PhysicalSortExpr>,
+}
+
+impl OrderingEquivalenceBuilder {
+    pub fn new(schema: SchemaRef) -> Self {
+        let eq_properties = EquivalenceProperties::new(schema.clone());
+        let ordering_eq_properties = OrderingEquivalenceProperties::new(schema);
+        Self {
+            eq_properties,
+            ordering_eq_properties,
+            existing_ordering: vec![],
+        }
+    }
+
+    pub fn extend(
+        mut self,
+        new_ordering_eq_properties: OrderingEquivalenceProperties,
+    ) -> Self {
+        self.ordering_eq_properties
+            .extend(new_ordering_eq_properties.classes().iter().cloned());
+        self
+    }
+
+    pub fn with_existing_ordering(
+        mut self,
+        existing_ordering: Option<Vec<PhysicalSortExpr>>,
+    ) -> Self {
+        if let Some(existing_ordering) = existing_ordering {
+            self.existing_ordering = existing_ordering;
+        }
+        self
+    }
+
+    pub fn with_equivalences(mut self, new_eq_properties: EquivalenceProperties) -> Self {
+        self.eq_properties = new_eq_properties;
+        self
+    }
+
+    pub fn add_equal_conditions(&mut self, new_equivalent_ordering: Vec<OrderedColumn>) {
+        let mut normalized_out_ordering = vec![];
+        for item in &self.existing_ordering {
+            // To account for ordering equivalences, first normalize the expression:
+            let normalized = normalize_expr_with_equivalence_properties(
+                item.expr.clone(),
+                self.eq_properties.classes(),
+            );
+            // Currently we only support ordering equivalences for `Column` expressions.
+            // TODO: Add support for ordering equivalence for all `PhysicalExpr`s.
+            if let Some(column) = normalized.as_any().downcast_ref::<Column>() {
+                normalized_out_ordering
+                    .push(OrderedColumn::new(column.clone(), item.options));
+            } else {
+                break;
+            }
+        }
+        // If there is an existing ordering, add new ordering as an equivalence:
+        if !normalized_out_ordering.is_empty() {
+            self.ordering_eq_properties.add_equal_conditions((
+                &normalized_out_ordering,
+                &new_equivalent_ordering,
+            ));
+        }
+    }
+
+    pub fn ordering_equivalence(&self) -> OrderingEquivalenceProperties {
+        self.ordering_eq_properties.clone()
+    }

Review Comment:
   Fixed it, thanks for the suggestion



-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org