This is an automated email from the ASF dual-hosted git repository.
mbutrovich pushed a commit to branch branch-54
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/branch-54 by this push:
new fcbc9bda1f [branch-54] Handle nulls in type coercion of higher-order
UDFs, map_extract, spark array_repeat (backport #23071) (#23629)
fcbc9bda1f is described below
commit fcbc9bda1fd004c349fa70f61f8c8b40448294ca
Author: gstvg <[email protected]>
AuthorDate: Sat Jul 18 08:29:16 2026 -0300
[branch-54] Handle nulls in type coercion of higher-order UDFs,
map_extract, spark array_repeat (backport #23071) (#23629)
## Which issue does this PR close?
- Backport of #23071 (clean cherry-pick of non-merge commits)
- Closes #23068 .
## Rationale for this change
Unresolved lambda variables and unspecified placeholders report a
`DataType::Null` data type, which causes an error when used as argument
of a function that doesn't handle nulls during type coercion
## What changes are included in this PR?
Handle null in `coerce_single_list_arg` for higher-order functions,
`map_extract` and in spark `array_repeat`. All other functions (scalar,
window, agg and higher-order) that uses user defined type coercion have
been checked to currently handle null args.
## Are these changes tested?
Test passing `array_filter` with an unresolved lambda variable to
`Dataframe::with_column` (which indirectly calls
`value_fields_with_higher_order_udf_and_lambdas` coercion via
`Projection::try_new`, `projection_schema`, `exprlist_to_fields`,
`Expr::to_field` ), and sqllogictests
## Are there any user-facing changes?
No
---
datafusion/core/tests/dataframe/mod.rs | 51 ++++++++++++++++++++--
datafusion/functions-nested/src/array_any_match.rs | 27 ++----------
datafusion/functions-nested/src/lambda_utils.rs | 1 +
datafusion/functions-nested/src/map_extract.rs | 10 +++++
datafusion/spark/src/function/array/repeat.rs | 8 ++--
.../test_files/array/array_any_match.slt | 6 +++
.../sqllogictest/test_files/array/array_filter.slt | 6 +++
.../test_files/array/array_transform.slt | 6 +++
datafusion/sqllogictest/test_files/map.slt | 6 +++
.../test_files/spark/array/array_repeat.slt | 7 +++
10 files changed, 98 insertions(+), 30 deletions(-)
diff --git a/datafusion/core/tests/dataframe/mod.rs
b/datafusion/core/tests/dataframe/mod.rs
index e55a373ada..807487bb88 100644
--- a/datafusion/core/tests/dataframe/mod.rs
+++ b/datafusion/core/tests/dataframe/mod.rs
@@ -39,6 +39,7 @@ use datafusion_functions_aggregate::expr_fn::{
array_agg, avg, avg_distinct, count, count_distinct, max, median, min, sum,
sum_distinct,
};
+use datafusion_functions_nested::expr_fn::{array_filter, array_transform,
make_array};
use datafusion_functions_nested::make_array::make_array_udf;
use datafusion_functions_window::expr_fn::{first_value, lead, row_number};
use insta::assert_snapshot;
@@ -78,8 +79,8 @@ use datafusion_expr::{
CreateMemoryTable, CreateView, DdlStatement, Expr, ExprFunctionExt,
ExprSchemable,
LogicalPlan, LogicalPlanBuilder, ScalarFunctionImplementation, SortExpr,
TableType,
WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition,
cast, col,
- create_udf, exists, in_subquery, lit, out_ref_col, placeholder,
scalar_subquery,
- when, wildcard,
+ create_udf, exists, in_subquery, lambda, lambda_var, lit, out_ref_col,
placeholder,
+ scalar_subquery, when, wildcard,
};
use datafusion_physical_expr::Partitioning;
use datafusion_physical_expr::aggregate::AggregateExprBuilder;
@@ -90,7 +91,9 @@ use datafusion_physical_plan::aggregates::{
AggregateExec, AggregateMode, PhysicalGroupBy,
};
use datafusion_physical_plan::empty::EmptyExec;
-use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties,
displayable};
+use datafusion_physical_plan::{
+ ExecutionPlan, ExecutionPlanProperties, collect, displayable,
+};
use datafusion::error::Result as DataFusionResult;
use datafusion::execution::options::JsonReadOptions;
@@ -7062,3 +7065,45 @@ async fn test_grouping_with_alias() -> Result<()> {
Ok(())
}
+
+#[tokio::test]
+async fn test_unresolved_lambda_variable() -> Result<()> {
+ let plan = table_with_mixed_lists()
+ .await?
+ .with_column(
+ "c",
+ array_transform(
+ make_array(vec![col("list")]),
+ lambda(
+ ["x"],
+ array_filter(
+ lambda_var("x"),
+ lambda(["y"], lambda_var("y").gt_eq(lit(2))),
+ ),
+ ),
+ ),
+ )?
+ .select_columns(&["list", "c"])?
+ .into_unoptimized_plan()
+ .resolve_lambda_variables()?
+ .data;
+
+ let session = SessionContext::new();
+ let exec = session.state().create_physical_plan(&plan).await?;
+ let context = session.task_ctx();
+ let results = collect(exec, context).await?;
+
+ let expected = [
+ "+-----------+----------+",
+ "| list | c |",
+ "+-----------+----------+",
+ "| [1, 2, 3] | [[2, 3]] |",
+ "| | [] |",
+ "| [] | [[]] |",
+ "| | [] |",
+ "+-----------+----------+",
+ ];
+ assert_batches_eq!(expected, &results);
+
+ Ok(())
+}
diff --git a/datafusion/functions-nested/src/array_any_match.rs
b/datafusion/functions-nested/src/array_any_match.rs
index a14cc8c260..e248dc95aa 100644
--- a/datafusion/functions-nested/src/array_any_match.rs
+++ b/datafusion/functions-nested/src/array_any_match.rs
@@ -37,6 +37,8 @@ use datafusion_expr::{
use datafusion_macros::user_doc;
use std::{fmt::Debug, sync::Arc};
+use crate::lambda_utils::coerce_single_list_arg;
+
make_higher_order_function_expr_and_func!(
ArrayAnyMatch,
array_any_match,
@@ -120,30 +122,7 @@ impl HigherOrderUDFImpl for ArrayAnyMatch {
}
fn coerce_value_types(&self, arg_types: &[DataType]) ->
Result<Vec<DataType>> {
- let [list] = arg_types else {
- return plan_err!(
- "{} function requires 1 value argument, got {}",
- self.name(),
- arg_types.len()
- );
- };
-
- let coerced = match list {
- DataType::List(_) | DataType::LargeList(_) => list.clone(),
- DataType::ListView(field) | DataType::FixedSizeList(field, _) => {
- DataType::List(Arc::clone(field))
- }
- DataType::LargeListView(field) =>
DataType::LargeList(Arc::clone(field)),
- _ => {
- return plan_err!(
- "{} expected a list as first argument, got {}",
- self.name(),
- list
- );
- }
- };
-
- Ok(vec![coerced])
+ coerce_single_list_arg(self.name(), arg_types)
}
fn lambda_parameters(
diff --git a/datafusion/functions-nested/src/lambda_utils.rs
b/datafusion/functions-nested/src/lambda_utils.rs
index 0f208ce5d2..ce596d0f1b 100644
--- a/datafusion/functions-nested/src/lambda_utils.rs
+++ b/datafusion/functions-nested/src/lambda_utils.rs
@@ -65,6 +65,7 @@ pub(crate) fn coerce_single_list_arg(
DataType::List(Arc::clone(field))
}
DataType::LargeListView(field) =>
DataType::LargeList(Arc::clone(field)),
+ DataType::Null => DataType::new_list(DataType::Null, true),
_ => return plan_err!("{name} expected a list as first argument, got
{list}"),
};
diff --git a/datafusion/functions-nested/src/map_extract.rs
b/datafusion/functions-nested/src/map_extract.rs
index aab0d013a4..45dc4b6f10 100644
--- a/datafusion/functions-nested/src/map_extract.rs
+++ b/datafusion/functions-nested/src/map_extract.rs
@@ -105,6 +105,11 @@ impl ScalarUDFImpl for MapExtract {
fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
let [map_type, _] = take_function_args(self.name(), arg_types)?;
+
+ if map_type.is_null() {
+ return Ok(DataType::Null);
+ }
+
let map_fields = get_map_entry_field(map_type)?;
Ok(DataType::List(Arc::new(Field::new_list_field(
map_fields.last().unwrap().data_type().clone(),
@@ -123,6 +128,10 @@ impl ScalarUDFImpl for MapExtract {
fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
let [map_type, _] = take_function_args(self.name(), arg_types)?;
+ if map_type.is_null() {
+ return Ok(arg_types.to_vec());
+ }
+
let field = get_map_entry_field(map_type)?;
Ok(vec![
map_type.clone(),
@@ -185,6 +194,7 @@ fn map_extract_inner(args: &[ArrayRef]) -> Result<ArrayRef>
{
let map_array = match map_arg.data_type() {
DataType::Map(_, _) => as_map_array(&map_arg)?,
+ DataType::Null => return Ok(Arc::clone(map_arg)),
_ => return exec_err!("The first argument in map_extract must be a
map"),
};
diff --git a/datafusion/spark/src/function/array/repeat.rs
b/datafusion/spark/src/function/array/repeat.rs
index da9b19a768..6effdf9a50 100644
--- a/datafusion/spark/src/function/array/repeat.rs
+++ b/datafusion/spark/src/function/array/repeat.rs
@@ -74,9 +74,11 @@ impl ScalarUDFImpl for SparkArrayRepeat {
// Coerce the second argument to Int64/UInt64 if it's a numeric type
let second = match second_type {
- DataType::Int8 | DataType::Int16 | DataType::Int32 |
DataType::Int64 => {
- DataType::Int64
- }
+ DataType::Int8
+ | DataType::Int16
+ | DataType::Int32
+ | DataType::Int64
+ | DataType::Null => DataType::Int64,
DataType::UInt8 | DataType::UInt16 | DataType::UInt32 |
DataType::UInt64 => {
DataType::UInt64
}
diff --git a/datafusion/sqllogictest/test_files/array/array_any_match.slt
b/datafusion/sqllogictest/test_files/array/array_any_match.slt
index 27f2a5339e..37aa47c55a 100644
--- a/datafusion/sqllogictest/test_files/array/array_any_match.slt
+++ b/datafusion/sqllogictest/test_files/array/array_any_match.slt
@@ -103,6 +103,12 @@ SELECT list_any_match([1, 2, 3], x -> x > 2);
----
true
+# null arg
+query B
+SELECT array_any_match(NULL, x -> x > 2);
+----
+NULL
+
statement ok
drop table t;
diff --git a/datafusion/sqllogictest/test_files/array/array_filter.slt
b/datafusion/sqllogictest/test_files/array/array_filter.slt
index f22cfb2198..9b564c5061 100644
--- a/datafusion/sqllogictest/test_files/array/array_filter.slt
+++ b/datafusion/sqllogictest/test_files/array/array_filter.slt
@@ -204,6 +204,12 @@ SELECT array_transform(array_filter(list, v -> v > 1), v
-> v * 3) FROM with_nul
[6]
NULL
+# null arg
+query ?
+SELECT array_filter(NULL, x -> x > 2);
+----
+NULL
+
statement ok
drop table t;
diff --git a/datafusion/sqllogictest/test_files/array/array_transform.slt
b/datafusion/sqllogictest/test_files/array/array_transform.slt
index c8c43588c8..5439d74411 100644
--- a/datafusion/sqllogictest/test_files/array/array_transform.slt
+++ b/datafusion/sqllogictest/test_files/array/array_transform.slt
@@ -393,6 +393,12 @@ physical_plan
02)--ProjectionExec: expr=[text@0 as text, list@1 as list, number@2 as number,
CASE WHEN number@2 > 30 THEN array_transform(make_array(make_array(list@1)),
(list) -> array_transform(list@3, (list) -> array_transform(list@4, (v) ->
number@2 + v@5 + array_element(list@4, 1)))) ELSE
array_transform(make_array(make_array(list@1)), (list) ->
array_transform(list@3, (list) -> array_transform(list@4, (v) -> number@2 +
array_element(list@4, 1)))) END as CASE WHEN t.number > Int64(30) THEN array_
[...]
03)----DataSourceExec: partitions=1, partition_sizes=[1]
+# null arg
+query ?
+SELECT array_transform(NULL, x -> x * 2);
+----
+NULL
+
query error
select array_transform();
----
diff --git a/datafusion/sqllogictest/test_files/map.slt
b/datafusion/sqllogictest/test_files/map.slt
index 62e70e6080..accf713544 100644
--- a/datafusion/sqllogictest/test_files/map.slt
+++ b/datafusion/sqllogictest/test_files/map.slt
@@ -642,6 +642,12 @@ select map_extract(MAP {1: 1, 2: 2, 3:3}, '1'),
map_extract(MAP {1: 1, 2: 2, 3:3
----
[1] [1] [1] [NULL] [1]
+# null arg
+query ?
+select map_extract(NULL, 'a');
+----
+NULL
+
# map_extract with columns
query ???
select map_extract(column1, 1), map_extract(column1, 5), map_extract(column1,
7) from map_array_table_1;
diff --git a/datafusion/sqllogictest/test_files/spark/array/array_repeat.slt
b/datafusion/sqllogictest/test_files/spark/array/array_repeat.slt
index 923e349140..d51767a264 100644
--- a/datafusion/sqllogictest/test_files/spark/array/array_repeat.slt
+++ b/datafusion/sqllogictest/test_files/spark/array/array_repeat.slt
@@ -112,3 +112,10 @@ FROM VALUES
[[123], [123]]
[[], []]
[[NULL], [NULL]]
+
+
+# null count
+query ?
+select array_repeat('a', NULL);
+----
+NULL
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
