andygrove commented on code in PR #7322:
URL: https://github.com/apache/datafusion-comet/pull/7322#discussion_r4418426063
##########
spark/src/main/scala/org/apache/comet/serde/arrays.scala:
##########
@@ -1041,15 +1042,133 @@ 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 binding && nativeLambdaSupported(expr, body) =>
+ // The body is made of expressions Comet evaluates natively, so run
the lambda over the
+ // flattened elements of the whole batch instead of once per element
in the JVM.
+ convertLambda(expr, lambdaVar, body, inputs, binding).orElse(
+ dispatch(expr, inputs, binding))
+ case _ => dispatch(expr, inputs, binding)
+ }
+ }
+
+ /**
+ * 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.
+ */
+ private def dispatch(
+ expr: ArrayFilter,
+ inputs: Seq[Attribute],
+ binding: Boolean): Option[ExprOuterClass.Expr] =
+ CometScalaUDF.emitJvmCodegenDispatch(expr, inputs, binding)
+
+ // 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. It needs
bound references: the
+ // native side re-indexes the columns the body captures, which it cannot do
for unbound ones.
+ private def nativeLambdaSupported(expr: ArrayFilter, body: Expression):
Boolean =
Review Comment:
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 a16bdba0c5. The tuning page now says to capture only
fixed-width columns.
##########
spark/src/main/scala/org/apache/comet/serde/arrays.scala:
##########
@@ -1041,15 +1042,133 @@ 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 binding && nativeLambdaSupported(expr, body) =>
+ // The body is made of expressions Comet evaluates natively, so run
the lambda over the
+ // flattened elements of the whole batch instead of once per element
in the JVM.
+ convertLambda(expr, lambdaVar, body, inputs, binding).orElse(
+ dispatch(expr, inputs, binding))
+ case _ => dispatch(expr, inputs, binding)
+ }
+ }
+
+ /**
+ * 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.
+ */
+ private def dispatch(
+ expr: ArrayFilter,
+ inputs: Seq[Attribute],
+ binding: Boolean): Option[ExprOuterClass.Expr] =
+ CometScalaUDF.emitJvmCodegenDispatch(expr, inputs, binding)
+
+ // 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. It needs
bound references: the
+ // native side re-indexes the columns the body captures, which it cannot do
for unbound ones.
+ private def nativeLambdaSupported(expr: ArrayFilter, body: Expression):
Boolean =
+ CometConf.COMET_NATIVE_LAMBDA_ENABLED.get() &&
+ body.deterministic &&
+ body.find(_.isInstanceOf[HigherOrderFunction]).isEmpty &&
+ shortCircuitSafe(body) &&
+ (expr.argument.dataType match {
+ case ArrayType(elementType, _) => isLambdaElementType(elementType)
+ case _ => false
+ })
+
+ /**
+ * Spark evaluates `a AND b`, `a OR b` and `x IN (a, b)` lazily, so `x != 0
AND 100 / x > 1`
+ * never divides by zero. The native side evaluates every operand over the
whole batch, so under
+ * ANSI mode the same lambda raises on the element Spark skips. `CASE WHEN`,
`IF` and `COALESCE`
+ * are not affected. When ANSI mode is on, a body that short-circuits stays
on the dispatcher
+ * unless every expression in it is one that cannot raise.
+ */
+ private def shortCircuitSafe(body: Expression): Boolean =
+ !SQLConf.get.ansiEnabled ||
Review Comment:
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 a16bdba0c5.
##########
spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda.sql:
##########
@@ -0,0 +1,164 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+-- `filter` with a lambda whose body Comet evaluates natively runs through
DataFusion's
+-- `array_filter`. Each query below must return what Spark returns, including
for NULL arrays,
+-- empty arrays, NULL elements and lambdas that capture a column.
+
+-- ConfigMatrix: spark.comet.exec.nativeLambda.enabled=true,false
Review Comment:
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 a16bdba0c5.
--
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]
