[andygrove]

Which issue does this PR close?

Closes #7321.

Rationale for this change

filter(arr, x -> ...) ran natively only for x IS NOT NULL, and exists(arr, x -> ...) never did. Any other lambda went through the JVM codegen dispatcher, which evaluates it one element at a time like Spark does, plus the transfer of each batch, so such queries ran at about Spark's speed. DataFusion 55 evaluates a lambda once over all the elements of a batch, so Comet can use it when the body is made of natively supported expressions. This PR does that for filter, with array_filter, and for exists, with array_any_match, which returns what Spark's exists returns under three-valued logic.

Release build, Spark 4.1, 1M rows of ten-element arrays (every fifth element is NULL), best time:

Query	Spark	Comet before	Comet now
filter(strs, x -> nvl(x, '') != '')	616 ms	547 ms	307 ms
array_remove(array_compact(strs), '')	587 ms	252 ms	252 ms
filter(ints, x -> x IS NOT NULL)	286 ms	124 ms	125 ms
filter(ints, x -> x > threshold)	346 ms	335 ms	110 ms

exists over 1M rows of ten-element arrays (every fifth element is NULL), Spark 3.5, release build, best time:

Query	Spark	Comet before	Comet now
exists(ints, x -> x > 90)	266 ms	234 ms	108 ms
exists(dbls, x -> x IS NULL OR x < 0 OR x > 90 OR x = 0)	266 ms	225 ms	94 ms
exists(ints, x -> x > threshold)	251 ms	196 ms	105 ms
size(dbls) < 1 OR exists(dbls, x -> x > 90)	355 ms	403 ms	110 ms
exists(strs, x -> x IS NULL OR x = '')	357 ms	183 ms	168 ms

filter(ints, x -> nvl(x, '') != '') is not in the first table because it no longer takes the lambda path: Comet runs it as array_compact (see the last change below). In a separate local run (Spark 3.5.5, local[4], 2M rows of ten-element int arrays with every fifth element NULL, best of 5) it took 0.453 s on the dispatcher, 0.281 s with the native lambda alone, and 0.086 s as array_compact, against 0.220 s on Spark.

What changes are included in this PR?

• A HigherOrderFunc and a LambdaVariable expression in the protobuf, and a planner arm that builds DataFusion's HigherOrderFunctionExpr with a LambdaExpr. The lambda body is planned against the input schema plus the lambda parameter, so captured columns keep their indexes.
• CometArrayFilter and CometArrayExists share the code that serializes a lambda (the NativeLambdaSerde trait). They serialize the body natively when all of these hold: the lambda has one parameter, the array has a primitive, decimal, string, binary, date or timestamp element type, the body is deterministic, has no nested lambda and uses the default string collation, it reads no captured column that is not fixed-width, it holds no TRY cast that can make a null map key and no cast between a date or timestamp and a string, date or timestamp (the native cast fails beyond chrono's year range, Native timestamp casts fail or panic for years beyond chrono's range #7471), and every subexpression of it converts natively. Otherwise they keep using the dispatcher, as before. The planner arm maps array_filter and array_any_match to DataFusion's functions.
• exists has three more conditions. spark.sql.legacy.followThreeValuedLogicInArrayExists must be true (the default): with it false Spark returns false where array_any_match returns NULL, so exists stays on the dispatcher. No expression in the body may raise, whatever operators it uses: Spark stops at the first element that matches and never evaluates the body on the elements after it, while array_any_match evaluates it on every element, so with ANSI on exists(array(1, 0), x -> 100 / x > 1) would raise where Spark returns true. It is the same check as the AND/OR/IN rule below. And the array argument must not contain a TRY cast whose map key cast can fail (CometCast.canProduceNullMapKey): it makes a null key, which Arrow cannot hold (Codegen dispatcher writes a null map key as the key type's default value #6172). The dispatcher runs such an argument in Spark's own code, where the map never reaches Arrow.
• The JVM codegen dispatcher now declines a tree that reads a lambda parameter bound outside the tree. Its kernel is fed only columns, so such a parameter would read as NULL. A body that needs the dispatcher therefore fails to convert natively and the whole filter stays on the dispatcher.
• DataFusion copies each column the lambda captures once per element of the array before it evaluates the body, and those copies are not reserved with the memory pool. A body that reads a captured string, binary, array, struct or map column stays on the dispatcher. Numeric, boolean, date, timestamp and decimal captures run natively.
• A filter body that uses AND, OR or IN stays on the dispatcher unless no expression in it can raise. Spark evaluates those lazily and the native operators evaluate every operand, so x != 0 AND 100 / x > 1 would otherwise raise on the elements Spark skips. This is the per-row case in Track evaluation masks for data-dependent errors beyond unbase64 operator fallback #6006. The check classifies each expression and fails closed: arithmetic and casts raise only in the ANSI eval mode they carry, and element_at raises for the index 0 in any mode, so an expression that is not known to be safe counts as able to raise. CASE, IF and COALESCE match Spark natively and are not restricted.
• NullShortCircuit evaluates the later arguments of functions such as array_contains on a projection of just the columns they read. It now projects and re-indexes a LambdaVariable like a Column, which filter(ints, x -> array_contains(other, x)) needs. A lambda inside the projected expression binds its own parameters, so those are not columns of the batch and are left alone. Without that, a higher-order function inside an expression that is evaluated for only some rows, such as the right operand of AND or OR, or an argument of array_contains, failed with project index N out of bounds. case_when.rs treats a LambdaVariable like a Column when it decides whether an expression can fail.
• filter(arr, x -> nvl(x, '') != '') on an array of integers, floats, decimals, dates or timestamps runs as array_compact, the way x IS NOT NULL does. Spark plans it as NOT (coalesce(CAST(x AS STRING), '') = ''), in either operand order. No such cast is empty, so the predicate holds exactly when x is not NULL, and the native lambda would otherwise cast every element to a string first. A string element has no cast and stays on the lambda, as does any other default value, comparison or extra condition.
• spark.comet.exec.nativeLambda.enabled (default true) turns the native path off.
• A new "Expression Tuning" page in the user guide covers lambdas: which forms run natively, the equivalent built-in array functions, and how nvl(x, '') != '' on a numeric array is handled. The filter and exists rows of the Supported Expressions page link to it. CometArrayExists calls the dispatcher itself, as CometArrayFilter does, so the generated Implementation column says Native for both. CometHigherOrderFunctionBenchmark produces the numbers above, filter and exists.

Not included: transform and forall still use the dispatcher. DataFusion 55.2 has array_transform but no all-match function.

How are these changes tested?

• array_filter_native_lambda.sql runs 42 queries against Spark with ANSI mode off: every supported element type including tinyint and smallint, NULL, empty and all-NULL arrays and batches, literal arrays, NaN and -0.0, captured columns, a parameter that shadows a column, nested filters, and bodies that must stay on the dispatcher. Each query pins its path with expect_native(filter) or expect_dispatch(filter). It also has queries with a lambda inside an argument evaluated for only some rows (array_contains(ints, size(filter(ints, x -> x > 1)))) and as the right operand of AND and OR.
• array_filter_native_lambda_ansi.sql covers the short-circuit queries under ANSI mode, and array_filter_native_lambda_captured.sql covers bodies that read captured columns of each kind and array functions that take the lambda parameter, with ANSI on and off. array_filter_native_lambda_unsupported_body.sql pins the bodies with a null-map-key cast or a date or timestamp cast to the dispatcher, for filter and exists, with a widening map cast and a timestamp to number cast on the native path. array_filter_native_lambda_disabled.sql checks that spark.comet.exec.nativeLambda.enabled=false runs every lambda on the dispatcher.
• array_filter_numeric_nvl.sql covers the array_compact rewrite: every numeric, date and timestamp element type, both operand orders, <> and coalesce, NULL, empty and all-NULL arrays, NaN and -0.0, the result inside size and array_max, and the string, other-default, = and extra-condition shapes that must keep their lambda semantics. Spark refuses coalesce(boolean, ''), so there is no boolean query.
• array_exists_native_lambda.sql runs 19 queries against Spark with ANSI mode off, 15 of them pinned with expect_native(exists): string, integer, double and float arrays, NULL, empty and all-NULL arrays, predicates that are NULL for some elements and false for the rest, captured columns, flatten as the argument, and exists as the right operand of OR and inside an argument of array_contains. Two bodies can raise on an element Spark never reaches and are pinned with expect_dispatch(exists): element_at with the index 0, and timestamp_seconds of a BIGINT, which raises on overflow. Two aggregate queries only compare with Spark, because the harness counts an expression inside an aggregate as dispatched as well as native.
• array_exists_native_lambda_ansi.sql runs six queries under ANSI mode. Four bodies raise only on an element after the first match (a division, an addition, a cast of a string, and a division inside CASE) and are pinned with expect_dispatch(exists). Two that cannot raise are pinned with expect_native(exists). Without the condition on exists, the division query raised DIVIDE_BY_ZERO natively, and the element_at query in the non-ANSI file raised INVALID_INDEX_OF_ZERO, where Spark returns true for both.
• array_exists_native_lambda_two_valued.sql checks that exists stays on the dispatcher with three-valued logic off. codegen_dispatch_null_map_key.sql pins the exists with a TRY map cast whose key cast can fail to the dispatcher, and the same query with a widening key cast to the native path.
• Rust unit tests cover the planner arm (the array_filter result, a non-array argument, an unknown function), array_any_match as the right operand of OR and AND, and NullShortCircuit with a lambda variable in a masked argument. Without the NullShortCircuit change, the last one fails with LambdaVariable references column … but batch only has 0 columns.
• A one-off comparison against Spark on 300k rows, with 25 lambda bodies and ANSI on and off, matched in every case where Spark does not raise.
• CI passes on this commit: the Comet suites on every Spark profile, the Rust tests and clippy, and the Spark 4.1 SQL suite through the run-spark-4.1-tests label. Locally, the lambda, exists and codegen_dispatch SQL files, CometCodegenHOFSuite, CometArrayExpressionSuite and GenerateDocsSuite pass on Spark 4.1, and the SQL files and CometCodegenHOFSuite pass on Spark 3.5.

[andygrove]

cc @hsiang-c

[comphead]

DataFusion 55.2's HigherOrderFunctionExpr calls remove_list_null_values before it evaluates the body, and list_values slices from the first offset to the last, so values under a NULL array or outside a sliced range never reach the lambda. A NULL predicate drops the element and a NULL array stays NULL, as in Spark's ArrayFilter, and captured columns are re-indexed correctly. Three comments.

[comphead]

DataFusion's lambda copies every captured column once per element before evaluating the body: evaluate_single_list_lambda in datafusion-functions-nested 55.2 spreads the captures with take_arrays(arrays, &list_values_row_number(&original_list)?). For an INT capture such as threshold, the shape the benchmark measures, that is cheap. For a string, binary or array capture the copy grows with the value size. filter(a, x -> array_contains(b, x)) over 100-element INT arrays copies b 100 times per row, roughly 330 MB for one 8192-row batch by my arithmetic, and as far as I can tell those allocations are not reserved with the memory pool. Spark reads the captured value in place, and DuckDB's list_filter references captures through a selection vector. Could the native path decline when the body captures a column that is not fixed-width, or could the benchmark add a captured string or array case to show where it stops paying off? Not measured.

[andygrove]

Good catch, and it applies to more than array_contains. Any captured string, binary, array, struct or map column is copied once per element before the body runs. I haven't checked the memory-pool accounting myself, but declining the native path for these captures sidesteps the question.

nativeLambdaSupported now declines when the body reads a captured column that isn't fixed-width, so those lambdas stay on the dispatcher. Numeric, boolean, date, timestamp and decimal captures still run natively, which is the case the benchmark measures. I didn't add a captured-string benchmark case, because that case now takes the dispatcher path and would only measure the old behaviour. Fixed in a16bdba. The tuning page now says to capture only fixed-width columns.

[comphead]

Some natively supported expressions raise with ANSI off too. Spark's ElementAt throws INVALID_INDEX_OF_ZERO for index 0 outside the failOnError branch, and Comet's one_based_index in list_extract.rs does the same. So with spark.sql.ansi.enabled=false, filter(ints, x -> x > 0 AND element_at(other, x) > 1) takes the native path, where Spark never evaluates element_at(other, 0) because And stops at x > 0. DataFusion's BinaryExpr pre-selects rows for the right side only when the left side is at most 20% true (PRE_SELECTION_THRESHOLD), so for array(0, 1, 2) it evaluates element_at on every element and raises. I have not run it. Could the check classify each node instead, treating ElementAt as raising in any mode and judging arithmetic and casts by their own evalMode, with that query in the non-ANSI fixture?

[andygrove]

You're right. The ANSI check missed expressions that raise in every mode, and element_at with the index 0 is one.

The check no longer looks at the session's ANSI setting. It classifies each expression and fails closed. Arithmetic and casts count as able to raise only when they carry the ANSI eval mode, and anything not on the known-safe list, including element_at, counts as able to raise. A body that combines AND, OR or IN with such an expression stays on the dispatcher. Non-ANSI bodies like x > 0 AND x % 3 = 0 still run natively.

I added your case to the non-ANSI fixture as v > 0 AND element_at(array(threshold, 5), v) > 1. I used a fixed-width array instead of the other column from your example, because a captured array now takes the dispatcher for the reason in the first thread, which would hide this rule. It is pinned with expect_dispatch(filter) and matches Spark. Fixed in a16bdba.

[comphead]

A plain query passes whether the lambda ran natively or in the dispatcher, so this file cannot show which path a query took, and the false arm runs the same queries through the dispatcher a second time. sql-file-tests.md documents query expect_native(...) and query expect_dispatch(...) for a serde that splits between the two paths. With expect_native(filter) on the native queries and expect_dispatch(filter) on the RLIKE, nested-lambda and two-parameter ones, here and in _ansi.sql, the ConfigMatrix and the three new CometCodegenHOFSuite tests could go, since they assert the same split through dispatcher counters. isLambdaElementType also admits tinyint and smallint elements, but the table has no column of either type.

[andygrove]

Agreed, a plain query couldn't show which path ran. Every query in array_filter_native_lambda.sql, the _ansi file and the new captured-column file now carries expect_native(filter) or expect_dispatch(filter). The nativeLambda ConfigMatrix is gone. A small array_filter_native_lambda_disabled.sql checks that the kill switch sends everything to the dispatcher. The three new CometCodegenHOFSuite tests and the assertNoCodegen helper I had added are removed, so those two files match apache/main again.

I also added tinyint and smallint columns, since isLambdaElementType admits both and nothing exercised them. Fixed in a16bdba.

[andygrove]

Thanks for the review. All three are addressed in a16bdba, and the PR description is updated to match. The Spark SQL suite is still running through the run-* labels.

[rich7420]

Component repro with the production projection functions: base passes; head fails with project index 2 out of bounds. Full Spark query not run.

[rich7420]

For array_contains(a, size(filter(b, x -> x > 1))), this adds the inner lambda parameter to the outer projection, causing an out-of-bounds error. Preserve lambda scope when collecting and remapping columns.

[andygrove]

Good catch. project() collected every LambdaVariable as a column of the outer batch, but a lambda inside the projected expression binds its own parameter, which is not a column of that batch, so batch.project was asked for an index past the end. I reproduced it with a higher-order function as the right operand of AND and OR, which goes through the same function, and the same applied to filter.

project() now skips a variable whose lambda is inside the expression it projects, and leaves it unchanged when it re-indexes. Lambda parameter names are unique per lambda (x#id), so collecting the names a lambda in the expression binds is enough to tell them apart from the free variables that the earlier fix needs, where a short-circuited argument inside a lambda body reads the lambda's own parameter.

Tests: your query, array_contains(ints, size(filter(ints, x -> x > 1))), and AND/OR with a lambda on the right are in array_filter_native_lambda.sql, and a Rust planner test covers AND and OR. I also compared these shapes against Spark with ANSI on and off. Fixed in 1656959.

[comphead]

#6485 this might be good fit for such things, but requires way more work and testing

[comphead]

I'll try to adopt this example for case study and see if we can get some benefits

[comphead]

Second pass at 66a1066. The three earlier points are addressed. The lambda-scope fix in project() also holds: DataFusion 55.2's HigherOrderFunctionExpr::evaluate keeps only the lambda.projection() indexes below batch.num_columns(), so an inner parameter that keeps its original index still resolves after the outer projection. One correctness point on exists, and a few smaller ones.

[comphead]

Spark's ArrayExists stops at the first matching element (while (i < arr.numElements && !exists), the same in 3.4.3, 3.5.8, 4.0.1 and main). array_any_match evaluates the body over every element of the batch first. This is the case shortCircuitSafe handles for AND, OR and IN, except that exists skips elements without any of them, so the gate passes.

With ANSI on, which is the default on Spark 4, exists(ints, x -> 100 / x > 1) over array(1, 0) returns true in Spark and should raise DIVIDE_BY_ZERO natively. With ANSI off, element_at with the index 0 behaves the same way. I have not run either.

Could exists also require !body.exists(mayRaise), whatever operators the body uses? That query would then go into an ANSI fixture as expect_dispatch(exists), and the tuning page bullet would cover exists. filter evaluates every element, so its current rule is enough.

[andygrove]

You're right. I ran both before the fix: with ANSI on, exists(ints, x -> 100 / x > 1) over array(1, 0) raised DIVIDE_BY_ZERO natively, and with ANSI off exists(ints, x -> element_at(array(threshold, 5), x) > 1) raised INVALID_INDEX_OF_ZERO. Spark returns true for both.

exists now also requires !body.exists(mayRaise), whatever operators the body uses. The new array_exists_native_lambda_ansi.sql pins your query, an ANSI addition that overflows, an ANSI cast of a string and a division inside CASE with expect_dispatch(exists). Each of them raises only on an element after the first match. Two bodies that cannot raise are pinned with expect_native(exists). The element_at query is in the non-ANSI file. The date-conversion query there moved to expect_dispatch(exists) as well, because timestamp_seconds of a BIGINT raises on overflow. The tuning page has a bullet for exists. Fixed in 3bb85cf.

[comphead]

CometCast.canProduceNullMapKey is already the #6172 predicate (a TRY cast whose key cast can fail), and CometBatchKernelCodegen.canHandle uses it. holdsMap also matches casts that cannot make a null key, such as the widening one in codegen_dispatch_null_map_key.sql. Reusing it, with an org.apache.comet.expressions.CometCast import:

Suggested change

	private def castsToMap(argument: Expression): Boolean = {
	def holdsMap(dt: DataType): Boolean = dt match {
	case _: MapType => true
	case ArrayType(element, _) => holdsMap(element)
	case StructType(fields) => fields.exists(f => holdsMap(f.dataType))
	case _ => false
	}
	argument.exists {
	case c: Cast => holdsMap(c.dataType)
	case _ => false
	}
	}
	private def castsToMap(argument: Expression): Boolean =
	argument.exists {
	case c: Cast => CometCast.canProduceNullMapKey(c)
	case _ => false
	}

The Scaladoc also says the dispatcher falls back to Spark for such an argument. But that file pins the exists(map_keys(try_cast(...)), ...) query with expect_dispatch(exists): the dispatcher runs it, because the map never leaves Spark's code.

[andygrove]

Agreed. It uses CometCast.canProduceNullMapKey now, and I renamed the helper to canProduceNullMapKey, since it no longer looks at the target type. The Scaladoc now says the dispatcher runs such an argument in Spark's own code. I added the widening cast to codegen_dispatch_null_map_key.sql as an expect_native(exists) query, next to the dispatched one. Fixed in 3bb85cf.

[comphead]

apply visits a node before its children, and a bound LambdaVariable always sits below its LambdaExpr. So bound is already complete when a variable is reached, and the two walks can be one:

Suggested change

	let mut bound: HashSet<String> = HashSet::new();
	expr.apply(|e| {
	if let Some(lambda) = e.downcast_ref::<LambdaExpr>() {
	bound.extend(lambda.params().iter().cloned());
	}
	Ok(TreeNodeRecursion::Continue)
	})?;
	let mut projection = vec![];
	expr.apply(|e| {
	if let Some(column) = e.downcast_ref::<Column>() {
	projection.push(column.index());
	} else if let Some(variable) = e.downcast_ref::<LambdaVariable>() {
	if !bound.contains(variable.name()) {
	projection.push(variable.index());
	}
	}
	Ok(TreeNodeRecursion::Continue)
	})?;
	// `apply` visits each lambda before the variables it binds.
	let mut bound: HashSet<String> = HashSet::new();
	let mut projection = vec![];
	expr.apply(|e| {
	if let Some(lambda) = e.downcast_ref::<LambdaExpr>() {
	bound.extend(lambda.params().iter().cloned());
	} else if let Some(column) = e.downcast_ref::<Column>() {
	projection.push(column.index());
	} else if let Some(variable) = e.downcast_ref::<LambdaVariable>() {
	if !bound.contains(variable.name()) {
	projection.push(variable.index());
	}
	}
	Ok(TreeNodeRecursion::Continue)
	})?;

[andygrove]

Done, it's one walk now, and your comment is folded into the existing one. Fixed in 3bb85cf.

[comphead]

The Implementation column is generated by GenerateDocs.classifySerde. It still says Codegen dispatch here because CometArrayExists extends CometCodegenDispatch, and the Notes cell is empty. Could the Notes say what the filter row says, with the same link to Expression Tuning? On that page, the capture bullet leaves out decimal, which isFixedWidth accepts.

[andygrove]

The Notes cell now matches the filter row, with the conditions that only apply to exists, and it has the same link. I also made CometArrayExists a plain CometExpressionSerde that calls the dispatcher itself, as CometArrayFilter does, so the generated Implementation column says Native for both. The capture bullet lists decimal now. Fixed in 3bb85cf.

[comphead]

The true block asserts what expect_native(exists) in array_exists_native_lambda.sql already pins. The false block asserts what array_exists_native_lambda_two_valued.sql pins with expect_dispatch(exists). As with the filter tests in the first round, could one copy go, preferably this test and the CometScalaUDFCodegen import?

Also, filter(ints, x -> nvl(x, '') != '') in array_filter_native_lambda.sql now runs as array_compact. Its comment about casting each element is stale, and the query repeats array_filter_numeric_nvl.sql.

[andygrove]

Removed the test and the CometScalaUDFCodegen import. I also removed the numeric nvl query and its stale comment from array_filter_native_lambda.sql, since array_filter_numeric_nvl.sql covers it. Fixed in 3bb85cf.

[comphead]

This half checks the three-valued result again. DataFusion's own array_any_match tests already cover it, and array_exists_native_lambda.sql covers it against Spark. expr and bound are also defined again in test_higher_order_as_right_operand. Could this test keep the planning checks (the array_filter result, the non-array argument, the unknown function), with the two helpers at module level?

[andygrove]

Done. test_higher_order_filter keeps the array_filter result and the two error cases, and expr and bound are module-level helpers that both tests use. Fixed in 3bb85cf.

[andygrove]

Thanks for the second pass. All six points are addressed in 3bb85cf, and the description is updated to match. The exists one was a real bug: both of your examples raised natively before the fix.

Thanks for the pointer to #6485 too. I'd be curious how a fused kernel does on a lambda body like these.

[rich7420]

Check canProduceNullMapKey in the body too. For [1L, 9999999999L], array_contains(map_keys(try_cast(map_from_arrays(array(x), array(1)) AS MAP<INT,INT>)), 1) returns [1] through Spark's filter, but native casting aborts on a NULL key.

[andygrove]

You are right, and I reproduced it: with [1L, 9999999999L] Comet aborted on the NULL key where Spark returns [1]. The check now covers the body as well as the argument, for filter and exists. array_filter_native_lambda_unsupported_body.sql pins your query with expect_dispatch, next to the same shape with a widening key cast on the native path.

[rich7420]

filter(dates, x -> CAST(x AS STRING) <> '') now hits #7471 for date_from_unix_date(100000000): Spark succeeds, native formatting fails. Keep these bodies on the dispatcher until the kernel supports Spark's date range.

[andygrove]

Agreed. A body with a cast between a date or timestamp and a string, date or timestamp now stays on the dispatcher, until the cast supports Spark's year range (#7471). That covers CAST(x AS STRING) on your date_from_unix_date(100000000), and the casts between a date and a timestamp, which the issue shows failing too. The fixture pins them with expect_dispatch, and a timestamp to number cast stays native. The tuning page has a bullet for it.