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