improve lambdas

Author: gstvg

datafusion/expr/src/expr.rs

@@ -3411,7 +3411,7 @@ impl Display for SchemaDisplay<'_> {
                 )
             }
             Expr::LambdaVariable(c) => {
-                write!(f, "{}", c.name)
+                f.write_str(&c.name)
             }
         }
     }

datafusion/expr/src/udlf.rs

@@ -19,14 +19,12 @@
 
 use crate::expr::schema_name_from_exprs_comma_separated_without_space;
 use crate::simplify::{ExprSimplifyResult, SimplifyContext};
-use crate::sort_properties::{ExprProperties, SortProperties};
 use crate::{ColumnarValue, Documentation, Expr};
 use arrow::array::{ArrayRef, RecordBatch};
 use arrow::datatypes::{DataType, Field, FieldRef, Schema};
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::{exec_err, not_impl_err, Result, ScalarValue};
 use datafusion_expr_common::dyn_eq::{DynEq, DynHash};
-use datafusion_expr_common::interval_arithmetic::Interval;
 use datafusion_expr_common::signature::Volatility;
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 use std::any::Any;
@@ -602,83 +600,6 @@ pub trait LambdaUDF: Debug + DynEq + DynHash + Send + Sync {
         }
     }
 
-    /// Computes the output [`Interval`] for a [`LambdaUDF`], given the input
-    /// intervals.
-    ///
-    /// # Parameters
-    ///
-    /// * `children` are the intervals for the children (inputs) of this function.
-    ///
-    /// # Example
-    ///
-    /// If the function is `ABS(a)`, and the input interval is `a: [-3, 2]`,
-    /// then the output interval would be `[0, 3]`.
-    fn evaluate_bounds(&self, _input: &[&Interval]) -> Result<Interval> {
-        // We cannot assume the input datatype is the same of output type.
-        Interval::make_unbounded(&DataType::Null)
-    }
-
-    /// Updates bounds for child expressions, given a known [`Interval`]s for this
-    /// function.
-    ///
-    /// This function is used to propagate constraints down through an
-    /// expression tree.
-    ///
-    /// # Parameters
-    ///
-    /// * `interval` is the currently known interval for this function.
-    /// * `inputs` are the current intervals for the inputs (children) of this function.
-    ///
-    /// # Returns
-    ///
-    /// A `Vec` of new intervals for the children, in order.
-    ///
-    /// If constraint propagation reveals an infeasibility for any child, returns
-    /// [`None`]. If none of the children intervals change as a result of
-    /// propagation, may return an empty vector instead of cloning `children`.
-    /// This is the default (and conservative) return value.
-    ///
-    /// # Example
-    ///
-    /// If the function is `ABS(a)`, the current `interval` is `[4, 5]` and the
-    /// input `a` is given as `[-7, 3]`, then propagation would return `[-5, 3]`.
-    fn propagate_constraints(
-        &self,
-        _interval: &Interval,
-        _inputs: &[&Interval],
-    ) -> Result<Option<Vec<Interval>>> {
-        Ok(Some(vec![]))
-    }
-
-    /// Calculates the [`SortProperties`] of this function based on its children's properties.
-    fn output_ordering(&self, inputs: &[ExprProperties]) -> Result<SortProperties> {
-        if !self.preserves_lex_ordering(inputs)? {
-            return Ok(SortProperties::Unordered);
-        }
-
-        let Some(first_order) = inputs.first().map(|p| &p.sort_properties) else {
-            return Ok(SortProperties::Singleton);
-        };
-
-        if inputs
-            .iter()
-            .skip(1)
-            .all(|input| &input.sort_properties == first_order)
-        {
-            Ok(*first_order)
-        } else {
-            Ok(SortProperties::Unordered)
-        }
-    }
-
-    /// Returns true if the function preserves lexicographical ordering based on
-    /// the input ordering.
-    ///
-    /// For example, `concat(a || b)` preserves lexicographical ordering, but `abs(a)` does not.
-    fn preserves_lex_ordering(&self, _inputs: &[ExprProperties]) -> Result<bool> {
-        Ok(false)
-    }
-
     /// Coerce arguments of a function call to types that the function can evaluate.
     ///
     /// See the [type coercion module](crate::type_coercion)

datafusion/optimizer/src/common_subexpr_eliminate.rs

@@ -30,7 +30,7 @@ use datafusion_common::alias::AliasGenerator;
 use datafusion_common::cse::{CSE, CSEController, FoundCommonNodes};
 use datafusion_common::tree_node::{Transformed, TreeNode};
 use datafusion_common::{Column, DFSchema, DFSchemaRef, Result, qualified_name};
-use datafusion_expr::expr::{Alias, ScalarFunction};
+use datafusion_expr::expr::{Alias, LambdaFunction, ScalarFunction};
 use datafusion_expr::logical_plan::{
     Aggregate, Filter, LogicalPlan, Projection, Sort, Window,
 };
@@ -651,12 +651,15 @@ impl CSEController for ExprCSEController<'_> {
 
     fn conditional_children(node: &Expr) -> Option<(Vec<&Expr>, Vec<&Expr>)> {
         match node {
-            // In case of `ScalarFunction`s we don't know which children are surely
+            // In case of `ScalarFunction`s and `LambdaFunction`s we don't know which children are surely
             // executed so start visiting all children conditionally and stop the
             // recursion with `TreeNodeRecursion::Jump`.
             Expr::ScalarFunction(ScalarFunction { func, args }) => {
                 func.conditional_arguments(args)
             }
+            Expr::LambdaFunction(LambdaFunction { func, args }) => {
+                func.conditional_arguments(args)
+            }
 
             // In case of `And` and `Or` the first child is surely executed, but we
             // account subexpressions as conditional in the second.
@@ -696,7 +699,8 @@ impl CSEController for ExprCSEController<'_> {
     }
 
     fn is_valid(node: &Expr) -> bool {
-        !node.is_volatile_node() && !matches!(node, Expr::LambdaVariable(_))
+        !node.is_volatile_node()
+            && !matches!(node, Expr::Lambda(_) | Expr::LambdaVariable(_))
     }
 
     fn is_ignored(&self, node: &Expr) -> bool {
@@ -726,6 +730,7 @@ impl CSEController for ExprCSEController<'_> {
                 | Expr::ScalarVariable(..)
                 | Expr::Alias(..)
                 | Expr::Wildcard { .. }
+                | Expr::Lambda(_)
                 | Expr::LambdaVariable(_)
         );
 

datafusion/physical-expr/src/expressions/lambda.rs

@@ -27,13 +27,13 @@ use arrow::{
     datatypes::{DataType, Schema},
     record_batch::RecordBatch,
 };
-use datafusion_common::{internal_err, tree_node::TreeNodeVisitor, HashSet, Result};
+use datafusion_common::{HashSet, Result, internal_err, tree_node::TreeNodeVisitor};
 use datafusion_common::{
     plan_err,
     tree_node::{TreeNode, TreeNodeRecursion},
 };
 use datafusion_expr::ColumnarValue;
-use hashbrown::{hash_map::EntryRef, HashMap};
+use hashbrown::{HashMap, hash_map::EntryRef};
 
 /// Represents a lambda with the given parameters names and body
 #[derive(Debug, Eq, Clone)]
@@ -100,6 +100,14 @@ impl LambdaExpr {
         &self.captured_columns
     }
 
+    /// Returns lambdas variables names that aren't of this lambda nor any other lambda down tree.
+    /// Example:
+    ///
+    /// `array_transform([[[1, 2, 3]]], a -> array_transform(a, b -> array_transform(b, c -> length(a) + length(b) + c)))`
+    ///
+    /// For the outermost lambda, this would return an empty hash set
+    /// For the middle one, `HashSet("a")`
+    /// And for the innermost, `HashSet("a", "b")`
     pub(crate) fn captured_variables(&self) -> &HashSet<String> {
         &self.captured_variables
     }
@@ -192,7 +200,7 @@ impl<'n> TreeNodeVisitor<'n> for Captures<'n> {
     fn f_down(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
         if let Some(lambda) = node.as_any().downcast_ref::<LambdaExpr>() {
             for param in &lambda.params {
-                *self.shadows.entry_ref(param.as_str()).or_default() += 1;
+                *self.shadows.entry(param.as_str()).or_default() += 1;
             }
         } else if let Some(lambda_variable) =
             node.as_any().downcast_ref::<LambdaVariable>()
@@ -230,3 +238,101 @@ impl<'n> TreeNodeVisitor<'n> for Captures<'n> {
         Ok(TreeNodeRecursion::Continue)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        LambdaFunctionExpr,
+        expressions::{Column, LambdaExpr, NoOp, lambda::lambda, lambda_variable},
+    };
+    use arrow::{
+        array::RecordBatch,
+        datatypes::{DataType, Field, FieldRef, Schema},
+    };
+    use datafusion_common::{HashSet, Result};
+    use datafusion_expr::{ColumnarValue, LambdaUDF};
+    use std::sync::Arc;
+
+    #[derive(Debug, Hash, Eq, PartialEq)]
+    struct DummyLambdaUDF;
+
+    impl LambdaUDF for DummyLambdaUDF {
+        fn as_any(&self) -> &dyn std::any::Any {
+            unimplemented!()
+        }
+
+        fn name(&self) -> &str {
+            "dummy_udlf"
+        }
+
+        fn signature(&self) -> &datafusion_expr::LambdaSignature {
+            unimplemented!()
+        }
+
+        fn lambdas_parameters(
+            &self,
+            _args: &[datafusion_expr::ValueOrLambda<FieldRef, ()>],
+        ) -> Result<Vec<Option<Vec<Field>>>> {
+            unimplemented!()
+        }
+
+        fn return_field_from_args(
+            &self,
+            _args: datafusion_expr::LambdaReturnFieldArgs,
+        ) -> Result<FieldRef> {
+            unimplemented!()
+        }
+
+        fn invoke_with_args(
+            &self,
+            _args: datafusion_expr::LambdaFunctionArgs,
+        ) -> Result<ColumnarValue> {
+            unimplemented!()
+        }
+    }
+
+    #[test]
+    fn test_lambda_captures() {
+        let null_field = Arc::new(Field::new("", DataType::Null, true));
+
+        //`var_b -> dummy_udlf(var_a, var_b, column@0, var_c -> var_c))`
+        let inner = LambdaExpr::try_new(
+            vec![String::from("var_b")],
+            Arc::new(LambdaFunctionExpr::new(
+                "dummy_udlf",
+                Arc::new(DummyLambdaUDF),
+                vec![
+                    lambda_variable("var_a", Arc::clone(&null_field)).unwrap(),
+                    lambda_variable("var_b", Arc::clone(&null_field)).unwrap(),
+                    Arc::new(Column::new("column", 0)),
+                    lambda(
+                        ["var_c"],
+                        lambda_variable("var_c", Arc::clone(&null_field)).unwrap(),
+                    )
+                    .unwrap(),
+                ],
+                Arc::clone(&null_field),
+                Arc::new(Default::default()),
+            )),
+        )
+        .unwrap();
+
+        assert_eq!(inner.captured_columns(), &HashSet::from([0]));
+        assert_eq!(
+            inner.captured_variables(),
+            &HashSet::from([String::from("var_a")])
+        );
+    }
+
+    #[test]
+    fn test_lambda_evaluate() {
+        let lambda = lambda(["a"], Arc::new(NoOp::new())).unwrap();
+        let batch = RecordBatch::new_empty(Arc::new(Schema::empty()));
+        assert!(lambda.evaluate(&batch).is_err());
+    }
+
+    #[test]
+    fn test_lambda_duplicate_name() {
+        assert!(lambda(["a", "a"], Arc::new(NoOp::new())).is_err());
+    }
+}

datafusion/physical-expr/src/expressions/mod.rs

@@ -23,11 +23,11 @@ mod case;
 mod cast;
 mod cast_column;
 mod column;
-mod lambda_variable;
 mod dynamic_filters;
 mod in_list;
 mod is_not_null;
 mod is_null;
+mod lambda_variable;
 mod lambda;
 mod like;
 mod literal;

datafusion/physical-expr/src/lambda_function.rs

@@ -41,8 +41,6 @@ use arrow::array::{Array, NullArray, RecordBatch};
 use arrow::datatypes::{DataType, Field, FieldRef, Schema};
 use datafusion_common::config::{ConfigEntry, ConfigOptions};
 use datafusion_common::{Result, ScalarValue, exec_err, internal_err};
-use datafusion_expr::interval_arithmetic::Interval;
-use datafusion_expr::sort_properties::ExprProperties;
 use datafusion_expr::type_coercion::functions::value_fields_with_lambda_udf;
 use datafusion_expr::{
     ColumnarValue, LambdaArgument, LambdaFunctionArgs, LambdaReturnFieldArgs, LambdaUDF,
@@ -243,7 +241,6 @@ fn sorted_config_entries(config_options: &ConfigOptions) -> Vec<ConfigEntry> {
 }
 
 impl PhysicalExpr for LambdaFunctionExpr {
-    /// Return a reference to Any that can be used for downcasting
     fn as_any(&self) -> &dyn Any {
         self
     }
@@ -404,34 +401,6 @@ impl PhysicalExpr for LambdaFunctionExpr {
         )))
     }
 
-    fn evaluate_bounds(&self, children: &[&Interval]) -> Result<Interval> {
-        self.fun.evaluate_bounds(children)
-    }
-
-    fn propagate_constraints(
-        &self,
-        interval: &Interval,
-        children: &[&Interval],
-    ) -> Result<Option<Vec<Interval>>> {
-        self.fun.propagate_constraints(interval, children)
-    }
-
-    fn get_properties(&self, children: &[ExprProperties]) -> Result<ExprProperties> {
-        let sort_properties = self.fun.output_ordering(children)?;
-        let preserves_lex_ordering = self.fun.preserves_lex_ordering(children)?;
-        let children_range = children
-            .iter()
-            .map(|props| &props.range)
-            .collect::<Vec<_>>();
-        let range = self.fun().evaluate_bounds(&children_range)?;
-
-        Ok(ExprProperties {
-            sort_properties,
-            range,
-            preserves_lex_ordering,
-        })
-    }
-
     fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result {
         write!(f, "{}(", self.name)?;
         for (i, expr) in self.args.iter().enumerate() {