rich7420 commented on code in PR #7322:
URL: https://github.com/apache/datafusion-comet/pull/7322#discussion_r4499955252
##########
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:
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.
##########
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:
`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.
--
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]
