andygrove commented on code in PR #7322:
URL: https://github.com/apache/datafusion-comet/pull/7322#discussion_r4526942111
##########
spark/src/main/scala/org/apache/comet/serde/arrays.scala:
##########
@@ -1026,7 +1029,128 @@ object CometFlatten extends
CometExpressionSerde[Flatten] with ArraysBase {
}
}
-object CometArrayFilter extends CometExpressionSerde[ArrayFilter] {
+/**
+ * What a serde needs to run the lambda of a higher-order function natively,
under a DataFusion
+ * higher-order function: the gate on the lambda and its body, and the
serialization.
+ */
+trait NativeLambdaSerde extends CometTypeShim with CometExprShim {
+
+ // The native lambda has the same semantics as Spark's `filter` for a
one-parameter lambda whose
+ // body is deterministic, contains no nested lambda, and whose parameter is
a type the native
+ // side can type-check against the list's element field. `exists` adds a
condition of its own,
+ // since Spark stops at the first match. It needs bound references: the
native side re-indexes
+ // the columns the body captures, which it cannot do for unbound ones.
+ protected def nativeLambdaSupported(argument: Expression, body: Expression):
Boolean =
+ CometConf.COMET_NATIVE_LAMBDA_ENABLED.get() &&
+ body.deterministic &&
+ !body.exists(_.isInstanceOf[HigherOrderFunction]) &&
+ capturesOnlyFixedWidthColumns(body) &&
+ shortCircuitSafe(body) &&
+ !canProduceNullMapKey(argument) &&
Review Comment:
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.
##########
spark/src/main/scala/org/apache/comet/serde/arrays.scala:
##########
@@ -1041,13 +1165,79 @@ object CometArrayFilter extends
CometExpressionSerde[ArrayFilter] {
// restore the native serde here (avoids per-batch JNI). Guard
requires the IsNotNull
// operand to be the lambda variable itself, not a captured column
(#4830).
CometArrayCompact.convert(expr, inputs, binding)
- case _ =>
- // General lambda: run Spark's own evaluation through the codegen
dispatcher so the result
- // matches Spark exactly, like the other higher-order functions
(`transform`, `exists`).
- // Falls back to Spark when the dispatcher is disabled.
- CometScalaUDF.emitJvmCodegenDispatch(expr, inputs, binding)
+ case LambdaFunction(body, Seq(lambdaVar: NamedLambdaVariable), _)
+ if isNotNullInDisguise(body, lambdaVar) =>
+ // `nvl(x, '') != ''` on a numeric array: same native serde, and no
cast of every element
+ // to a string.
+ CometArrayCompact.convert(expr, inputs, binding)
+ case LambdaFunction(body, Seq(lambdaVar: NamedLambdaVariable), _)
+ if binding && nativeLambdaSupported(expr.argument, body) =>
Review Comment:
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
