Add display implementation for equivalence properties, sort, etc

Author: alamb

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

@@ -223,3 +223,22 @@ pub fn down_cast_any_ref(any: &dyn Any) -> &dyn Any {
         any
     }
 }
+
+/// Writes a list of [`PhysicalExpr`]s to a `std::fmt::Formatter`.
+///
+/// Example output: `[a + 1, b]`
+pub fn format_physical_expr_list(
+    f: &mut std::fmt::Formatter<'_>,
+    ordering: &[Arc<dyn PhysicalExpr>],
+) -> std::fmt::Result {
+    let mut iter = ordering.iter();
+    write!(f, "[")?;
+    if let Some(expr) = iter.next() {
+        write!(f, "{}", expr)?;
+    }
+    for expr in iter {
+        write!(f, ", {}", expr)?;
+    }
+    write!(f, "]")?;
+    Ok(())
+}

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

@@ -30,6 +30,51 @@ use datafusion_common::Result;
 use datafusion_expr_common::columnar_value::ColumnarValue;
 
 /// Represents Sort operation for a column in a RecordBatch
+///
+/// Example:
+/// ```no_run
+/// # use std::any::Any;
+/// # use std::fmt::Display;
+/// # use std::hash::Hasher;
+/// # use std::sync::Arc;
+/// # use arrow::array::RecordBatch;
+/// # use datafusion_common::Result;
+/// # use arrow::compute::SortOptions;
+/// # use arrow::datatypes::{DataType, Schema};
+/// # use datafusion_expr_common::columnar_value::ColumnarValue;
+/// # use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
+/// # use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
+/// # // this crate doesn't have a physical expression implementation
+/// # // so make a really simple one
+/// # #[derive(Clone, Debug, PartialEq, Eq, Hash)]
+/// # struct MyPhysicalExpr;
+/// # impl PhysicalExpr for MyPhysicalExpr {
+/// #  fn as_any(&self) -> &dyn Any {todo!() }
+/// #  fn data_type(&self, input_schema: &Schema) -> Result<DataType> {todo!()}
+/// #  fn nullable(&self, input_schema: &Schema) -> Result<bool> {todo!() }
+/// #  fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {todo!() }
+/// #  fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {todo!()}
+/// #  fn with_new_children(self: Arc<Self>, children: Vec<Arc<dyn PhysicalExpr>>) -> Result<Arc<dyn PhysicalExpr>> {todo!()}
+/// #  fn dyn_hash(&self, _state: &mut dyn Hasher) {todo!()}
+/// # }
+/// # impl Display for MyPhysicalExpr {
+/// #    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { write!(f, "a") }
+/// # }
+/// # impl PartialEq<dyn Any> for MyPhysicalExpr {
+/// #    fn eq(&self, _other: &dyn Any) -> bool { true }
+/// # }
+/// # fn col(name: &str) -> Arc<dyn PhysicalExpr> { Arc::new(MyPhysicalExpr) }
+/// // Sort by a ASC
+/// let options = SortOptions::default();
+/// let sort_expr = PhysicalSortExpr::new(col("a"), options);
+/// assert_eq!(sort_expr.to_string(), "a ASC NULLS LAST");
+///
+/// // Sort by a DESC NULLS FIRST
+/// let sort_expr = PhysicalSortExpr::new_default(col("a"))
+///   .desc()
+///   .nulls_first();
+/// assert_eq!(sort_expr.to_string(), "a DESC NULLS FIRST");
+/// ```
 #[derive(Clone, Debug)]
 pub struct PhysicalSortExpr {
     /// Physical expression representing the column to sort
@@ -43,6 +88,35 @@ impl PhysicalSortExpr {
     pub fn new(expr: Arc<dyn PhysicalExpr>, options: SortOptions) -> Self {
         Self { expr, options }
     }
+
+    /// Create a new PhysicalSortExpr with default [`SortOptions`]
+    pub fn new_default(expr: Arc<dyn PhysicalExpr>) -> Self {
+        Self::new(expr, SortOptions::default())
+    }
+
+    /// Set the sort sort options to ASC
+    pub fn asc(mut self) -> Self {
+        self.options.descending = false;
+        self
+    }
+
+    /// Set the sort sort options to DESC
+    pub fn desc(mut self) -> Self {
+        self.options.descending = true;
+        self
+    }
+
+    /// Set the sort sort options to NULLS FIRST
+    pub fn nulls_first(mut self) -> Self {
+        self.options.nulls_first = true;
+        self
+    }
+
+    /// Set the sort sort options to NULLS LAST
+    pub fn nulls_last(mut self) -> Self {
+        self.options.nulls_first = false;
+        self
+    }
 }
 
 impl PartialEq for PhysicalSortExpr {
@@ -60,7 +134,7 @@ impl Hash for PhysicalSortExpr {
     }
 }
 
-impl std::fmt::Display for PhysicalSortExpr {
+impl Display for PhysicalSortExpr {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         write!(f, "{} {}", self.expr, to_str(&self.options))
     }

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

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::fmt::Display;
 use std::sync::Arc;
 
 use super::{add_offset_to_expr, collapse_lex_req, ProjectionMapping};
@@ -27,6 +28,7 @@ use crate::{
 
 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.
@@ -101,6 +103,19 @@ impl ConstExpr {
     }
 }
 
+/// Display implementation for `ConstExpr`
+///
+/// Example `c` or `c(across_partitions)`
+impl Display for ConstExpr {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "{}", self.expr)?;
+        if self.across_partitions {
+            write!(f, "(across_partitions)")?;
+        }
+        Ok(())
+    }
+}
+
 impl From<Arc<dyn PhysicalExpr>> for ConstExpr {
     fn from(expr: Arc<dyn PhysicalExpr>) -> Self {
         Self::new(expr)
@@ -224,6 +239,14 @@ impl EquivalenceClass {
     }
 }
 
+impl Display for EquivalenceClass {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "[")?;
+        format_physical_expr_list(f, &self.exprs)?;
+        write!(f, "]")
+    }
+}
+
 /// An `EquivalenceGroup` is a collection of `EquivalenceClass`es where each
 /// class represents a distinct equivalence class in a relation.
 #[derive(Debug, Clone)]
@@ -575,6 +598,20 @@ impl EquivalenceGroup {
     }
 }
 
+impl Display for EquivalenceGroup {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "[")?;
+        let mut iter = self.iter();
+        if let Some(cls) = iter.next() {
+            write!(f, "{}", cls)?;
+        }
+        for cls in iter {
+            write!(f, ", {}", cls)?;
+        }
+        write!(f, "]")
+    }
+}
+
 #[cfg(test)]
 mod tests {
 

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

@@ -15,13 +15,13 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::fmt::Display;
 use std::hash::Hash;
 use std::sync::Arc;
 
-use arrow_schema::SortOptions;
-
 use crate::equivalence::add_offset_to_expr;
 use crate::{LexOrdering, PhysicalExpr, PhysicalSortExpr};
+use arrow_schema::SortOptions;
 
 /// An `OrderingEquivalenceClass` object keeps track of different alternative
 /// orderings than can describe a schema. For example, consider the following table:
@@ -219,6 +219,21 @@ fn resolve_overlap(orderings: &mut [LexOrdering], idx: usize, pre_idx: usize) ->
     false
 }
 
+impl Display for OrderingEquivalenceClass {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "[")?;
+        let mut iter = self.orderings.iter();
+        if let Some(ordering) = iter.next() {
+            write!(f, "{}", PhysicalSortExpr::format_list(ordering))?;
+        }
+        for ordering in iter {
+            write!(f, "{}", PhysicalSortExpr::format_list(ordering))?;
+        }
+        write!(f, "]")?;
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;

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

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::fmt::Display;
 use std::hash::{Hash, Hasher};
 use std::sync::Arc;
 
@@ -71,7 +72,7 @@ use itertools::Itertools;
 /// specifies the output sorted by EITHER `a ASC` or `b DESC`, the sort can be
 /// avoided.
 ///
-/// # Example equivalent equivalent expressions
+/// # Example equivalent expressions
 ///
 /// Similarly, consider the table below:
 ///
@@ -92,6 +93,32 @@ use itertools::Itertools;
 /// `Hash(a)` and output partitioning is `Hash(b)`, then DataFusion avoids
 /// repartitioning the data as the existing partitioning satisfies the
 /// requirement.
+///
+/// # Code Example
+/// ```
+/// use std::sync::Arc;
+/// use arrow_schema::{Schema, Field, DataType, SchemaRef};
+/// use datafusion_physical_expr::{ConstExpr, EquivalenceProperties};
+/// use datafusion_physical_expr::expressions::col;
+///
+/// let schema: SchemaRef = Arc::new(Schema::new(vec![
+///    Field::new("a", DataType::Int32, false),
+///    Field::new("b", DataType::Int32, false),
+///    Field::new("c", DataType::Int32, false),
+/// ]));
+///
+/// let col_a = col("a", &schema).unwrap();
+/// let col_b = col("b", &schema).unwrap();
+/// let col_c = col("c", &schema).unwrap();
+///
+/// // This object represents data that is sorted by a ASC, c DESC
+/// // with a single constant value of b
+/// let mut eq_properties = EquivalenceProperties::new(schema);
+/// eq_properties.add_constants(vec![ConstExpr::from(col_b)]);
+/// eq_properties.add_new_ordering(vec![col_a.asc(), col_c.desc()]);
+///
+/// assert_eq!(eq_properties.to_string(), "ff")
+/// ```
 #[derive(Debug, Clone)]
 pub struct EquivalenceProperties {
     /// Collection of equivalence classes that store expressions with the same
@@ -203,6 +230,11 @@ impl EquivalenceProperties {
         self.oeq_class.add_new_orderings(orderings);
     }
 
+    /// Adds a single ordering to the existing ordering equivalence class.
+    pub fn add_new_ordering(&mut self, ordering: LexOrdering) {
+        self.add_new_orderings([ordering]);
+    }
+
     /// Incorporates the given equivalence group to into the existing
     /// equivalence group within.
     pub fn add_equivalence_group(&mut self, other_eq_group: EquivalenceGroup) {
@@ -1059,6 +1091,37 @@ impl EquivalenceProperties {
     }
 }
 
+/// More readable display version of the `EquivalenceProperties`.
+///
+/// Format:
+/// ```text
+/// order: [[a ASC, b ASC], [a ASC, c ASC]], eq: [[a = b], [a = c]], const: [a = 1]
+/// ```
+impl Display for EquivalenceProperties {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        if self.eq_group.is_empty()
+            && self.oeq_class.is_empty()
+            && self.constants.is_empty()
+        {
+            return write!(f, "No properties");
+        }
+        if !self.oeq_class.is_empty() {
+            write!(f, "order: {}", self.oeq_class)?;
+        }
+        if !self.eq_group.is_empty() {
+            write!(f, ", eq: {}", self.eq_group)?;
+        }
+        if !self.constants.is_empty() {
+            write!(f, ", const: [")?;
+            for c in &self.constants {
+                write!(f, ", {}", c)?;
+            }
+            write!(f, "]")?;
+        }
+        Ok(())
+    }
+}
+
 /// Calculates the properties of a given [`ExprPropertiesNode`].
 ///
 /// Order information can be retrieved as:
@@ -3032,21 +3095,31 @@ mod tests {
         // Check whether constants are same
         let lhs_constants = lhs.constants();
         let rhs_constants = rhs.constants();
-        assert_eq!(lhs_constants.len(), rhs_constants.len(), "{}", err_msg);
+        assert_eq!(
+            lhs_constants.len(),
+            rhs_constants.len(),
+            "{err_msg}\nlhs: {lhs}\nrhs: {rhs}"
+        );
         for rhs_constant in rhs_constants {
             assert!(
                 const_exprs_contains(lhs_constants, rhs_constant.expr()),
-                "{}",
-                err_msg
+                "{err_msg}\nlhs: {lhs}\nrhs: {rhs}"
             );
         }
 
         // Check whether orderings are same.
         let lhs_orderings = lhs.oeq_class();
         let rhs_orderings = &rhs.oeq_class.orderings;
-        assert_eq!(lhs_orderings.len(), rhs_orderings.len(), "{}", err_msg);
+        assert_eq!(
+            lhs_orderings.len(),
+            rhs_orderings.len(),
+            "{err_msg}\nlhs: {lhs}\nrhs: {rhs}"
+        );
         for rhs_ordering in rhs_orderings {
-            assert!(lhs_orderings.contains(rhs_ordering), "{}", err_msg);
+            assert!(
+                lhs_orderings.contains(rhs_ordering),
+                "{err_msg}\nlhs: {lhs}\nrhs: {rhs}"
+            );
         }
     }
 }