diff --git a/datafusion/expr/src/higher_order_function.rs b/datafusion/expr/src/higher_order_function.rs index c300be8f6fcfe..fc0e5c8f71090 100644 --- a/datafusion/expr/src/higher_order_function.rs +++ b/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>, ) -> Result> { - 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; diff --git a/datafusion/functions-nested/src/array_any_match.rs b/datafusion/functions-nested/src/array_any_match.rs index 0f620f18bd8f2..97029eb422058 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 { + 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::::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::::from_iter([ + "a", "b", "c", "a", + ])); + let result = run_any_match_eq(values, lit("c"))?; + assert_eq!( + result.as_any().downcast_ref::().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::().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 d11c3e7435fde..685992a0672d8 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/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))) })? diff --git a/datafusion/substrait/tests/cases/logical_plans.rs b/datafusion/substrait/tests/cases/logical_plans.rs index 522381de6efdf..d1c930d297909 100644 --- a/datafusion/substrait/tests/cases/logical_plans.rs +++ b/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(()) + } } 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 0000000000000..0b0b144c84a98 --- /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 0000000000000..5ee822f6f19ae --- /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" + ] + } + } + ] +} diff --git a/datafusion/substrait/tests/testdata/test_plans/nested_any_match_dictionary_list_element.substrait.json b/datafusion/substrait/tests/testdata/test_plans/nested_any_match_dictionary_list_element.substrait.json new file mode 100644 index 0000000000000..cfb0c3461bf5c --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/nested_any_match_dictionary_list_element.substrait.json @@ -0,0 +1,208 @@ +{ + "version": { + "minorNumber": 85, + "producer": "datafusion" + }, + "extensions": [ + { + "extensionFunction": { + "extensionUrnReference": 4294967295, + "functionAnchor": 1, + "name": "equal" + } + }, + { + "extensionFunction": { + "extensionUrnReference": 4294967295, + "functionAnchor": 2, + "name": "array_any_match" + } + }, + { + "extensionFunction": { + "extensionUrnReference": 4294967295, + "name": "array_transform" + } + } + ], + "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": 2, + "arguments": [ + { + "value": { + "scalarFunction": { + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": { + "field": 1 + } + }, + "rootReference": {} + } + } + }, + { + "value": { + "lambda": { + "parameters": { + "types": [ + { + "string": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + }, + "body": { + "selection": { + "directReference": { + "structField": {} + }, + "lambdaParameterReference": {} + } + } + } + } + } + ], + "outputType": { + "list": { + "type": { + "string": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + }, + { + "value": { + "lambda": { + "parameters": { + "types": [ + { + "string": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + }, + "body": { + "scalarFunction": { + "functionReference": 1, + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": {} + }, + "lambdaParameterReference": {} + } + } + }, + { + "value": { + "literal": { + "string": "c" + } + } + } + ], + "outputType": { + "bool": { + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + } + } + } + ], + "outputType": { + "bool": { + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + ] + } + }, + "names": [ + "matched" + ] + } + } + ] +}