fix: cast lambda variables to their declared type instead of erroring…

datafusion/expr/src/higher_order_function.rs

@@ -1219,24 +1219,25 @@ fn resolve_higher_order_function(
     // a map of lambda variable name => a never empty stack of fields [ [..shadowed], in_scope ]
     vars: &mut HashMap<String, Vec<FieldRef>>,
 ) -> Result<Transformed<Expr>> {
-    let args = if !vars.is_empty() {
-        /*  if this is a nested lambda, we must resolve non-lambda args before invoking
-            lambda_parameters because it will invoke ExprSchemable::to_field for every
-            non-lambda parameter, and if one them contains a lambda variable, it will fail
-            due to it being unresolved. Example query:
+    /*  non-lambda args are resolved before invoking lambda_parameters because it will
+        invoke ExprSchemable::to_field for every non-lambda parameter, and if one of them
+        contains a lambda variable, it will fail due to it being unresolved. Example query:
 
-            array_transform([[1, 2]], a -> array_transform(a, b -> b+1))
+        array_transform([[1, 2]], a -> array_transform(a, b -> b+1))
 
-            the nested array_transform's lambda_parameters will call Lambdavariable::to_field
-            on it's first argument, the variable `a`, which must be resolved
-        */
-        args.map_elements(|arg| match arg {
-            Expr::Lambda(_) => Ok(Transformed::no(arg)),
-            _ => resolve_lambda_variables(arg, schema, vars),
-        })?
-    } else {
-        Transformed::no(args)
-    };
+        the nested array_transform's lambda_parameters will call Lambdavariable::to_field
+        on it's first argument, the variable `a`, which must be resolved
+
+        They are also resolved when no variable is in scope yet, as an argument can be a
+        higher order function of its own, whose variables this is the only chance to
+        resolve: the traversal jumps over the arguments once this returns. Example query:
+
+        array_any_match(array_transform(tags, x -> x), y -> y = 'c')
+    */
+    let args = args.map_elements(|arg| match arg {
+        Expr::Lambda(_) => Ok(Transformed::no(arg)),
+        _ => resolve_lambda_variables(arg, schema, vars),
+    })?;
 
     let transformed = args.transformed;
     let mut args = args.data;

datafusion/functions-nested/src/array_any_match.rs

@@ -191,11 +191,14 @@ mod tests {
     use std::{collections::HashMap, sync::Arc};
 
     use arrow::{
-        array::{ArrayRef, BooleanArray, Int32Array, ListArray, RecordBatch},
+        array::{
+            ArrayRef, BooleanArray, DictionaryArray, Int32Array, ListArray, RecordBatch,
+            StringViewArray,
+        },
         buffer::{NullBuffer, OffsetBuffer},
-        datatypes::{DataType, Field},
+        datatypes::{DataType, Field, UInt32Type},
     };
-    use datafusion_common::{DFSchema, Result};
+    use datafusion_common::{DFSchema, Result, ScalarValue};
     use datafusion_expr::{
         Expr, HigherOrderReturnFieldArgs, HigherOrderUDFImpl, ValueOrLambda, col,
         execution_props::ExecutionProps,
@@ -490,4 +493,74 @@ mod tests {
         );
         Ok(())
     }
+
+    /// Evaluates `any_match(list, x -> x = needle)` over two rows, `[a, b]` and
+    /// `[c, a]`, built from `values`, with the lambda variable declared as the element
+    /// type. `needle` carries its own type, as these expressions are planned directly
+    /// and so are never type coerced.
+    fn run_any_match_eq(values: ArrayRef, needle: Expr) -> Result<ArrayRef> {
+        let element_type = values.data_type().clone();
+        let element_field = Arc::new(Field::new_list_field(element_type.clone(), true));
+        let list = ListArray::new(
+            Arc::clone(&element_field),
+            OffsetBuffer::<i32>::from_lengths(vec![2, 2]),
+            values,
+            None,
+        );
+
+        let schema = DFSchema::from_unqualified_fields(
+            vec![Field::new("list", DataType::List(element_field), true)].into(),
+            HashMap::new(),
+        )?;
+
+        create_physical_expr(
+            &Expr::HigherOrderFunction(HigherOrderFunction::new(
+                array_any_match_higher_order_function(),
+                vec![
+                    col("list"),
+                    lambda(
+                        ["x"],
+                        Expr::LambdaVariable(LambdaVariable::new(
+                            "x".to_string(),
+                            Some(Arc::new(Field::new("x", element_type, true))),
+                        ))
+                        .eq(needle),
+                    ),
+                ],
+            )),
+            &schema,
+            &ExecutionProps::new(),
+            &PhysicalPlanningContext::default(),
+        )?
+        .evaluate(&RecordBatch::try_new(
+            Arc::clone(schema.inner()),
+            vec![Arc::new(list)],
+        )?)?
+        .into_array(2)
+    }
+
+    #[test]
+    fn test_any_match_on_dictionary_encoded_elements() -> Result<()> {
+        let values: ArrayRef = Arc::new(DictionaryArray::<UInt32Type>::from_iter([
+            "a", "b", "c", "a",
+        ]));
+        let result = run_any_match_eq(values, lit("c"))?;
+        assert_eq!(
+            result.as_any().downcast_ref::<BooleanArray>().unwrap(),
+            &BooleanArray::from(vec![Some(false), Some(true)])
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn test_any_match_on_string_view_elements() -> Result<()> {
+        let values: ArrayRef = Arc::new(StringViewArray::from(vec!["a", "b", "c", "a"]));
+        let needle = lit(ScalarValue::Utf8View(Some("c".to_string())));
+        let result = run_any_match_eq(values, needle)?;
+        assert_eq!(
+            result.as_any().downcast_ref::<BooleanArray>().unwrap(),
+            &BooleanArray::from(vec![Some(false), Some(true)])
+        );
+        Ok(())
+    }
 }

datafusion/optimizer/src/analyzer/type_coercion.rs

@@ -163,6 +163,16 @@ fn analyze_internal(
     // apply coercion rewrite all expressions in the plan individually
     plan.map_expressions(|expr| {
         let original_name = name_preserver.save(&expr);
+
+        // A lambda variable carries the field recorded when the plan was built, which
+        // need not be the one its function derives from the arguments. Resolve them
+        // before coercing, so lambda bodies are coerced against the types they receive.
+        let expr = if expr.exists(|e| Ok(matches!(e, Expr::HigherOrderFunction(_))))? {
+            expr.resolve_lambda_variables(&schema)?.data
+        } else {
+            expr
+        };
+
         expr.rewrite(&mut expr_rewrite)
             .map(|transformed| transformed.update_data(|e| original_name.restore(e)))
     })?

datafusion/substrait/tests/cases/logical_plans.rs

@@ -340,4 +340,97 @@ mod tests {
         DataFrame::new(ctx.state(), plan).show().await?;
         Ok(())
     }
+
+    // The read schema declares the list element as dictionary encoded while the lambda
+    // carries its own plain parameter type, as a producer that cannot express dictionary
+    // encoding in Substrait emits. Physical planning reconciles the two with a cast.
+    #[tokio::test]
+    async fn higher_order_function_with_dictionary_encoded_lambda_parameter() -> Result<()>
+    {
+        let proto_plan = read_json(
+            "tests/testdata/test_plans/any_match_dictionary_list_element.substrait.json",
+        );
+        let ctx = add_plan_schemas_to_ctx(SessionContext::new(), &proto_plan)?;
+        let plan = from_substrait_plan(&ctx.state(), &proto_plan).await?;
+
+        assert_snapshot!(
+        plan,
+        @r#"
+        Projection: array_any_match(t.tags, (p0) -> p0 = Utf8("c")) AS matched
+          TableScan: t
+        "#
+        );
+
+        let df = DataFrame::new(ctx.state(), plan);
+
+        // Without resolution this fails in physical planning, where the field recorded
+        // on the parameter does not match the one the schema carries.
+        df.clone().show().await?;
+
+        // The lambda parameter is resolved to the encoding the list actually carries, so
+        // the comparison coerces the literal rather than decoding every element.
+        assert_snapshot!(
+        df.into_optimized_plan()?,
+        @r#"
+        Projection: array_any_match(t.tags, (p0) -> p0 = Dictionary(UInt32, Utf8("c"))) AS matched
+          TableScan: t projection=[tags]
+        "#
+        );
+        Ok(())
+    }
+
+    // Substrait has no string view type, so a producer records a plain parameter for a
+    // lambda over a column the plan itself declares as a view.
+    #[tokio::test]
+    async fn higher_order_function_with_string_view_lambda_parameter() -> Result<()> {
+        let proto_plan = read_json(
+            "tests/testdata/test_plans/any_match_string_view_list_element.substrait.json",
+        );
+        let ctx = add_plan_schemas_to_ctx(SessionContext::new(), &proto_plan)?;
+        let plan = from_substrait_plan(&ctx.state(), &proto_plan).await?;
+
+        let df = DataFrame::new(ctx.state(), plan);
+
+        // Without resolution this fails in physical planning, where the field recorded
+        // on the parameter does not match the one the schema carries.
+        df.clone().show().await?;
+
+        // The lambda parameter is resolved to the view the list actually carries, so the
+        // comparison coerces the literal rather than materializing every element.
+        assert_snapshot!(
+        df.into_optimized_plan()?,
+        @r#"
+        Projection: array_any_match(t.tags, (p0) -> p0 = Utf8View("c")) AS matched
+          TableScan: t projection=[tags]
+        "#
+        );
+        Ok(())
+    }
+
+    // A higher order function can be the value argument of another one, where its own
+    // lambda parameter has to be resolved too.
+    #[tokio::test]
+    async fn nested_higher_order_function_with_dictionary_encoded_lambda_parameter()
+    -> Result<()> {
+        let proto_plan = read_json(
+            "tests/testdata/test_plans/nested_any_match_dictionary_list_element.substrait.json",
+        );
+        let ctx = add_plan_schemas_to_ctx(SessionContext::new(), &proto_plan)?;
+        let plan = from_substrait_plan(&ctx.state(), &proto_plan).await?;
+
+        let df = DataFrame::new(ctx.state(), plan);
+
+        // Without resolution of the inner function this fails in physical planning.
+        df.clone().show().await?;
+
+        // Both parameters are resolved to the encoding the list actually carries.
+        assert_snapshot!(
+        df.into_optimized_plan()?,
+        @r#"
+        Projection: array_any_match(array_transform(t.tags, (p0) -> p0), (p1) -> p1 = Dictionary(UInt32, Utf8("c"))) AS matched
+          TableScan: t projection=[tags]
+        "#
+        );
+        Ok(())
+    }
 }