This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25418-63ab2a855eea79f33b11ebeb6408537b7017c54f in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit adb084051a8de6739726977201fab1f513cc30de Author: Lía Adriana <[email protected]> AuthorDate: Thu Sep 24 07:19:13 2026 +0000 fix: cast lambda variables to their declared type instead of erroring (#25418) ## Which issue does this PR close? - Closes #25411 ## Rationale for this change `create_physical_expr` required an `Expr::LambdaVariable`'s field to equal the planning schema's field and returned a plan error otherwise. The two can legitimately differ when a plan producer records a type that Substrait cannot express, such as dictionary encoding or a string view (Utf8view), since the recorded field and the field derived by `lambda_parameters` come from independent sources. ## What changes are included in this PR? Cast the variable to the recorded type instead, so the declared type is enforced rather than asserted. This is the same approach already used for `ScalarFunction` arguments, which the `TypeCoercion` analyzer casts to their coerced types instead of requiring them to match. Incompatible types still fail, as a cast error. ## What is the testing strategy for this PR? I added - Unit tests covering dictionary-encoded and `Utf8View` list elements with the lambda variable declared as plain `Utf8`, plus plan-shape tests pinning that the cast is inserted on a mismatch and omitted when the types already agree. - A Substrait integration test with a plan whose read schema declares a dictionary-encoded list element while the lambda carries its own plain `Utf8` parameter type. ## Are there any user-facing changes? No, this is a bug fix --- datafusion/functions-nested/src/array_any_match.rs | 79 ++++++++++- datafusion/optimizer/src/analyzer/type_coercion.rs | 10 ++ datafusion/substrait/tests/cases/logical_plans.rs | 66 +++++++++ ...ny_match_dictionary_list_element.substrait.json | 158 +++++++++++++++++++++ ...y_match_string_view_list_element.substrait.json | 147 +++++++++++++++++++ 5 files changed, 457 insertions(+), 3 deletions(-) diff --git a/datafusion/functions-nested/src/array_any_match.rs b/datafusion/functions-nested/src/array_any_match.rs index 0f620f18bd..97029eb422 100644 --- a/datafusion/functions-nested/src/array_any_match.rs +++ b/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(()) + } } diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index 1b2b2ae14c..c45bc5678a 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -155,6 +155,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))) })? diff --git a/datafusion/substrait/tests/cases/logical_plans.rs b/datafusion/substrait/tests/cases/logical_plans.rs index 522381de6e..6b92308a80 100644 --- a/datafusion/substrait/tests/cases/logical_plans.rs +++ b/datafusion/substrait/tests/cases/logical_plans.rs @@ -340,4 +340,70 @@ 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(()) + } } diff --git a/datafusion/substrait/tests/testdata/test_plans/any_match_dictionary_list_element.substrait.json b/datafusion/substrait/tests/testdata/test_plans/any_match_dictionary_list_element.substrait.json new file mode 100644 index 0000000000..0b0b144c84 --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/any_match_dictionary_list_element.substrait.json @@ -0,0 +1,158 @@ +{ + "version": { + "minorNumber": 85, + "producer": "datafusion" + }, + "extensions": [ + { + "extensionFunction": { + "extensionUrnReference": 4294967295, + "name": "equal" + } + }, + { + "extensionFunction": { + "extensionUrnReference": 4294967295, + "functionAnchor": 1, + "name": "array_any_match" + } + } + ], + "relations": [ + { + "root": { + "input": { + "project": { + "common": { + "emit": { + "outputMapping": [ + 2 + ] + } + }, + "input": { + "read": { + "baseSchema": { + "names": [ + "id", + "tags" + ], + "struct": { + "types": [ + { + "i32": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "list": { + "type": { + "map": { + "key": { + "i32": { + "typeVariationReference": 1, + "nullability": "NULLABILITY_REQUIRED" + } + }, + "value": { + "string": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "typeVariationReference": 1, + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": [ + "t" + ] + } + } + }, + "expressions": [ + { + "scalarFunction": { + "functionReference": 1, + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": { + "field": 1 + } + }, + "rootReference": {} + } + } + }, + { + "value": { + "lambda": { + "parameters": { + "types": [ + { + "string": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + }, + "body": { + "scalarFunction": { + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": {} + }, + "lambdaParameterReference": {} + } + } + }, + { + "value": { + "literal": { + "string": "c" + } + } + } + ], + "outputType": { + "bool": { + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + } + } + } + ], + "outputType": { + "bool": { + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + ] + } + }, + "names": [ + "matched" + ] + } + } + ] +} diff --git a/datafusion/substrait/tests/testdata/test_plans/any_match_string_view_list_element.substrait.json b/datafusion/substrait/tests/testdata/test_plans/any_match_string_view_list_element.substrait.json new file mode 100644 index 0000000000..5ee822f6f1 --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/any_match_string_view_list_element.substrait.json @@ -0,0 +1,147 @@ +{ + "version": { + "minorNumber": 85, + "producer": "datafusion" + }, + "extensions": [ + { + "extensionFunction": { + "extensionUrnReference": 4294967295, + "name": "equal" + } + }, + { + "extensionFunction": { + "extensionUrnReference": 4294967295, + "functionAnchor": 1, + "name": "array_any_match" + } + } + ], + "relations": [ + { + "root": { + "input": { + "project": { + "common": { + "emit": { + "outputMapping": [ + 2 + ] + } + }, + "input": { + "read": { + "baseSchema": { + "names": [ + "id", + "tags" + ], + "struct": { + "types": [ + { + "i32": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "list": { + "type": { + "string": { + "typeVariationReference": 2, + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": [ + "t" + ] + } + } + }, + "expressions": [ + { + "scalarFunction": { + "functionReference": 1, + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": { + "field": 1 + } + }, + "rootReference": {} + } + } + }, + { + "value": { + "lambda": { + "parameters": { + "types": [ + { + "string": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + }, + "body": { + "scalarFunction": { + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": {} + }, + "lambdaParameterReference": {} + } + } + }, + { + "value": { + "literal": { + "string": "c" + } + } + } + ], + "outputType": { + "bool": { + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + } + } + } + ], + "outputType": { + "bool": { + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + ] + } + }, + "names": [ + "matched" + ] + } + } + ] +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
