Account for constant equivalence properties in union, tests

Author: alamb

datafusion/physical-expr-common/src/sort_expr.rs

@@ -96,28 +96,48 @@ impl PhysicalSortExpr {
     }
 
     /// Set the sort sort options to ASC
-    pub fn asc(mut self) -> Self {
-        self.options.descending = false;
-        self
+    pub fn asc(self) -> Self {
+        self.with_descending(false)
     }
 
     /// Set the sort sort options to DESC
-    pub fn desc(mut self) -> Self {
-        self.options.descending = true;
-        self
+    pub fn desc(self) -> Self {
+        self.with_descending(true)
     }
 
-    /// Set the sort sort options to NULLS FIRST
-    pub fn nulls_first(mut self) -> Self {
-        self.options.nulls_first = true;
+    /// set the sort options `descending` flag
+    pub fn with_descending(mut self, descending: bool) -> Self {
+        self.options.descending = descending;
         self
     }
 
-    /// Set the sort sort options to NULLS LAST
-    pub fn nulls_last(mut self) -> Self {
-        self.options.nulls_first = false;
+    /// Set the sort options to NULLS FIRST
+    pub fn nulls_first(self) -> Self {
+        self.with_nulls_first(true)
+    }
+
+    /// Set the sort options to NULLS LAST
+    pub fn nulls_last(self) -> Self {
+        self.with_nulls_first(false)
+    }
+
+    /// set the sort options `nulls_first` flag
+    pub fn with_nulls_first(mut self, nulls_first: bool) -> Self {
+        self.options.nulls_first = nulls_first;
         self
     }
+
+    /// return the underlying PhysicalExpr
+    pub fn expr(&self) -> &Arc<dyn PhysicalExpr> {
+        &self.expr
+    }
+}
+
+/// Access the PhysicalSortExpr as a PhysicalExpr
+impl AsRef<dyn PhysicalExpr> for PhysicalSortExpr {
+    fn as_ref(&self) -> &(dyn PhysicalExpr + 'static) {
+        self.expr.as_ref()
+    }
 }
 
 impl PartialEq for PhysicalSortExpr {

datafusion/physical-expr/src/equivalence/class.rs

@@ -30,7 +30,6 @@ use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
 use datafusion_common::JoinType;
 use datafusion_physical_expr_common::physical_expr::format_physical_expr_list;
 
-#[derive(Debug, Clone)]
 /// A structure representing a expression known to be constant in a physical execution plan.
 ///
 /// The `ConstExpr` struct encapsulates an expression that is constant during the execution
@@ -41,9 +40,10 @@ use datafusion_physical_expr_common::physical_expr::format_physical_expr_list;
 ///
 /// - `expr`: Constant expression for a node in the physical plan.
 ///
-/// - `across_partitions`: A boolean flag indicating whether the constant expression is
-///   valid across partitions. If set to `true`, the constant expression has same value for all partitions.
-///   If set to `false`, the constant expression may have different values for different partitions.
+/// - `across_partitions`: A boolean flag indicating whether the constant
+///   expression is the same across partitions. If set to `true`, the constant
+///   expression has same value for all partitions. If set to `false`, the
+///   constant expression may have different values for different partitions.
 ///
 /// # Example
 ///
@@ -56,11 +56,22 @@ use datafusion_physical_expr_common::physical_expr::format_physical_expr_list;
 /// // create a constant expression from a physical expression
 /// let const_expr = ConstExpr::from(col);
 /// ```
+#[derive(Debug, Clone)]
 pub struct ConstExpr {
+    /// The  expression that is known to be constant (e.g. a `Column`)
     expr: Arc<dyn PhysicalExpr>,
+    /// Does the constant have the same value across all partitions? See
+    /// struct docs for more details
     across_partitions: bool,
 }
 
+impl PartialEq for ConstExpr {
+    fn eq(&self, other: &Self) -> bool {
+        self.across_partitions == other.across_partitions
+            && self.expr.eq(other.expr.as_any())
+    }
+}
+
 impl ConstExpr {
     /// Create a new constant expression from a physical expression.
     ///
@@ -74,11 +85,17 @@ impl ConstExpr {
         }
     }
 
+    /// Set the `across_partitions` flag
+    ///
+    /// See struct docs for more details
     pub fn with_across_partitions(mut self, across_partitions: bool) -> Self {
         self.across_partitions = across_partitions;
         self
     }
 
+    /// Is the  expression the same across all partitions?
+    ///
+    /// See struct docs for more details
     pub fn across_partitions(&self) -> bool {
         self.across_partitions
     }
@@ -101,6 +118,31 @@ impl ConstExpr {
             across_partitions: self.across_partitions,
         })
     }
+
+    /// Returns true if this constant expression is equal to the given expression
+    pub fn eq_expr(&self, other: impl AsRef<dyn PhysicalExpr>) -> bool {
+        self.expr.eq(other.as_ref().as_any())
+    }
+
+    /// Returns a [`Display`]able list of `ConstExpr`.
+    pub fn format_list(input: &[ConstExpr]) -> impl Display + '_ {
+        struct DisplayableList<'a>(&'a [ConstExpr]);
+        impl<'a> Display for DisplayableList<'a> {
+            fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+                let mut first = true;
+                for const_expr in self.0 {
+                    if first {
+                        first = false;
+                    } else {
+                        write!(f, ",")?;
+                    }
+                    write!(f, "{}", const_expr)?;
+                }
+                Ok(())
+            }
+        }
+        DisplayableList(input)
+    }
 }
 
 /// Display implementation for `ConstExpr`