From ad8d1eb3d1b04441826998993fad809ac0e939cf Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Fri, 5 Jul 2024 08:26:51 -0400 Subject: [PATCH] Minor: Add `ConstExpr`::from, use in places --- .../sort_preserving_repartition_fuzz.rs | 2 +- .../physical-expr/src/equivalence/class.rs | 28 +++++++++++++++++++ .../physical-expr/src/equivalence/mod.rs | 2 +- .../physical-expr/src/equivalence/ordering.rs | 6 ++-- .../src/equivalence/properties.rs | 22 +++++++-------- datafusion/physical-plan/src/filter.rs | 6 ++-- datafusion/physical-plan/src/union.rs | 2 +- datafusion/physical-plan/src/windows/mod.rs | 4 +-- 8 files changed, 47 insertions(+), 25 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs b/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs index f00d17a06ffc..ceae13a469f0 100644 --- a/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs @@ -80,7 +80,7 @@ mod sp_repartition_fuzz_tests { // Define a and f are aliases eq_properties.add_equal_conditions(col_a, col_f)?; // Column e has constant value. - eq_properties = eq_properties.add_constants([ConstExpr::new(col_e.clone())]); + eq_properties = eq_properties.add_constants([ConstExpr::from(col_e)]); // Randomly order columns for sorting let mut rng = StdRng::seed_from_u64(seed); diff --git a/datafusion/physical-expr/src/equivalence/class.rs b/datafusion/physical-expr/src/equivalence/class.rs index 4c0edd2a5d9a..e483f935b75c 100644 --- a/datafusion/physical-expr/src/equivalence/class.rs +++ b/datafusion/physical-expr/src/equivalence/class.rs @@ -42,12 +42,28 @@ use datafusion_common::JoinType; /// - `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. +/// +/// # Example +/// +/// ```rust +/// # use datafusion_physical_expr::ConstExpr; +/// # use datafusion_physical_expr_common::expressions::lit; +/// let col = lit(5); +/// // Create a constant expression from a physical expression ref +/// let const_expr = ConstExpr::from(&col); +/// // create a constant expression from a physical expression +/// let const_expr = ConstExpr::from(col); +/// ``` pub struct ConstExpr { expr: Arc, across_partitions: bool, } impl ConstExpr { + /// Create a new constant expression from a physical expression. + /// + /// Note you can also use `ConstExpr::from` to create a constant expression + /// from a reference as well pub fn new(expr: Arc) -> Self { Self { expr, @@ -85,6 +101,18 @@ impl ConstExpr { } } +impl From> for ConstExpr { + fn from(expr: Arc) -> Self { + Self::new(expr) + } +} + +impl From<&Arc> for ConstExpr { + fn from(expr: &Arc) -> Self { + Self::new(Arc::clone(expr)) + } +} + /// Checks whether `expr` is among in the `const_exprs`. pub fn const_exprs_contains( const_exprs: &[ConstExpr], diff --git a/datafusion/physical-expr/src/equivalence/mod.rs b/datafusion/physical-expr/src/equivalence/mod.rs index 1ed9a4ac217f..83f94057f740 100644 --- a/datafusion/physical-expr/src/equivalence/mod.rs +++ b/datafusion/physical-expr/src/equivalence/mod.rs @@ -205,7 +205,7 @@ mod tests { // Define a and f are aliases eq_properties.add_equal_conditions(col_a, col_f)?; // Column e has constant value. - eq_properties = eq_properties.add_constants([ConstExpr::new(Arc::clone(col_e))]); + eq_properties = eq_properties.add_constants([ConstExpr::from(col_e)]); // Randomly order columns for sorting let mut rng = StdRng::seed_from_u64(seed); diff --git a/datafusion/physical-expr/src/equivalence/ordering.rs b/datafusion/physical-expr/src/equivalence/ordering.rs index d71075dc77e1..c4b8a5c46563 100644 --- a/datafusion/physical-expr/src/equivalence/ordering.rs +++ b/datafusion/physical-expr/src/equivalence/ordering.rs @@ -556,9 +556,9 @@ mod tests { let eq_group = EquivalenceGroup::new(eq_group); eq_properties.add_equivalence_group(eq_group); - let constants = constants.into_iter().map(|expr| { - ConstExpr::new(Arc::clone(expr)).with_across_partitions(true) - }); + let constants = constants + .into_iter() + .map(|expr| ConstExpr::from(expr).with_across_partitions(true)); eq_properties = eq_properties.add_constants(constants); let reqs = convert_to_sort_exprs(&reqs); diff --git a/datafusion/physical-expr/src/equivalence/properties.rs b/datafusion/physical-expr/src/equivalence/properties.rs index 9a6a17f58c1f..d9d19c0bcf47 100644 --- a/datafusion/physical-expr/src/equivalence/properties.rs +++ b/datafusion/physical-expr/src/equivalence/properties.rs @@ -213,13 +213,13 @@ impl EquivalenceProperties { // Left expression is constant, add right as constant if !const_exprs_contains(&self.constants, right) { self.constants - .push(ConstExpr::new(Arc::clone(right)).with_across_partitions(true)); + .push(ConstExpr::from(right).with_across_partitions(true)); } } else if self.is_expr_constant(right) { // Right expression is constant, add left as constant if !const_exprs_contains(&self.constants, left) { self.constants - .push(ConstExpr::new(Arc::clone(left)).with_across_partitions(true)); + .push(ConstExpr::from(left).with_across_partitions(true)); } } @@ -300,7 +300,7 @@ impl EquivalenceProperties { { if !const_exprs_contains(&self.constants, &expr) { let const_expr = - ConstExpr::new(expr).with_across_partitions(across_partitions); + ConstExpr::from(expr).with_across_partitions(across_partitions); self.constants.push(const_expr); } } @@ -404,7 +404,7 @@ impl EquivalenceProperties { // we add column `a` as constant to the algorithm state. This enables us // to deduce that `(b + c) ASC` is satisfied, given `a` is constant. eq_properties = eq_properties - .add_constants(std::iter::once(ConstExpr::new(normalized_req.expr))); + .add_constants(std::iter::once(ConstExpr::from(normalized_req.expr))); } true } @@ -832,9 +832,8 @@ impl EquivalenceProperties { && !const_exprs_contains(&projected_constants, target) { // Expression evaluates to single value - projected_constants.push( - ConstExpr::new(Arc::clone(target)).with_across_partitions(true), - ); + projected_constants + .push(ConstExpr::from(target).with_across_partitions(true)); } } projected_constants @@ -927,8 +926,8 @@ impl EquivalenceProperties { // 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 { - eq_properties = eq_properties - .add_constants(std::iter::once(ConstExpr::new(Arc::clone(expr)))); + eq_properties = + eq_properties.add_constants(std::iter::once(ConstExpr::from(expr))); search_indices.shift_remove(idx); } // Add new ordered section to the state. @@ -2147,8 +2146,7 @@ mod tests { let col_h = &col("h", &test_schema)?; // Add column h as constant - eq_properties = - eq_properties.add_constants(vec![ConstExpr::new(Arc::clone(col_h))]); + eq_properties = eq_properties.add_constants(vec![ConstExpr::from(col_h)]); let test_cases = vec![ // TEST CASE 1 @@ -2458,7 +2456,7 @@ mod tests { for case in cases { let mut properties = base_properties .clone() - .add_constants(case.constants.into_iter().map(ConstExpr::new)); + .add_constants(case.constants.into_iter().map(ConstExpr::from)); for [left, right] in &case.equal_conditions { properties.add_equal_conditions(left, right)? } diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index c141958c1171..ab7a63e44550 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -173,13 +173,11 @@ impl FilterExec { // Filter evaluates to single value for all partitions if input_eqs.is_expr_constant(binary.left()) { res_constants.push( - ConstExpr::new(binary.right().clone()) - .with_across_partitions(true), + ConstExpr::from(binary.right()).with_across_partitions(true), ) } else if input_eqs.is_expr_constant(binary.right()) { res_constants.push( - ConstExpr::new(binary.left().clone()) - .with_across_partitions(true), + ConstExpr::from(binary.left()).with_across_partitions(true), ) } } diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 3f88eb4c3732..867cddeb7b41 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -188,7 +188,7 @@ fn calculate_union_eq_properties( // TODO: Check whether constant expressions evaluates the same value or not for each partition let across_partitions = false; return Some( - ConstExpr::new(meet_constant.owned_expr()) + ConstExpr::from(meet_constant.owned_expr()) .with_across_partitions(across_partitions), ); } diff --git a/datafusion/physical-plan/src/windows/mod.rs b/datafusion/physical-plan/src/windows/mod.rs index d798eaacc787..0622aad74cad 100644 --- a/datafusion/physical-plan/src/windows/mod.rs +++ b/datafusion/physical-plan/src/windows/mod.rs @@ -598,9 +598,7 @@ pub fn get_window_mode( options: None, })); // Treat partition by exprs as constant. During analysis of requirements are satisfied. - let const_exprs = partitionby_exprs - .iter() - .map(|expr| ConstExpr::new(expr.clone())); + let const_exprs = partitionby_exprs.iter().map(ConstExpr::from); let partition_by_eqs = input_eqs.add_constants(const_exprs); let order_by_reqs = PhysicalSortRequirement::from_sort_exprs(orderby_keys); let reverse_order_by_reqs =