From a30b862fdc636d8c5166da849681c4636379c2ef Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 29 Oct 2026 14:12:07 -0600 Subject: [PATCH 1/4] feat: run filter lambdas natively with DataFusion's array_filter A filter lambda other than `x IS NOT NULL` went through the JVM codegen dispatcher, which evaluates it one element at a time like Spark does, plus the transfer of each batch. A body made of natively supported expressions now runs once over all the elements of a batch through DataFusion's array_filter. The native path applies to a one-parameter lambda over an array of a primitive, decimal, string, binary, date or timestamp element type. These stay on the dispatcher: nested lambdas, nondeterministic bodies, bodies that call out to the JVM, non-default string collations, and under ANSI mode a body that short-circuits (AND, OR, IN) around an operation that can raise, since the native operators evaluate every operand. The new spark.comet.exec.nativeLambda.enabled turns the native path off. Adds a lambda tuning section to the user guide and a benchmark. Closes #7321. --- docs/source/user-guide/latest/expressions.md | 2 +- docs/source/user-guide/latest/index.rst | 1 + docs/source/user-guide/latest/tuning.md | 2 + .../user-guide/latest/tuning/expressions.md | 75 ++++++++ native/core/src/execution/planner.rs | 63 ++++++- native/proto/src/proto/expr.proto | 23 +++ .../scala/org/apache/comet/CometConf.scala | 13 ++ .../apache/comet/serde/QueryPlanSerde.scala | 1 + .../scala/org/apache/comet/serde/arrays.scala | 133 +++++++++++++- .../array/array_filter_native_lambda.sql | 164 ++++++++++++++++++ .../array/array_filter_native_lambda_ansi.sql | 67 +++++++ .../apache/comet/CometCodegenHOFSuite.scala | 74 ++++++++ .../CometHigherOrderFunctionBenchmark.scala | 96 ++++++++++ 13 files changed, 704 insertions(+), 10 deletions(-) create mode 100644 docs/source/user-guide/latest/tuning/expressions.md create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_ansi.sql create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometHigherOrderFunctionBenchmark.scala diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index f66a705afac..3596fdb2daf 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -388,7 +388,7 @@ to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enab | `aggregate` | ✅ | Codegen dispatch | | | `array_sort` | ✅ | Native | The default comparator runs natively; custom comparators and arrays with non-default string collations use the JVM codegen dispatcher | | `exists` | ✅ | Codegen dispatch | | -| `filter` | ✅ | Native | General lambda routed through the JVM codegen dispatcher; the `array_compact` form runs natively | +| `filter` | ✅ | Native | A one-parameter lambda whose body is natively supported runs natively, as does the `array_compact` form; other lambdas use the JVM codegen dispatcher. See [Expression Tuning](tuning/expressions.md#lambdas-and-higher-order-functions) | | `forall` | ✅ | Codegen dispatch | | | `map_filter` | ✅ | Codegen dispatch | | | `map_zip_with` | ✅ | Codegen dispatch | | diff --git a/docs/source/user-guide/latest/index.rst b/docs/source/user-guide/latest/index.rst index 1217d3ea798..bde1ebe05eb 100644 --- a/docs/source/user-guide/latest/index.rst +++ b/docs/source/user-guide/latest/index.rst @@ -96,6 +96,7 @@ the navigation menu to read more. Shuffle Celeborn Shuffle Operators + Expressions Row/Columnar Transitions .. toctree:: diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index 9b22839b568..f6de447d1eb 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -35,6 +35,8 @@ guide is split into the following pages: `spark.comet.scan.icebergNative.dataFileConcurrencyLimit`. - [Operator Tuning](tuning/operators.md): joins, adaptive partial aggregation, and sorting on floating-point values. +- [Expression Tuning](tuning/expressions.md): lambdas and higher-order functions such as `filter`, and the + built-in array functions that replace them. - [Reducing Row/Columnar Conversion Overhead](tuning/transitions.md): stages in which many operators fall back to Spark, and the experimental direct columnar-to-row converter. diff --git a/docs/source/user-guide/latest/tuning/expressions.md b/docs/source/user-guide/latest/tuning/expressions.md new file mode 100644 index 00000000000..1a22c4fbd13 --- /dev/null +++ b/docs/source/user-guide/latest/tuning/expressions.md @@ -0,0 +1,75 @@ + + +# Expression Tuning + +## Lambdas and Higher-Order Functions + +A higher-order function takes a lambda: `filter(arr, x -> ...)`, `transform`, `exists`, `forall`, `aggregate`, +`zip_with` and `map_filter`. How fast Comet runs one depends on whether it can run the lambda natively. + +| Lambda | How Comet runs it | +| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `filter(arr, x -> x IS NOT NULL)` (what `array_compact` expands to) | Natively, as `array_compact`. | +| `filter` with a one-parameter lambda whose body uses only natively supported expressions, over an array of primitive, decimal, string, binary, date or timestamp elements | Natively. The lambda runs once over all the elements of a batch, vectorized, instead of once per element. | +| Any other lambda, and every other higher-order function | Through the [JVM codegen dispatcher](../scala_java_udfs.md). Spark's interpreter evaluates the lambda for each element inside a Comet kernel, so expect Spark's speed, not more. | + +The dispatcher has no advantage over Spark for a lambda, because Spark also evaluates a lambda one element at a time. +It adds the cost of handing each batch to the JVM and back. The rest of the query stays native either way. + +To get the most out of Comet when a query uses lambdas: + +- **Prefer a built-in array function to a lambda.** Most of them run natively with no per-element evaluation, for + example `array_compact`, `array_remove`, `array_distinct` and `array_position`. See the + [Supported Expressions](../expressions.md) tables for the caveats of each. + `array_compact(arr)` returns what `filter(arr, x -> x IS NOT NULL)` does for any array, and + `array_remove(array_compact(arr), '')` returns what `filter(arr, x -> nvl(x, '') != '')` does for an array of + strings with the default collation. +- **Do not test a numeric array for nulls with a string comparison.** With ANSI mode off, Spark casts every element + of `filter(int_arr, x -> nvl(x, '') != '')` to a string to compare it with `''`; with ANSI mode on the query does + not analyze. Write `x IS NOT NULL`, or call `array_compact`. +- **Keep the lambda body to natively supported expressions.** A body that calls a regular expression function, a + UDF or another lambda cannot be split from the dispatcher, so the whole lambda runs there. Use + [`spark.comet.explain.codegen.enabled`](../understanding-comet-plans.md#sparkcometexplaincodegenenabled) to see which + expressions of a query went through the dispatcher. +- **Under ANSI mode, keep `AND`, `OR` and `IN` away from operations that can raise.** Spark evaluates them lazily, so + `x != 0 AND 100 / x > 1` never divides by zero, but the native lambda evaluates every operand for every element. A + body that combines them with an operation that can raise, such as a division, an addition that can overflow or a + cast, stays on the dispatcher. `CASE WHEN`, `IF` and `COALESCE` are lazy natively too, so + `CASE WHEN x = 0 THEN false ELSE 100 / x > 1 END` runs natively. +- **Use a lambda of one parameter.** `filter(arr, (x, i) -> i > 0)`, which also receives the element's index, runs + through the dispatcher. + +`spark.comet.exec.nativeLambda.enabled=false` makes Comet run every lambda through the dispatcher. + +### Example + +`CometHigherOrderFunctionBenchmark` filters 1 million rows of ten-element arrays, in which every fifth element is +`NULL`. These are the best times of the Spark baseline and of Comet on an Apple M-series laptop, with +`spark.comet.exec.nativeLambda.enabled` off and at its default: + +| Query | Spark | Comet, nativeLambda off | Comet, default | +| ----------------------------------------------------------- | ------ | ----------------------- | -------------- | +| `filter(strs, x -> nvl(x, '') != '')` (string array) | 616 ms | 547 ms | 307 ms | +| `array_remove(array_compact(strs), '')` (same result) | 587 ms | 252 ms | 252 ms | +| `filter(ints, x -> nvl(x, '') != '')` (integer array) | 508 ms | 456 ms | 331 ms | +| `filter(ints, x -> x IS NOT NULL)` | 286 ms | 124 ms | 125 ms | +| `filter(ints, x -> x > threshold)`, `threshold` is a column | 346 ms | 335 ms | 110 ms | + +Rewriting a lambda to built-in functions is still the fastest form, but a native lambda gets most of the way there. diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 899d8ff3320..59b03509288 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -66,8 +66,11 @@ use datafusion::{ execution::FunctionRegistry, logical_expr::Operator as DataFusionOperator, physical_expr::{ - expressions::{BinaryExpr, CastExpr, Column, IsNullExpr, Literal as DataFusionLiteral}, - PhysicalExpr, PhysicalSortExpr, ScalarFunctionExpr, + expressions::{ + BinaryExpr, CastExpr, Column, IsNullExpr, LambdaExpr, LambdaVariable, + Literal as DataFusionLiteral, + }, + HigherOrderFunctionExpr, PhysicalExpr, PhysicalSortExpr, ScalarFunctionExpr, }, physical_plan::{ empty::EmptyExec, limit::LocalLimitExec, sorts::sort::SortExec, ExecutionPlan, @@ -93,6 +96,7 @@ use datafusion::common::scalar::ScalarStructBuilder; use datafusion::common::ScalarValue; use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::object_store::ObjectStoreUrl; +use datafusion::functions_nested::array_filter::array_filter_higher_order_function; use datafusion::logical_expr::type_coercion::functions::fields_with_udf; use datafusion::logical_expr::{ReturnFieldArgs, ScalarUDF, TypeSignature}; use datafusion::physical_expr::expressions::Literal; @@ -1248,6 +1252,61 @@ impl PhysicalPlanner { )); Ok(expr) } + ExprStruct::HigherOrderFunc(hof) => { + let array = + self.create_expr(hof.array.as_ref().unwrap(), Arc::clone(&input_schema))?; + // The lambda parameter takes the element field of the array as planned, since + // DataFusion checks it against the field of the list it evaluates. + let array_field = array.return_field(input_schema.as_ref())?; + let element_field = match array_field.data_type() { + DataType::List(f) | DataType::LargeList(f) => f, + other => { + return Err(GeneralError(format!( + "{} expects an array argument, got {other}", + hof.func + ))) + } + }; + // The lambda body sees the enclosing schema with the parameter appended, so + // `BoundReference`s keep their indexes and a `LambdaVariable` is found by name. + let mut fields: Vec = input_schema.fields().iter().cloned().collect(); + fields.push(Arc::new( + element_field.as_ref().clone().with_name(&hof.param_name), + )); + let lambda_schema = Arc::new(Schema::new_with_metadata( + fields, + input_schema.metadata().clone(), + )); + let body = self.create_expr(hof.body.as_ref().unwrap(), lambda_schema)?; + let lambda: Arc = + Arc::new(LambdaExpr::try_new(vec![hof.param_name.clone()], body)?); + let func = match hof.func.as_str() { + "array_filter" => array_filter_higher_order_function(), + other => { + return Err(GeneralError(format!( + "Unsupported higher-order function '{other}'" + ))) + } + }; + Ok(Arc::new(HigherOrderFunctionExpr::try_new_with_schema( + func, + vec![array, lambda], + input_schema.as_ref(), + Arc::new(ConfigOptions::default()), + )?)) + } + ExprStruct::LambdaVariable(var) => { + let idx = input_schema.index_of(&var.name).map_err(|_| { + GeneralError(format!( + "Lambda variable '{}' is not in scope. Schema: {input_schema}", + var.name + )) + })?; + Ok(Arc::new(LambdaVariable::new( + idx, + Arc::clone(&input_schema.fields()[idx]), + ))) + } expr => Err(GeneralError(format!("Not implemented: {expr:?}"))), } } diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto index 9f631491c4f..a34188664ff 100644 --- a/native/proto/src/proto/expr.proto +++ b/native/proto/src/proto/expr.proto @@ -96,6 +96,8 @@ message Expr { Uuid uuid = 74; NativeScalarUdf native_scalar_udf = 75; AtLeastNNonNulls at_least_n_non_nulls = 76; + HigherOrderFunc higher_order_func = 77; + LambdaVariable lambda_variable = 78; } reserved 20; @@ -739,6 +741,27 @@ message ArraysZip { repeated string names = 2; } +// A DataFusion higher-order function applied to one array argument and a +// one-parameter lambda, e.g. `array_filter(array, x -> body)`. The lambda body +// is an ordinary expression tree in which `LambdaVariable` nodes refer to the +// lambda parameter and `BoundReference` nodes refer to columns of the +// enclosing operator's input. +message HigherOrderFunc { + // DataFusion higher-order function name, e.g. "array_filter". + string func = 1; + Expr array = 2; + // Unique name of the lambda parameter. `LambdaVariable.name` nodes in the + // body resolve against it. + string param_name = 3; + Expr body = 4; +} + +// Reference to a lambda parameter from inside a `HigherOrderFunc` body. The +// element type is taken from the array argument at planning time. +message LambdaVariable { + string name = 1; +} + // Scalar UDF dispatched to the JVM via JNI. Native side exports input arrays // through Arrow C Data Interface, calls CometUdfBridge.evaluate, and imports // the result. diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 6003703a630..6c1f5d436cd 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -767,6 +767,19 @@ object CometConf extends ShimCometConf { .toSequence .createWithDefault(Nil) + val COMET_NATIVE_LAMBDA_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.exec.nativeLambda.enabled") + .category(CATEGORY_EXEC) + .doc("Whether to evaluate the lambda of a higher-order function natively when its body " + + "is made only of expressions that Comet supports natively. The lambda then runs " + + "vectorized over all the elements of a batch instead of once per element in the JVM. " + + "Currently applies to `filter` over an array of a primitive, decimal, string, binary, " + + "date or timestamp element type, with a lambda of one parameter. Other lambdas run " + + "through the JVM codegen dispatcher " + + "(see `spark.comet.exec.scalaUDF.codegen.enabled`).") + .booleanConf + .createWithDefault(true) + val COMET_SCALA_UDF_CODEGEN_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.exec.scalaUDF.codegen.enabled") .category(CATEGORY_EXEC) diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index c8b89b0742c..b51a729b92c 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -87,6 +87,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { classOf[Size] -> CometSize, classOf[ArraysZip] -> CometArraysZip, classOf[ArrayTransform] -> CometArrayTransform, + classOf[NamedLambdaVariable] -> CometNamedLambdaVariable, classOf[ArrayExists] -> CometArrayExists, classOf[ArrayForAll] -> CometArrayForAll, classOf[ArrayAggregate] -> CometArrayAggregate, diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index 07edb45b741..bb73e73d70d 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -23,12 +23,13 @@ import scala.annotation.tailrec import scala.jdk.CollectionConverters._ import org.apache.spark.SPARK_VERSION -import org.apache.spark.sql.catalyst.expressions.{And, ArrayAggregate, ArrayAppend, ArrayContains, ArrayDistinct, ArrayExcept, ArrayExists, ArrayFilter, ArrayForAll, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayPosition, ArrayRemove, ArraySort, ArraysOverlap, ArraysZip, ArrayTransform, ArrayUnion, Attribute, BoundReference, Cast, CreateArray, ElementAt, EmptyRow, Expression, Flatten, GetArrayItem, GreaterThan, If, IsNotNull, IsNull, LambdaFunction, LessThan, Literal, NamedLambdaVariable, Reverse, Sequence, Size, Slice, SortArray, ZipWith} +import org.apache.spark.sql.catalyst.expressions.{And, ArrayAggregate, ArrayAppend, ArrayContains, ArrayDistinct, ArrayExcept, ArrayExists, ArrayFilter, ArrayForAll, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayPosition, ArrayRemove, ArraySort, ArraysOverlap, ArraysZip, ArrayTransform, ArrayUnion, Attribute, AttributeReference, BinaryComparison, BoundReference, CaseWhen, Cast, Coalesce, Contains, CreateArray, ElementAt, EmptyRow, EndsWith, Expression, Flatten, GetArrayItem, GreaterThan, HigherOrderFunction, If, In, InSet, IsNaN, IsNotNull, IsNull, LambdaFunction, Length, LessThan, Literal, Lower, NamedLambdaVariable, Not, Or, Reverse, Sequence, Size, Slice, SortArray, StartsWith, Upper, ZipWith} import org.apache.spark.sql.catalyst.util.GenericArrayData import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ +import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.withFallbackReason import org.apache.comet.DataTypeSupport.{deepNullable, isComplexType} import org.apache.comet.serde.QueryPlanSerde._ @@ -1026,7 +1027,7 @@ object CometFlatten extends CometExpressionSerde[Flatten] with ArraysBase { } } -object CometArrayFilter extends CometExpressionSerde[ArrayFilter] { +object CometArrayFilter extends CometExpressionSerde[ArrayFilter] with CometTypeShim { override def getSupportLevel(expr: ArrayFilter): SupportLevel = Compatible() @@ -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 || + !body.exists { + case _: And | _: Or | _: In | _: InSet => true + case _ => false + } || !body.exists(!neverRaises(_)) + + private def neverRaises(e: Expression): Boolean = e match { + case _: Literal | _: NamedLambdaVariable | _: AttributeReference | _: BoundReference | + _: BinaryComparison | _: IsNull | _: IsNotNull | _: Not | _: And | _: Or | _: In | + _: InSet | _: IsNaN | _: Coalesce | _: If | _: CaseWhen | _: Length | _: Upper | + _: Lower | _: StartsWith | _: EndsWith | _: Contains => + true + case _ => false + } + + private def isLambdaElementType(dt: DataType): Boolean = dt match { + case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType | + DateType | TimestampType | TimestampNTZType | BinaryType | _: DecimalType => + true + case _: StringType => !hasNonDefaultStringCollation(dt) + case _ => false + } + + private def convertLambda( + expr: ArrayFilter, + lambdaVar: NamedLambdaVariable, + body: Expression, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = { + for { + arrayProto <- exprToProtoInternal(expr.argument, inputs, binding) + bodyProto <- exprToProtoInternal(body, inputs, binding) + // A subexpression the serde routed through the JVM codegen dispatcher cannot be part of the + // lambda. Its call out to the JVM would receive the lambda parameter as an argument, and + // the whole lambda is cheaper to run in the dispatcher than to split across the boundary. + if !CometArrayFilterProto.containsJvmScalarUdf(bodyProto) + } yield { + val hof = ExprOuterClass.HigherOrderFunc + .newBuilder() + .setFunc("array_filter") + .setArray(arrayProto) + .setParamName(CometNamedLambdaVariable.uniqueName(lambdaVar)) + .setBody(bodyProto) + ExprOuterClass.Expr.newBuilder().setHigherOrderFunc(hof).build() } } } +object CometArrayFilterProto { + + /** Whether `message`, or any message nested in it, is a call out to the JVM. */ + def containsJvmScalarUdf(message: com.google.protobuf.Message): Boolean = + message.isInstanceOf[ExprOuterClass.JvmScalarUdf] || + message.getAllFields.values.asScala.exists { + case m: com.google.protobuf.Message => containsJvmScalarUdf(m) + case l: java.util.List[_] => + l.asScala.exists { + case m: com.google.protobuf.Message => containsJvmScalarUdf(m) + case _ => false + } + case _ => false + } +} + +/** + * A lambda parameter inside the body of a higher-order function that Comet runs natively. It is + * only reachable from such a body: the enclosing serde decides whether to serialize the body. + */ +object CometNamedLambdaVariable extends CometExpressionSerde[NamedLambdaVariable] { + + /** + * The native side resolves a lambda variable by name, so the name must be unique per lambda. + */ + def uniqueName(v: NamedLambdaVariable): String = s"${v.name}#${v.exprId.id}" + + override def convert( + expr: NamedLambdaVariable, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = { + val variable = ExprOuterClass.LambdaVariable.newBuilder().setName(uniqueName(expr)) + Some(ExprOuterClass.Expr.newBuilder().setLambdaVariable(variable).build()) + } +} + object CometSize extends CometExpressionSerde[Size] { override def getSupportLevel(expr: Size): SupportLevel = { diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda.sql b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda.sql new file mode 100644 index 00000000000..5f142c834b2 --- /dev/null +++ b/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 + +statement +CREATE TABLE test_filter_lambda( + id int, + strs array, + ints array, + longs array, + dbls array, + flts array, + decs array, + dates array, + tss array, + bins array, + bools array, + ntzs array, + x int, + threshold int, + needle string +) USING parquet + +statement +INSERT INTO test_filter_lambda VALUES + (1, array('a', NULL, '', 'b', ' ', NULL, 'abc'), array(1, NULL, 0, 5, -3), array(1L, NULL, 0L, 9000000000L), + array(1.5, NULL, CAST('NaN' AS DOUBLE), -0.0, 0.0, CAST('Infinity' AS DOUBLE)), + array(CAST(1.5 AS FLOAT), NULL, CAST('NaN' AS FLOAT), CAST(2.5 AS FLOAT)), + array(CAST(1.25 AS DECIMAL(10, 2)), NULL, CAST(-3.50 AS DECIMAL(10, 2)), CAST(0 AS DECIMAL(10, 2))), + array(DATE '2020-01-01', NULL, DATE '2024-02-29'), + array(TIMESTAMP '2020-01-01 00:00:00', NULL, TIMESTAMP '2024-02-29 12:34:56'), + array(X'01', NULL, X'', X'FF'), + array(true, NULL, false), + array(TIMESTAMP_NTZ '2020-01-01 00:00:00', NULL, TIMESTAMP_NTZ '2024-02-29 12:34:56'), + 3, 2, 'a'), + (2, array(), array(), array(), array(), array(), array(), array(), array(), array(), array(), array(), 0, 0, NULL), + (3, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL), + (4, array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), + array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), 1, NULL, ''), + (5, array('x', 'y'), array(7, 8, 9), array(7L, 8L), array(7.0, 8.0), array(CAST(7 AS FLOAT)), + array(CAST(7 AS DECIMAL(10, 2))), array(DATE '2000-01-01'), array(TIMESTAMP '2000-01-01 00:00:00'), + array(X'AA'), array(true), array(TIMESTAMP_NTZ '2000-01-01 00:00:00'), 8, 7, 'y') + +-- The customer shape: drop NULL and empty-string elements. +query +SELECT id, filter(strs, x -> nvl(x, '') != '') FROM test_filter_lambda + +query +SELECT id, filter(strs, x -> x IS NOT NULL AND x != '') FROM test_filter_lambda + +query +SELECT id, filter(strs, x -> coalesce(x, '') <> '' AND length(x) > 1) FROM test_filter_lambda + +-- On a numeric array `nvl(x, '')` makes the lambda cast each element to a string. +query +SELECT id, filter(ints, x -> nvl(x, '') != '') FROM test_filter_lambda + +query +SELECT id, filter(ints, x -> x > 0) FROM test_filter_lambda + +query +SELECT id, filter(ints, x -> x % 2 = 1 OR x IS NULL) FROM test_filter_lambda + +query +SELECT id, filter(longs, x -> x > 0) FROM test_filter_lambda + +query +SELECT id, filter(dbls, x -> x > 1.0) FROM test_filter_lambda + +query +SELECT id, filter(dbls, x -> isnan(x) OR x = 0.0) FROM test_filter_lambda + +query +SELECT id, filter(flts, x -> x > 2) FROM test_filter_lambda + +query +SELECT id, filter(decs, x -> x > 0) FROM test_filter_lambda + +query +SELECT id, filter(dates, x -> x > DATE '2021-01-01') FROM test_filter_lambda + +query +SELECT id, filter(tss, x -> x > TIMESTAMP '2021-01-01 00:00:00') FROM test_filter_lambda + +query +SELECT id, filter(ntzs, x -> x > TIMESTAMP_NTZ '2021-01-01 00:00:00') FROM test_filter_lambda + +query +SELECT id, filter(bins, x -> length(x) > 0) FROM test_filter_lambda + +query +SELECT id, filter(bools, x -> x) FROM test_filter_lambda + +query +SELECT id, filter(bools, x -> NOT x) FROM test_filter_lambda + +-- Constant results: every element kept, or none. +query +SELECT id, filter(ints, x -> true) FROM test_filter_lambda + +query +SELECT id, filter(ints, x -> false) FROM test_filter_lambda + +query +SELECT id, filter(ints, x -> CAST(NULL AS BOOLEAN)) FROM test_filter_lambda + +-- Captured columns, one value per row. +query +SELECT id, filter(ints, e -> e > threshold) FROM test_filter_lambda + +query +SELECT id, filter(strs, e -> e = needle) FROM test_filter_lambda + +query +SELECT id, filter(ints, e -> e + x > threshold + 1) FROM test_filter_lambda + +-- The lambda parameter shadows a column of the same name. +query +SELECT id, filter(ints, x -> x > 0) FROM test_filter_lambda WHERE x IS NOT NULL + +-- A captured column that the lambda body does not mention, and a parameter the body ignores. +query +SELECT id, filter(ints, e -> threshold > 1) FROM test_filter_lambda + +-- Filters over a filter, and over a computed array. +query +SELECT id, filter(filter(ints, a -> a IS NOT NULL), b -> b > 0) FROM test_filter_lambda + +query +SELECT id, filter(array(x, threshold, 5), v -> v > 1) FROM test_filter_lambda + +query +SELECT id, size(filter(ints, v -> v > 0)) FROM test_filter_lambda + +-- A body that cannot run natively stays on the codegen dispatcher and must still match Spark. +query +SELECT id, filter(strs, v -> v RLIKE '^a') FROM test_filter_lambda + +-- A nested lambda runs through the dispatcher. +query +SELECT id, filter(ints, v -> exists(ints, w -> w = v + 1)) FROM test_filter_lambda + +-- A two-parameter lambda runs through the dispatcher. +query +SELECT id, filter(ints, (v, i) -> v > i) FROM test_filter_lambda diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_ansi.sql b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_ansi.sql new file mode 100644 index 00000000000..21e35a56819 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_ansi.sql @@ -0,0 +1,67 @@ +-- 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. + +-- Spark evaluates AND, OR and IN lazily, so a guard before a division keeps it from raising on +-- the elements the guard rules out. The native lambda evaluates every operand over every +-- element, so under ANSI mode a lambda that short-circuits must stay on the dispatcher unless +-- nothing in it can raise. None of these queries may raise. + +-- Config: spark.sql.ansi.enabled=true +-- ConfigMatrix: spark.comet.exec.nativeLambda.enabled=true,false + +statement +CREATE TABLE test_filter_lambda_ansi(id int, ints array, c int) USING parquet + +statement +INSERT INTO test_filter_lambda_ansi VALUES + (1, array(0, 1, NULL, 5, 100), 0), + (2, array(), 1), + (3, NULL, 2), + (4, array(0, 0), 3), + (5, array(2147483647, 0, -7), 4) + +query +SELECT id, filter(ints, x -> x != 0 AND 100 / x > 1) FROM test_filter_lambda_ansi + +query +SELECT id, filter(ints, x -> x = 0 OR 100 / x > 1) FROM test_filter_lambda_ansi + +query +SELECT id, filter(ints, x -> NOT (x = 0) AND 100 / x > 1) FROM test_filter_lambda_ansi + +query +SELECT id, filter(ints, x -> x IN (0, 100 / x)) FROM test_filter_lambda_ansi + +query +SELECT id, filter(ints, x -> x IS NOT NULL AND x > 0 AND x < 1000 AND x + 1 > 0 AND 100 / x > 1) FROM test_filter_lambda_ansi + +-- CASE, IF and COALESCE are lazy on the native side too. +query +SELECT id, filter(ints, x -> CASE WHEN x = 0 THEN false ELSE 100 / x > 1 END) FROM test_filter_lambda_ansi + +query +SELECT id, filter(ints, x -> if(x = 0, false, 100 / x > 1)) FROM test_filter_lambda_ansi + +query +SELECT id, filter(ints, x -> coalesce(x, 100 / x) > 1) FROM test_filter_lambda_ansi + +-- Nothing here can raise, so these stay native. +query +SELECT id, filter(ints, x -> x IS NOT NULL AND x > c) FROM test_filter_lambda_ansi + +query +SELECT id, filter(ints, x -> x IS NULL OR x != 0) FROM test_filter_lambda_ansi diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala index 49fadd16e0e..d9f0b7a00c4 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala @@ -23,6 +23,8 @@ import org.apache.spark.SparkConf import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.comet.udf.codegen.CometScalaUDFCodegen + /** * Higher-order function regression coverage for the codegen dispatcher. * @@ -123,4 +125,76 @@ class CometCodegenHOFSuite checkSparkAnswerAndOperator(sql(q)) } } + + private def withStringArrayTable(f: => Unit): Unit = { + withTable("t") { + sql("CREATE TABLE t (a ARRAY, c INT) USING parquet") + sql("INSERT INTO t VALUES (array('a', NULL, '', 'b'), 1), (array(), 2), (NULL, 3)") + f + } + } + + private def dispatcherActivity(f: => Unit): Long = { + CometScalaUDFCodegen.resetStats() + f + val stats = CometScalaUDFCodegen.stats() + stats.compileCount + stats.cacheHitCount + } + + test("ArrayFilter with a natively supported lambda body does not use the dispatcher") { + withStringArrayTable { + Seq( + "SELECT filter(a, x -> nvl(x, '') != '') FROM t", + "SELECT filter(a, x -> x IS NOT NULL AND length(x) > c) FROM t").foreach { query => + val native = dispatcherActivity { + checkSparkAnswerAndOperator(sql(query)) + } + assert(native == 0, s"expected no dispatcher activity for $query") + + withSQLConf(CometConf.COMET_NATIVE_LAMBDA_ENABLED.key -> "false") { + val dispatched = dispatcherActivity { + checkSparkAnswerAndOperator(sql(query)) + } + assert(dispatched >= 1, s"expected the dispatcher to run $query") + } + } + } + } + + test("ArrayFilter keeps a lambda the native side cannot evaluate on the dispatcher") { + withStringArrayTable { + Seq( + // The body calls out to the JVM, so it cannot be split from the dispatcher. + "SELECT filter(a, x -> x RLIKE '^a') FROM t", + // A nested lambda. + "SELECT filter(a, x -> exists(a, y -> y = x)) FROM t", + // A lambda of two parameters. + "SELECT filter(a, (x, i) -> i > 0) FROM t").foreach { query => + val activity = dispatcherActivity { + checkSparkAnswerAndOperator(sql(query)) + } + assert(activity >= 1, s"expected the dispatcher to run $query") + } + } + } + + test( + "ArrayFilter keeps a lambda that short-circuits an operation that can raise on the dispatcher") { + withArrayIntTable("(array(0, 1, 5)), (array()), (null)") { + val query = "SELECT filter(a, x -> x != 0 AND 100 / x > 1) FROM t" + withSQLConf("spark.sql.ansi.enabled" -> "true") { + // Spark skips `100 / x` for the element 0. The native side would raise on it. + val ansi = dispatcherActivity { + checkSparkAnswerAndOperator(sql(query)) + } + assert(ansi >= 1, "expected the dispatcher to run the lambda under ANSI mode") + } + withSQLConf("spark.sql.ansi.enabled" -> "false") { + val nonAnsi = dispatcherActivity { + checkSparkAnswerAndOperator(sql(query)) + } + assert(nonAnsi == 0, "expected the lambda to run natively when ANSI mode is off") + } + } + } } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHigherOrderFunctionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHigherOrderFunctionBenchmark.scala new file mode 100644 index 00000000000..29960b1b0cd --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHigherOrderFunctionBenchmark.scala @@ -0,0 +1,96 @@ +/* + * 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. + */ + +package org.apache.spark.sql.benchmark + +import org.apache.comet.CometConf + +// spotless:off +/** + * Benchmark to measure `filter` with a lambda, and the `array_compact` / `array_remove` forms a + * lambda can often be rewritten to. To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometHigherOrderFunctionBenchmark + * }}} + * Results will be written to "spark/benchmarks/CometHigherOrderFunctionBenchmark-**results.txt". + */ +// spotless:on +object CometHigherOrderFunctionBenchmark extends CometBenchmarkBase { + + /** + * The arrays hold ten elements. Every fifth element is NULL and every seventh is an empty + * string. + */ + private def prepareFilterTable(f: => Unit): Unit = { + withTempPath { dir => + withTempTable("parquetV1Table") { + val strElems = (0 until 10).map { i => + s"""CASE WHEN (value + $i) % 5 = 0 THEN NULL + | WHEN (value + $i) % 7 = 0 THEN '' + | ELSE cast((value + $i) % 100 as string) END""".stripMargin + } + val intElems = (0 until 10).map { i => + s"CASE WHEN (value + $i) % 5 = 0 THEN NULL ELSE cast((value + $i) % 100 as int) END" + } + prepareTable( + dir, + spark.sql(s"""SELECT + | array(${strElems.mkString(",\n")}) AS str_arr, + | array(${intElems.mkString(",\n")}) AS int_arr, + | cast(value % 3 as int) AS threshold + |FROM $tbl""".stripMargin)) + f + } + } + } + + def filterLambdaBenchmark(values: Int): Unit = { + prepareFilterTable { + Seq( + "string: filter(x -> nvl(x, '') != '')" -> + "SELECT filter(str_arr, x -> nvl(x, '') != '') FROM parquetV1Table", + "string: array_remove(array_compact(arr), '')" -> + "SELECT array_remove(array_compact(str_arr), '') FROM parquetV1Table", + "int: filter(x -> nvl(x, '') != '')" -> + "SELECT filter(int_arr, x -> nvl(x, '') != '') FROM parquetV1Table", + "int: filter(x -> x IS NOT NULL)" -> + "SELECT filter(int_arr, x -> x IS NOT NULL) FROM parquetV1Table", + "int: filter(x -> x > threshold)" -> + "SELECT filter(int_arr, x -> x > threshold) FROM parquetV1Table").foreach { + case (name, query) => + runExpressionBenchmark(name, values.toLong, query) + if (query.contains("->")) { + // The same lambda through the JVM codegen dispatcher, which is what Comet did before + // it evaluated lambdas natively. + runExpressionBenchmark( + s"$name [dispatcher]", + values.toLong, + query, + Map(CometConf.COMET_NATIVE_LAMBDA_ENABLED.key -> "false")) + } + } + } + } + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + runBenchmarkWithTable("FilterLambda", 1024 * 1024) { v => + filterLambdaBenchmark(v) + } + } +} From f543dcae0608134924de85c75ed710b562b20dd4 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 29 Oct 2026 15:26:57 -0600 Subject: [PATCH 2/4] fix: keep a lambda parameter readable by short-circuited array function arguments NullShortCircuit evaluates the later arguments of a function such as array_contains on a projection of just the columns they read. It re-indexed only Column, so a lambda parameter passed as such an argument, as in filter(ints, x -> array_contains(other, x)), was dropped from the projection and the query failed. Treat LambdaVariable like Column when projecting, and like Column in case_when's infallible check. The JVM codegen dispatcher now declines a tree that reads a lambda parameter bound outside it. Its kernel is fed only columns, so the parameter read as NULL. The filter serde no longer walks the serialized body to find such a call. Also simplify the planner arm test, share assertNoCodegen between the codegen suites, and tidy the benchmark. --- native/core/src/execution/planner.rs | 105 +++++++++++++++++- .../src/conditional_funcs/case_when.rs | 4 +- .../conditional_funcs/null_short_circuit.rs | 70 +++++++++--- .../scala/org/apache/comet/CometConf.scala | 14 +-- .../apache/comet/serde/CometScalaUDF.scala | 16 ++- .../scala/org/apache/comet/serde/arrays.scala | 27 +---- .../array/array_filter_native_lambda.sql | 20 ++++ .../array_filter_native_lambda_captured.sql | 73 ++++++++++++ .../apache/comet/CometCodegenAssertions.scala | 10 ++ .../apache/comet/CometCodegenHOFSuite.scala | 29 ++--- .../CometHigherOrderFunctionBenchmark.scala | 27 +++-- 11 files changed, 308 insertions(+), 87 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_captured.sql diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 59b03509288..ddbf4091eb4 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1269,7 +1269,7 @@ impl PhysicalPlanner { }; // The lambda body sees the enclosing schema with the parameter appended, so // `BoundReference`s keep their indexes and a `LambdaVariable` is found by name. - let mut fields: Vec = input_schema.fields().iter().cloned().collect(); + let mut fields: Vec = input_schema.fields().to_vec(); fields.push(Arc::new( element_field.as_ref().clone().with_name(&hof.param_name), )); @@ -5416,6 +5416,109 @@ mod tests { } } + #[test] + fn test_higher_order_filter() -> Result<(), DataFusionError> { + use arrow::array::ListBuilder; + use arrow::array::StringBuilder; + use datafusion_comet_proto::spark_expression::{ + BinaryExpr, BoundReference, HigherOrderFunc, LambdaVariable, + }; + + fn expr(expr_struct: ExprStruct) -> spark_expression::Expr { + spark_expression::Expr { + expr_struct: Some(expr_struct), + query_context: None, + expr_id: None, + } + } + fn bound(index: i32) -> spark_expression::Expr { + expr(Bound(BoundReference { + index, + datatype: None, + })) + } + fn variable(name: &str) -> spark_expression::Expr { + expr(ExprStruct::LambdaVariable(LambdaVariable { + name: name.to_string(), + })) + } + fn string_lists(rows: Vec>>>) -> ListArray { + let mut builder = ListBuilder::new(StringBuilder::new()); + for row in rows { + match row { + Some(values) => { + values + .into_iter() + .for_each(|v| builder.values().append_option(v)); + builder.append(true); + } + None => builder.append(false), + } + } + builder.finish() + } + + // array_filter(a, x#1 -> x#1 = c) over a: array and c: string + fn filter(func: &str, array: spark_expression::Expr) -> spark_expression::Expr { + expr(ExprStruct::HigherOrderFunc(Box::new(HigherOrderFunc { + func: func.to_string(), + array: Some(Box::new(array)), + param_name: "x#1".to_string(), + body: Some(Box::new(expr(ExprStruct::Eq(Box::new(BinaryExpr { + left: Some(Box::new(variable("x#1"))), + right: Some(Box::new(bound(1))), + }))))), + }))) + } + + let list = string_lists(vec![ + Some(vec![Some("a"), None, Some("b"), Some("a")]), + Some(vec![]), + None, + Some(vec![Some("b")]), + ]); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", list.data_type().clone(), true), + Field::new("c", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(list), + Arc::new(StringArray::from(vec!["a", "a", "a", "x"])), + ], + )?; + + let session_ctx = SessionContext::new(); + let planner = PhysicalPlanner::new(Arc::from(session_ctx), 0); + + // The lambda keeps the elements equal to the captured column, and drops NULL elements. + let physical = planner + .create_expr(&filter("array_filter", bound(0)), Arc::clone(&schema)) + .unwrap(); + let result = physical.evaluate(&batch)?.into_array(batch.num_rows())?; + let expected = string_lists(vec![ + Some(vec![Some("a"), Some("a")]), + Some(vec![]), + None, + Some(vec![]), + ]); + assert_eq!(&result, &(Arc::new(expected) as ArrayRef)); + + // The argument must be an array, and only array_filter is implemented. + let not_array = planner.create_expr(&filter("array_filter", bound(1)), Arc::clone(&schema)); + assert!(not_array + .unwrap_err() + .to_string() + .contains("expects an array argument")); + let unknown = planner.create_expr(&filter("array_zip", bound(0)), Arc::clone(&schema)); + assert!(unknown + .unwrap_err() + .to_string() + .contains("Unsupported higher-order function")); + Ok(()) + } + #[test] fn test_create_array() { let session_ctx = SessionContext::new(); diff --git a/native/spark-expr/src/conditional_funcs/case_when.rs b/native/spark-expr/src/conditional_funcs/case_when.rs index 6eed2a951a2..8830ab1b067 100644 --- a/native/spark-expr/src/conditional_funcs/case_when.rs +++ b/native/spark-expr/src/conditional_funcs/case_when.rs @@ -34,7 +34,7 @@ use datafusion::common::{ use datafusion::logical_expr::type_coercion::binary::type_union_coercion; use datafusion::logical_expr::{ColumnarValue, Operator}; use datafusion::physical_expr::expressions::{ - BinaryExpr, CaseExpr, Column, IsNotNullExpr, IsNullExpr, Literal, NotExpr, + BinaryExpr, CaseExpr, Column, IsNotNullExpr, IsNullExpr, LambdaVariable, Literal, NotExpr, }; use datafusion::physical_expr::PhysicalExpr; use std::fmt::{Display, Formatter}; @@ -451,7 +451,7 @@ pub(super) fn cannot_fail(expr: &Arc, input_schema: &Schema) - } fn infallible(expr: &Arc, input_schema: &Schema, parsing: bool) -> bool { - if expr.is::() || expr.is::() { + if expr.is::() || expr.is::() || expr.is::() { return true; } let node_is_infallible = if let Some(binary) = expr.downcast_ref::() { diff --git a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs index 232b9725cf6..b6adddb0d25 100644 --- a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs +++ b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs @@ -20,11 +20,10 @@ use arrow::array::{new_null_array, Array, BooleanArray}; use arrow::buffer::NullBuffer; use arrow::datatypes::{DataType, FieldRef, Schema}; use arrow::record_batch::{RecordBatch, RecordBatchOptions}; -use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion::common::{internal_datafusion_err, internal_err, Result}; use datafusion::logical_expr::ColumnarValue; -use datafusion::physical_expr::expressions::{Column, Literal}; -use datafusion::physical_expr::utils::collect_columns; +use datafusion::physical_expr::expressions::{Column, LambdaVariable, Literal}; use datafusion::physical_expr::PhysicalExpr; use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; @@ -281,27 +280,46 @@ impl Arg { } /// The input columns that `expr` reads, and `expr` reading them from a batch of just those columns. +/// +/// A [`LambdaVariable`] reads a column of the batch its lambda is evaluated on exactly as a +/// [`Column`] does, so it is projected and re-indexed the same way. fn project(expr: &Arc) -> Result<(Vec, Arc)> { - let mut projection = collect_columns(expr) - .iter() - .map(Column::index) - .collect::>(); + let mut projection = vec![]; + expr.apply(|e| { + if let Some(column) = e.downcast_ref::() { + projection.push(column.index()); + } else if let Some(variable) = e.downcast_ref::() { + projection.push(variable.index()); + } + Ok(TreeNodeRecursion::Continue) + })?; projection.sort_unstable(); projection.dedup(); + let position = |old: usize| { + projection + .binary_search(&old) + .map_err(|_| internal_datafusion_err!("column {old} is not projected")) + }; let projected = Arc::clone(expr) .transform_down(|e| { - let Some(column) = e.downcast_ref::() else { - return Ok(Transformed::no(e)); - }; - let index = projection - .binary_search(&column.index()) - .map_err(|_| internal_datafusion_err!("{column} is not projected"))?; - if index == column.index() { - return Ok(Transformed::no(e)); + if let Some(column) = e.downcast_ref::() { + let new = position(column.index())?; + if new != column.index() { + return Ok(Transformed::yes( + Arc::new(Column::new(column.name(), new)) as Arc + )); + } + } else if let Some(variable) = e.downcast_ref::() { + let new = position(variable.index())?; + if new != variable.index() { + return Ok(Transformed::yes(Arc::new(LambdaVariable::new( + new, + Arc::clone(variable.field()), + )) + as Arc)); + } } - Ok(Transformed::yes( - Arc::new(Column::new(column.name(), index)) as Arc, - )) + Ok(Transformed::no(e)) })? .data; Ok((projection, projected)) @@ -445,6 +463,22 @@ mod tests { assert_eq!(actual.as_ref(), expected_positions().slice(1, 4).as_ref()); } + #[test] + fn masked_argument_may_read_a_lambda_variable() { + // The string `s` is read through a lambda variable, as inside the body of a lambda. The + // masked argument is evaluated on a projection of the batch, which has to keep it. + let batch = test_batch(); + let schema = batch.schema(); + let variable = Arc::new(LambdaVariable::new(2, Arc::clone(&schema.fields()[2]))); + let element = ansi_cast(variable, DataType::Int32); + let guarded = + NullShortCircuit::wrap(array_position(col("a", &schema).unwrap(), element), &schema) + .unwrap(); + assert!(guarded.is::()); + let actual = evaluate(&guarded, &batch).unwrap(); + assert_eq!(actual.as_ref(), expected_positions().as_ref()); + } + #[test] fn later_argument_still_fails_where_the_array_is_not_null() { let batch = test_batch(); diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 6c1f5d436cd..790d7e46ef4 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -770,13 +770,13 @@ object CometConf extends ShimCometConf { val COMET_NATIVE_LAMBDA_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.exec.nativeLambda.enabled") .category(CATEGORY_EXEC) - .doc("Whether to evaluate the lambda of a higher-order function natively when its body " + - "is made only of expressions that Comet supports natively. The lambda then runs " + - "vectorized over all the elements of a batch instead of once per element in the JVM. " + - "Currently applies to `filter` over an array of a primitive, decimal, string, binary, " + - "date or timestamp element type, with a lambda of one parameter. Other lambdas run " + - "through the JVM codegen dispatcher " + - "(see `spark.comet.exec.scalaUDF.codegen.enabled`).") + .doc( + "Whether to evaluate the lambda of a higher-order function natively when its body " + + "is made only of expressions that Comet supports natively. The lambda then runs " + + "vectorized over all the elements of a batch instead of once per element in the JVM. " + + "Currently applies to `filter`. Lambdas that cannot run natively use the JVM codegen " + + "dispatcher (see `spark.comet.exec.scalaUDF.codegen.enabled`). The Expression Tuning " + + "guide lists which lambdas run natively.") .booleanConf .createWithDefault(true) diff --git a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala index 7244875e2cb..0e13d23d94c 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala @@ -22,7 +22,7 @@ package org.apache.comet.serde import scala.util.control.NonFatal import org.apache.spark.SparkEnv -import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSeq, BindReferences, Expression, ExpressionSet, If, IsNull, KnownNotNull, Literal, Or, RuntimeReplaceable, ScalaUDF} +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSeq, BindReferences, Expression, ExpressionSet, If, IsNull, KnownNotNull, LambdaFunction, Literal, NamedLambdaVariable, Or, RuntimeReplaceable, ScalaUDF} import org.apache.spark.sql.types.BinaryType import org.apache.comet.CometConf @@ -173,6 +173,20 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { case None => } + // The kernel is fed only column references, so a lambda parameter bound by a lambda outside + // `target` would never be supplied and would read as NULL. + val boundByTarget = target.collect { case l: LambdaFunction => + l.arguments.map(_.exprId) + }.flatten + target.collectFirst { + case v: NamedLambdaVariable if !boundByTarget.contains(v.exprId) => v + } match { + case Some(v) => + return Left( + s"$exprName: codegen dispatch: lambda variable ${v.name} is bound outside the expression") + case None => + } + // Bind against only the AttributeReferences the tree actually reads, so ordinals align with // the data args we ship. val attrs = target.collect { case a: AttributeReference => a }.distinct diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index bb73e73d70d..fd0cea1ec13 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -1070,7 +1070,7 @@ object CometArrayFilter extends CometExpressionSerde[ArrayFilter] with CometType private def nativeLambdaSupported(expr: ArrayFilter, body: Expression): Boolean = CometConf.COMET_NATIVE_LAMBDA_ENABLED.get() && body.deterministic && - body.find(_.isInstanceOf[HigherOrderFunction]).isEmpty && + !body.exists(_.isInstanceOf[HigherOrderFunction]) && shortCircuitSafe(body) && (expr.argument.dataType match { case ArrayType(elementType, _) => isLambdaElementType(elementType) @@ -1082,7 +1082,8 @@ object CometArrayFilter extends CometExpressionSerde[ArrayFilter] with CometType * 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. + * unless every expression in it is one that cannot raise. `infallible` in `case_when.rs` is the + * native counterpart of `neverRaises`. */ private def shortCircuitSafe(body: Expression): Boolean = !SQLConf.get.ansiEnabled || @@ -1116,11 +1117,9 @@ object CometArrayFilter extends CometExpressionSerde[ArrayFilter] with CometType binding: Boolean): Option[ExprOuterClass.Expr] = { for { arrayProto <- exprToProtoInternal(expr.argument, inputs, binding) + // A subexpression that needs the JVM codegen dispatcher cannot be part of the lambda: the + // dispatcher declines a tree that reads the lambda parameter, so the body fails to convert. bodyProto <- exprToProtoInternal(body, inputs, binding) - // A subexpression the serde routed through the JVM codegen dispatcher cannot be part of the - // lambda. Its call out to the JVM would receive the lambda parameter as an argument, and - // the whole lambda is cheaper to run in the dispatcher than to split across the boundary. - if !CometArrayFilterProto.containsJvmScalarUdf(bodyProto) } yield { val hof = ExprOuterClass.HigherOrderFunc .newBuilder() @@ -1133,22 +1132,6 @@ object CometArrayFilter extends CometExpressionSerde[ArrayFilter] with CometType } } -object CometArrayFilterProto { - - /** Whether `message`, or any message nested in it, is a call out to the JVM. */ - def containsJvmScalarUdf(message: com.google.protobuf.Message): Boolean = - message.isInstanceOf[ExprOuterClass.JvmScalarUdf] || - message.getAllFields.values.asScala.exists { - case m: com.google.protobuf.Message => containsJvmScalarUdf(m) - case l: java.util.List[_] => - l.asScala.exists { - case m: com.google.protobuf.Message => containsJvmScalarUdf(m) - case _ => false - } - case _ => false - } -} - /** * A lambda parameter inside the body of a higher-order function that Comet runs natively. It is * only reachable from such a body: the enclosing serde decides whether to serialize the body. diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda.sql b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda.sql index 5f142c834b2..b23976fc65b 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda.sql @@ -162,3 +162,23 @@ SELECT id, filter(ints, v -> exists(ints, w -> w = v + 1)) FROM test_filter_lamb -- A two-parameter lambda runs through the dispatcher. query SELECT id, filter(ints, (v, i) -> v > i) FROM test_filter_lambda + +-- Literal arrays, with and without a captured column. No input column feeds the array. +query +SELECT id, filter(array(1, NULL, 3, 5), v -> v > 1) FROM test_filter_lambda + +query +SELECT id, filter(array('a', NULL, '', 'b'), v -> nvl(v, '') != '') FROM test_filter_lambda + +query +SELECT id, filter(array(1, 2, 3, 4), v -> v > threshold) FROM test_filter_lambda + +-- Batches of one shape: every array NULL, every array empty, and no rows at all. +query +SELECT id, filter(ints, v -> v > 0) FROM test_filter_lambda WHERE id = 3 + +query +SELECT id, filter(strs, v -> nvl(v, '') != '') FROM test_filter_lambda WHERE id = 2 + +query +SELECT id, filter(ints, v -> v > 0) FROM test_filter_lambda WHERE id < 0 diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_captured.sql b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_captured.sql new file mode 100644 index 00000000000..11b98f5560a --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_captured.sql @@ -0,0 +1,73 @@ +-- 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. + +-- A lambda body that reads a captured column or struct, or calls an array function on a captured +-- array with the lambda parameter as an argument. The native side re-indexes the columns that +-- the body reads and evaluates the later arguments of those functions on a projection of them. + +-- ConfigMatrix: spark.sql.ansi.enabled=true,false + +statement +CREATE TABLE sc2(id int, ints array, other array, m map, s struct, c int) USING parquet + +statement +INSERT INTO sc2 VALUES (1, array(1, 2, 3, NULL), array(2, 3), map(1, 10, 2, 20), named_struct('a', 2, 'b', 'x'), 1), + (2, array(5, 6), array(), map(), named_struct('a', 9, 'b', 'y'), 2), + (3, NULL, array(1), NULL, NULL, 3), + (4, array(1, 2), NULL, map(2, 5), named_struct('a', 1, 'b', NULL), 4) + +query +SELECT id, filter(ints, x -> array_contains(other, x)) FROM sc2 + +query +SELECT id, filter(ints, x -> x < size(other) + 2) FROM sc2 + +query +SELECT id, filter(ints, x -> try_element_at(other, 1) = x) FROM sc2 + +query +SELECT id, filter(ints, x -> element_at(map(1, 1, 2, 2), 2) = x) FROM sc2 + +query +SELECT id, filter(ints, x -> s.a = x) FROM sc2 + +query +SELECT id, filter(ints, x -> coalesce(try_element_at(other, 1), 0) = x) FROM sc2 + +query +SELECT id, filter(ints, x -> x = c AND s.b IS NOT NULL) FROM sc2 + +query +SELECT id, filter(ints, x -> m[x] > 5) FROM sc2 + +query +SELECT id, filter(ints, x -> array_position(other, x) > 0) FROM sc2 + +query +SELECT id, filter(ints, x -> array_contains(other, x + 1)) FROM sc2 + +query +SELECT id, filter(ints, x -> array_contains(other, cast(x as bigint))) FROM sc2 + +query +SELECT id, filter(ints, x -> try_element_at(other, x) IS NOT NULL) FROM sc2 + +query +SELECT id, filter(ints, x -> CASE WHEN x > c THEN x < size(other) ELSE array_contains(other, x) END) FROM sc2 + +query +SELECT id, filter(ints, x -> arrays_overlap(other, array(x, c))) FROM sc2 diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenAssertions.scala b/spark/src/test/scala/org/apache/comet/CometCodegenAssertions.scala index bce8bfc5986..fcd81dc6925 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenAssertions.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenAssertions.scala @@ -42,6 +42,16 @@ trait CometCodegenAssertions { s"expected codegen dispatcher activity, got $after") } + /** Asserts the dispatcher did not run during `f`, guarding against a silent fallback to it. */ + protected def assertNoCodegen(f: => Unit): Unit = { + CometScalaUDFCodegen.resetStats() + f + val after = CometScalaUDFCodegen.stats() + assert( + after.compileCount + after.cacheHitCount == 0, + s"expected no codegen dispatcher activity, got $after") + } + /** * Asserts the composed subtree fused into one kernel signature, not N (one per sub-expression). * Uses the JVM-wide signature set rather than `compileCount` because per-task `boundExpr` diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala index d9f0b7a00c4..100ba0b6158 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala @@ -23,8 +23,6 @@ import org.apache.spark.SparkConf import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper -import org.apache.comet.udf.codegen.CometScalaUDFCodegen - /** * Higher-order function regression coverage for the codegen dispatcher. * @@ -134,28 +132,18 @@ class CometCodegenHOFSuite } } - private def dispatcherActivity(f: => Unit): Long = { - CometScalaUDFCodegen.resetStats() - f - val stats = CometScalaUDFCodegen.stats() - stats.compileCount + stats.cacheHitCount - } - test("ArrayFilter with a natively supported lambda body does not use the dispatcher") { withStringArrayTable { Seq( "SELECT filter(a, x -> nvl(x, '') != '') FROM t", "SELECT filter(a, x -> x IS NOT NULL AND length(x) > c) FROM t").foreach { query => - val native = dispatcherActivity { + assertNoCodegen { checkSparkAnswerAndOperator(sql(query)) } - assert(native == 0, s"expected no dispatcher activity for $query") - withSQLConf(CometConf.COMET_NATIVE_LAMBDA_ENABLED.key -> "false") { - val dispatched = dispatcherActivity { + assertCodegenRan { checkSparkAnswerAndOperator(sql(query)) } - assert(dispatched >= 1, s"expected the dispatcher to run $query") } } } @@ -164,16 +152,15 @@ class CometCodegenHOFSuite test("ArrayFilter keeps a lambda the native side cannot evaluate on the dispatcher") { withStringArrayTable { Seq( - // The body calls out to the JVM, so it cannot be split from the dispatcher. + // The body needs the dispatcher, which cannot see the lambda parameter. "SELECT filter(a, x -> x RLIKE '^a') FROM t", // A nested lambda. "SELECT filter(a, x -> exists(a, y -> y = x)) FROM t", // A lambda of two parameters. "SELECT filter(a, (x, i) -> i > 0) FROM t").foreach { query => - val activity = dispatcherActivity { + assertCodegenRan { checkSparkAnswerAndOperator(sql(query)) } - assert(activity >= 1, s"expected the dispatcher to run $query") } } } @@ -182,18 +169,16 @@ class CometCodegenHOFSuite "ArrayFilter keeps a lambda that short-circuits an operation that can raise on the dispatcher") { withArrayIntTable("(array(0, 1, 5)), (array()), (null)") { val query = "SELECT filter(a, x -> x != 0 AND 100 / x > 1) FROM t" + // Spark skips `100 / x` for the element 0. The native side would raise on it. withSQLConf("spark.sql.ansi.enabled" -> "true") { - // Spark skips `100 / x` for the element 0. The native side would raise on it. - val ansi = dispatcherActivity { + assertCodegenRan { checkSparkAnswerAndOperator(sql(query)) } - assert(ansi >= 1, "expected the dispatcher to run the lambda under ANSI mode") } withSQLConf("spark.sql.ansi.enabled" -> "false") { - val nonAnsi = dispatcherActivity { + assertNoCodegen { checkSparkAnswerAndOperator(sql(query)) } - assert(nonAnsi == 0, "expected the lambda to run natively when ANSI mode is off") } } } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHigherOrderFunctionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHigherOrderFunctionBenchmark.scala index 29960b1b0cd..741b624cc8f 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHigherOrderFunctionBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHigherOrderFunctionBenchmark.scala @@ -65,8 +65,6 @@ object CometHigherOrderFunctionBenchmark extends CometBenchmarkBase { Seq( "string: filter(x -> nvl(x, '') != '')" -> "SELECT filter(str_arr, x -> nvl(x, '') != '') FROM parquetV1Table", - "string: array_remove(array_compact(arr), '')" -> - "SELECT array_remove(array_compact(str_arr), '') FROM parquetV1Table", "int: filter(x -> nvl(x, '') != '')" -> "SELECT filter(int_arr, x -> nvl(x, '') != '') FROM parquetV1Table", "int: filter(x -> x IS NOT NULL)" -> @@ -75,22 +73,23 @@ object CometHigherOrderFunctionBenchmark extends CometBenchmarkBase { "SELECT filter(int_arr, x -> x > threshold) FROM parquetV1Table").foreach { case (name, query) => runExpressionBenchmark(name, values.toLong, query) - if (query.contains("->")) { - // The same lambda through the JVM codegen dispatcher, which is what Comet did before - // it evaluated lambdas natively. - runExpressionBenchmark( - s"$name [dispatcher]", - values.toLong, - query, - Map(CometConf.COMET_NATIVE_LAMBDA_ENABLED.key -> "false")) - } + // The same lambda through the JVM codegen dispatcher, which is what Comet did before it + // evaluated lambdas natively. + runExpressionBenchmark( + s"$name [dispatcher]", + values.toLong, + query, + Map(CometConf.COMET_NATIVE_LAMBDA_ENABLED.key -> "false")) } + // The built-in functions that replace a lambda have no native-lambda variant to compare. + runExpressionBenchmark( + "string: array_remove(array_compact(arr), '')", + values.toLong, + "SELECT array_remove(array_compact(str_arr), '') FROM parquetV1Table") } } override def runCometBenchmark(mainArgs: Array[String]): Unit = { - runBenchmarkWithTable("FilterLambda", 1024 * 1024) { v => - filterLambdaBenchmark(v) - } + runBenchmarkWithTable("FilterLambda", 1024 * 1024)(filterLambdaBenchmark) } } From a16bdba0c57c1e2a03c1c249a08cf011c8e1add1 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 29 Oct 2026 15:44:49 -0600 Subject: [PATCH 3/4] fix: decline the native filter lambda for wide captures and for bodies that can raise in any mode DataFusion copies each column a lambda captures once per element of the array, so a captured string, binary or nested column is copied as many times as the array has elements. A body that reads one now stays on the dispatcher. The check for a body that short-circuits (AND, OR, IN) no longer depends on the ANSI setting. It classifies each expression: arithmetic and casts by the eval mode they carry, and anything not known to be safe, such as element_at, which raises for the index 0 in any mode, as able to raise. The SQL tests pin the path of each query with expect_native and expect_dispatch, replacing the ConfigMatrix over the native lambda flag and the codegen suite tests, and cover tinyint and smallint elements. --- .../user-guide/latest/tuning/expressions.md | 10 +- .../scala/org/apache/comet/serde/arrays.scala | 70 ++++++++---- .../array/array_filter_native_lambda.sql | 100 ++++++++++-------- .../array/array_filter_native_lambda_ansi.sql | 21 ++-- .../array_filter_native_lambda_captured.sql | 51 ++++++--- .../array_filter_native_lambda_disabled.sql | 32 ++++++ .../apache/comet/CometCodegenAssertions.scala | 10 -- .../apache/comet/CometCodegenHOFSuite.scala | 59 ----------- 8 files changed, 193 insertions(+), 160 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_disabled.sql diff --git a/docs/source/user-guide/latest/tuning/expressions.md b/docs/source/user-guide/latest/tuning/expressions.md index 1a22c4fbd13..134e439cc82 100644 --- a/docs/source/user-guide/latest/tuning/expressions.md +++ b/docs/source/user-guide/latest/tuning/expressions.md @@ -48,11 +48,15 @@ To get the most out of Comet when a query uses lambdas: UDF or another lambda cannot be split from the dispatcher, so the whole lambda runs there. Use [`spark.comet.explain.codegen.enabled`](../understanding-comet-plans.md#sparkcometexplaincodegenenabled) to see which expressions of a query went through the dispatcher. -- **Under ANSI mode, keep `AND`, `OR` and `IN` away from operations that can raise.** Spark evaluates them lazily, so +- **Keep `AND`, `OR` and `IN` away from operations that can raise.** Spark evaluates them lazily, so `x != 0 AND 100 / x > 1` never divides by zero, but the native lambda evaluates every operand for every element. A - body that combines them with an operation that can raise, such as a division, an addition that can overflow or a - cast, stays on the dispatcher. `CASE WHEN`, `IF` and `COALESCE` are lazy natively too, so + body that combines them with an operation that can raise stays on the dispatcher. That covers a division, an + addition or a cast under ANSI mode, which raise only in that mode, and `element_at`, which raises for the index 0 in + any mode. `CASE WHEN`, `IF` and `COALESCE` are lazy natively too, so `CASE WHEN x = 0 THEN false ELSE 100 / x > 1 END` runs natively. +- **Capture only numeric, boolean, date or timestamp columns.** DataFusion copies each column the lambda reads once per + element of the array before it evaluates the lambda, so a captured string, binary, array, struct or map column is + copied as many times as the array has elements. A lambda that reads one stays on the dispatcher. - **Use a lambda of one parameter.** `filter(arr, (x, i) -> i > 0)`, which also receives the element's index, runs through the dispatcher. diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index fd0cea1ec13..88ca3089c6e 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -23,7 +23,7 @@ import scala.annotation.tailrec import scala.jdk.CollectionConverters._ import org.apache.spark.SPARK_VERSION -import org.apache.spark.sql.catalyst.expressions.{And, ArrayAggregate, ArrayAppend, ArrayContains, ArrayDistinct, ArrayExcept, ArrayExists, ArrayFilter, ArrayForAll, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayPosition, ArrayRemove, ArraySort, ArraysOverlap, ArraysZip, ArrayTransform, ArrayUnion, Attribute, AttributeReference, BinaryComparison, BoundReference, CaseWhen, Cast, Coalesce, Contains, CreateArray, ElementAt, EmptyRow, EndsWith, Expression, Flatten, GetArrayItem, GreaterThan, HigherOrderFunction, If, In, InSet, IsNaN, IsNotNull, IsNull, LambdaFunction, Length, LessThan, Literal, Lower, NamedLambdaVariable, Not, Or, Reverse, Sequence, Size, Slice, SortArray, StartsWith, Upper, ZipWith} +import org.apache.spark.sql.catalyst.expressions.{Add, And, ArrayAggregate, ArrayAppend, ArrayContains, ArrayDistinct, ArrayExcept, ArrayExists, ArrayFilter, ArrayForAll, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayPosition, ArrayRemove, ArraySort, ArraysOverlap, ArraysZip, ArrayTransform, ArrayUnion, Attribute, AttributeReference, BinaryComparison, BoundReference, CaseWhen, Cast, Coalesce, Contains, CreateArray, Divide, ElementAt, EmptyRow, EndsWith, EvalMode, Expression, Flatten, GetArrayItem, GreaterThan, HigherOrderFunction, If, In, InSet, IntegralDivide, IsNaN, IsNotNull, IsNull, LambdaFunction, Length, LessThan, Literal, Lower, Multiply, NamedLambdaVariable, Not, Or, Remainder, Reverse, Sequence, Size, Slice, SortArray, StartsWith, Substring, Subtract, Upper, ZipWith} import org.apache.spark.sql.catalyst.util.GenericArrayData import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.internal.SQLConf @@ -32,6 +32,7 @@ import org.apache.spark.sql.types._ import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.withFallbackReason import org.apache.comet.DataTypeSupport.{deepNullable, isComplexType} +import org.apache.comet.expressions.CometEvalMode import org.apache.comet.serde.QueryPlanSerde._ import org.apache.comet.shims.{CometExprShim, CometTypeShim} @@ -1027,7 +1028,10 @@ object CometFlatten extends CometExpressionSerde[Flatten] with ArraysBase { } } -object CometArrayFilter extends CometExpressionSerde[ArrayFilter] with CometTypeShim { +object CometArrayFilter + extends CometExpressionSerde[ArrayFilter] + with CometTypeShim + with CometExprShim { override def getSupportLevel(expr: ArrayFilter): SupportLevel = Compatible() @@ -1071,42 +1075,72 @@ object CometArrayFilter extends CometExpressionSerde[ArrayFilter] with CometType CometConf.COMET_NATIVE_LAMBDA_ENABLED.get() && body.deterministic && !body.exists(_.isInstanceOf[HigherOrderFunction]) && + capturesOnlyFixedWidthColumns(body) && shortCircuitSafe(body) && (expr.argument.dataType match { case ArrayType(elementType, _) => isLambdaElementType(elementType) case _ => false }) + /** + * DataFusion copies each column the body captures once per element of its row's array before it + * evaluates the body, so a captured string, binary or nested value is copied as many times as + * the array has elements, and those copies are not reserved with the memory pool. + */ + private def capturesOnlyFixedWidthColumns(body: Expression): Boolean = + body + .collect { + case a: AttributeReference => a.dataType + case b: BoundReference => b.dataType + } + .forall(isFixedWidth) + /** * 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. `infallible` in `case_when.rs` is the - * native counterpart of `neverRaises`. + * never divides by zero. The native side evaluates every operand over the whole batch, so the + * same lambda raises on the element Spark skips. `CASE WHEN`, `IF` and `COALESCE` are not + * affected. A body that short-circuits stays on the dispatcher unless none of its expressions + * can raise. */ private def shortCircuitSafe(body: Expression): Boolean = - !SQLConf.get.ansiEnabled || - !body.exists { - case _: And | _: Or | _: In | _: InSet => true - case _ => false - } || !body.exists(!neverRaises(_)) + !body.exists { + case _: And | _: Or | _: In | _: InSet => true + case _ => false + } || !body.exists(mayRaise) - private def neverRaises(e: Expression): Boolean = e match { + /** + * Whether `e` can raise on an input value. This fails closed: an expression is assumed to raise + * unless it is listed here, since some raise in any mode (`element_at` with the index 0), and + * arithmetic and casts raise only in the ANSI eval mode they carry. `infallible` in + * `case_when.rs` is the native counterpart. + */ + private def mayRaise(e: Expression): Boolean = e match { case _: Literal | _: NamedLambdaVariable | _: AttributeReference | _: BoundReference | _: BinaryComparison | _: IsNull | _: IsNotNull | _: Not | _: And | _: Or | _: In | _: InSet | _: IsNaN | _: Coalesce | _: If | _: CaseWhen | _: Length | _: Upper | - _: Lower | _: StartsWith | _: EndsWith | _: Contains => + _: Lower | _: Substring | _: StartsWith | _: EndsWith | _: Contains => + false + case c: Cast => evalMode(c) == CometEvalMode.ANSI + case a: Add => a.evalMode == EvalMode.ANSI + case a: Subtract => a.evalMode == EvalMode.ANSI + case a: Multiply => a.evalMode == EvalMode.ANSI + case a: Divide => a.evalMode == EvalMode.ANSI + case a: IntegralDivide => a.evalMode == EvalMode.ANSI + case a: Remainder => a.evalMode == EvalMode.ANSI + case _ => true + } + + private def isFixedWidth(dt: DataType): Boolean = dt match { + case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType | + DateType | TimestampType | TimestampNTZType | _: DecimalType => true case _ => false } private def isLambdaElementType(dt: DataType): Boolean = dt match { - case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType | - DateType | TimestampType | TimestampNTZType | BinaryType | _: DecimalType => - true + case BinaryType => true case _: StringType => !hasNonDefaultStringCollation(dt) - case _ => false + case _ => isFixedWidth(dt) } private def convertLambda( diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda.sql b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda.sql index b23976fc65b..bfe07fb050e 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda.sql @@ -19,7 +19,7 @@ -- `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 +-- Config: spark.sql.ansi.enabled=false statement CREATE TABLE test_filter_lambda( @@ -35,6 +35,8 @@ CREATE TABLE test_filter_lambda( bins array, bools array, ntzs array, + tinyints array, + smallints array, x int, threshold int, needle string @@ -51,134 +53,148 @@ INSERT INTO test_filter_lambda VALUES array(X'01', NULL, X'', X'FF'), array(true, NULL, false), array(TIMESTAMP_NTZ '2020-01-01 00:00:00', NULL, TIMESTAMP_NTZ '2024-02-29 12:34:56'), + array(CAST(1 AS TINYINT), NULL, CAST(-128 AS TINYINT), CAST(127 AS TINYINT)), + array(CAST(1 AS SMALLINT), NULL, CAST(-32768 AS SMALLINT), CAST(32767 AS SMALLINT)), 3, 2, 'a'), - (2, array(), array(), array(), array(), array(), array(), array(), array(), array(), array(), array(), 0, 0, NULL), - (3, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL), + (2, array(), array(), array(), array(), array(), array(), array(), array(), array(), array(), array(), array(), array(), 0, 0, NULL), + (3, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL), (4, array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), - array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), 1, NULL, ''), + array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), + array(NULL, NULL), array(NULL, NULL), 1, NULL, ''), (5, array('x', 'y'), array(7, 8, 9), array(7L, 8L), array(7.0, 8.0), array(CAST(7 AS FLOAT)), array(CAST(7 AS DECIMAL(10, 2))), array(DATE '2000-01-01'), array(TIMESTAMP '2000-01-01 00:00:00'), - array(X'AA'), array(true), array(TIMESTAMP_NTZ '2000-01-01 00:00:00'), 8, 7, 'y') + array(X'AA'), array(true), array(TIMESTAMP_NTZ '2000-01-01 00:00:00'), array(CAST(7 AS TINYINT)), array(CAST(7 AS SMALLINT)), 8, 7, 'y') -- The customer shape: drop NULL and empty-string elements. -query +query expect_native(filter) SELECT id, filter(strs, x -> nvl(x, '') != '') FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(strs, x -> x IS NOT NULL AND x != '') FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(strs, x -> coalesce(x, '') <> '' AND length(x) > 1) FROM test_filter_lambda -- On a numeric array `nvl(x, '')` makes the lambda cast each element to a string. -query +query expect_native(filter) SELECT id, filter(ints, x -> nvl(x, '') != '') FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(ints, x -> x > 0) FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(ints, x -> x % 2 = 1 OR x IS NULL) FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(longs, x -> x > 0) FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(dbls, x -> x > 1.0) FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(dbls, x -> isnan(x) OR x = 0.0) FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(flts, x -> x > 2) FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(decs, x -> x > 0) FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(dates, x -> x > DATE '2021-01-01') FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(tss, x -> x > TIMESTAMP '2021-01-01 00:00:00') FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(ntzs, x -> x > TIMESTAMP_NTZ '2021-01-01 00:00:00') FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(bins, x -> length(x) > 0) FROM test_filter_lambda -query +query expect_native(filter) +SELECT id, filter(tinyints, x -> x > 0) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, filter(smallints, x -> x < 0 OR x IS NULL) FROM test_filter_lambda + +query expect_native(filter) SELECT id, filter(bools, x -> x) FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(bools, x -> NOT x) FROM test_filter_lambda -- Constant results: every element kept, or none. -query +query expect_native(filter) SELECT id, filter(ints, x -> true) FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(ints, x -> false) FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(ints, x -> CAST(NULL AS BOOLEAN)) FROM test_filter_lambda -- Captured columns, one value per row. -query +query expect_native(filter) SELECT id, filter(ints, e -> e > threshold) FROM test_filter_lambda -query +query expect_dispatch(filter) SELECT id, filter(strs, e -> e = needle) FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(ints, e -> e + x > threshold + 1) FROM test_filter_lambda -- The lambda parameter shadows a column of the same name. -query +query expect_native(filter) SELECT id, filter(ints, x -> x > 0) FROM test_filter_lambda WHERE x IS NOT NULL -- A captured column that the lambda body does not mention, and a parameter the body ignores. -query +query expect_native(filter) SELECT id, filter(ints, e -> threshold > 1) FROM test_filter_lambda -- Filters over a filter, and over a computed array. -query +query expect_native(filter) SELECT id, filter(filter(ints, a -> a IS NOT NULL), b -> b > 0) FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(array(x, threshold, 5), v -> v > 1) FROM test_filter_lambda -query +query expect_native(filter) SELECT id, size(filter(ints, v -> v > 0)) FROM test_filter_lambda -- A body that cannot run natively stays on the codegen dispatcher and must still match Spark. -query +query expect_dispatch(filter) SELECT id, filter(strs, v -> v RLIKE '^a') FROM test_filter_lambda -- A nested lambda runs through the dispatcher. -query +query expect_dispatch(filter) SELECT id, filter(ints, v -> exists(ints, w -> w = v + 1)) FROM test_filter_lambda -- A two-parameter lambda runs through the dispatcher. -query +query expect_dispatch(filter) SELECT id, filter(ints, (v, i) -> v > i) FROM test_filter_lambda -- Literal arrays, with and without a captured column. No input column feeds the array. -query +query expect_native(filter) SELECT id, filter(array(1, NULL, 3, 5), v -> v > 1) FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(array('a', NULL, '', 'b'), v -> nvl(v, '') != '') FROM test_filter_lambda -query +query expect_native(filter) SELECT id, filter(array(1, 2, 3, 4), v -> v > threshold) FROM test_filter_lambda -- Batches of one shape: every array NULL, every array empty, and no rows at all. -query +query expect_native(filter) SELECT id, filter(ints, v -> v > 0) FROM test_filter_lambda WHERE id = 3 -query +query expect_native(filter) SELECT id, filter(strs, v -> nvl(v, '') != '') FROM test_filter_lambda WHERE id = 2 -query +query expect_native(filter) SELECT id, filter(ints, v -> v > 0) FROM test_filter_lambda WHERE id < 0 + +-- `element_at` raises for the index 0 whatever the ANSI setting, and Spark skips it for the +-- elements `x > 0` rules out, so this lambda stays on the dispatcher. +query expect_dispatch(filter) +SELECT id, filter(ints, v -> v > 0 AND element_at(array(threshold, 5), v) > 1) FROM test_filter_lambda diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_ansi.sql b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_ansi.sql index 21e35a56819..80a0390bb39 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_ansi.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_ansi.sql @@ -21,7 +21,6 @@ -- nothing in it can raise. None of these queries may raise. -- Config: spark.sql.ansi.enabled=true --- ConfigMatrix: spark.comet.exec.nativeLambda.enabled=true,false statement CREATE TABLE test_filter_lambda_ansi(id int, ints array, c int) USING parquet @@ -34,34 +33,34 @@ INSERT INTO test_filter_lambda_ansi VALUES (4, array(0, 0), 3), (5, array(2147483647, 0, -7), 4) -query +query expect_dispatch(filter) SELECT id, filter(ints, x -> x != 0 AND 100 / x > 1) FROM test_filter_lambda_ansi -query +query expect_dispatch(filter) SELECT id, filter(ints, x -> x = 0 OR 100 / x > 1) FROM test_filter_lambda_ansi -query +query expect_dispatch(filter) SELECT id, filter(ints, x -> NOT (x = 0) AND 100 / x > 1) FROM test_filter_lambda_ansi -query +query expect_dispatch(filter) SELECT id, filter(ints, x -> x IN (0, 100 / x)) FROM test_filter_lambda_ansi -query +query expect_dispatch(filter) SELECT id, filter(ints, x -> x IS NOT NULL AND x > 0 AND x < 1000 AND x + 1 > 0 AND 100 / x > 1) FROM test_filter_lambda_ansi -- CASE, IF and COALESCE are lazy on the native side too. -query +query expect_native(filter) SELECT id, filter(ints, x -> CASE WHEN x = 0 THEN false ELSE 100 / x > 1 END) FROM test_filter_lambda_ansi -query +query expect_native(filter) SELECT id, filter(ints, x -> if(x = 0, false, 100 / x > 1)) FROM test_filter_lambda_ansi -query +query expect_native(filter) SELECT id, filter(ints, x -> coalesce(x, 100 / x) > 1) FROM test_filter_lambda_ansi -- Nothing here can raise, so these stay native. -query +query expect_native(filter) SELECT id, filter(ints, x -> x IS NOT NULL AND x > c) FROM test_filter_lambda_ansi -query +query expect_native(filter) SELECT id, filter(ints, x -> x IS NULL OR x != 0) FROM test_filter_lambda_ansi diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_captured.sql b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_captured.sql index 11b98f5560a..6d3ece99e28 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_captured.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_captured.sql @@ -15,9 +15,11 @@ -- specific language governing permissions and limitations -- under the License. --- A lambda body that reads a captured column or struct, or calls an array function on a captured --- array with the lambda parameter as an argument. The native side re-indexes the columns that --- the body reads and evaluates the later arguments of those functions on a projection of them. +-- A lambda body that reads a captured column, or calls an array function with the lambda +-- parameter as an argument. The native side re-indexes the columns that the body reads and +-- evaluates the later arguments of those functions on a projection of them. A captured string, +-- binary or nested column is copied once per element by DataFusion, so a body that reads one +-- stays on the dispatcher. -- ConfigMatrix: spark.sql.ansi.enabled=true,false @@ -30,44 +32,59 @@ INSERT INTO sc2 VALUES (1, array(1, 2, 3, NULL), array(2, 3), map(1, 10, 2, 20), (3, NULL, array(1), NULL, NULL, 3), (4, array(1, 2), NULL, map(2, 5), named_struct('a', 1, 'b', NULL), 4) -query +query expect_dispatch(filter) SELECT id, filter(ints, x -> array_contains(other, x)) FROM sc2 -query +query expect_dispatch(filter) SELECT id, filter(ints, x -> x < size(other) + 2) FROM sc2 -query +query expect_dispatch(filter) SELECT id, filter(ints, x -> try_element_at(other, 1) = x) FROM sc2 -query +query expect_native(filter) SELECT id, filter(ints, x -> element_at(map(1, 1, 2, 2), 2) = x) FROM sc2 -query +query expect_dispatch(filter) SELECT id, filter(ints, x -> s.a = x) FROM sc2 -query +query expect_dispatch(filter) SELECT id, filter(ints, x -> coalesce(try_element_at(other, 1), 0) = x) FROM sc2 -query +query expect_dispatch(filter) SELECT id, filter(ints, x -> x = c AND s.b IS NOT NULL) FROM sc2 -query +query expect_dispatch(filter) SELECT id, filter(ints, x -> m[x] > 5) FROM sc2 -query +query expect_dispatch(filter) SELECT id, filter(ints, x -> array_position(other, x) > 0) FROM sc2 -query +query expect_dispatch(filter) SELECT id, filter(ints, x -> array_contains(other, x + 1)) FROM sc2 -query +query expect_dispatch(filter) SELECT id, filter(ints, x -> array_contains(other, cast(x as bigint))) FROM sc2 -query +query expect_dispatch(filter) SELECT id, filter(ints, x -> try_element_at(other, x) IS NOT NULL) FROM sc2 -query +query expect_dispatch(filter) SELECT id, filter(ints, x -> CASE WHEN x > c THEN x < size(other) ELSE array_contains(other, x) END) FROM sc2 -query +query expect_dispatch(filter) SELECT id, filter(ints, x -> arrays_overlap(other, array(x, c))) FROM sc2 + +query expect_native(filter) +SELECT id, filter(ints, x -> array_contains(array(c, 2, 3), x)) FROM sc2 + +query expect_native(filter) +SELECT id, filter(ints, x -> array_contains(array(c, 2, 3), x + 1)) FROM sc2 + +query expect_native(filter) +SELECT id, filter(ints, x -> array_position(array(c, 2), x) > 0) FROM sc2 + +query expect_native(filter) +SELECT id, filter(ints, x -> try_element_at(array(c, 5), x) > 1) FROM sc2 + +query expect_native(filter) +SELECT id, filter(ints, x -> arrays_overlap(array(c, 2), array(x, c))) FROM sc2 diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_disabled.sql b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_disabled.sql new file mode 100644 index 00000000000..639f89a85c6 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_disabled.sql @@ -0,0 +1,32 @@ +-- 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. + +-- With the native lambda off, every lambda runs on the dispatcher. + +-- Config: spark.comet.exec.nativeLambda.enabled=false + +statement +CREATE TABLE test_filter_lambda_off(id int, strs array, ints array, c int) USING parquet + +statement +INSERT INTO test_filter_lambda_off VALUES (1, array('a', NULL, '', 'b'), array(1, NULL, 3), 2), (2, array(), array(), 1), (3, NULL, NULL, 0) + +query expect_dispatch(filter) +SELECT id, filter(strs, x -> nvl(x, '') != '') FROM test_filter_lambda_off + +query expect_dispatch(filter) +SELECT id, filter(ints, x -> x > c) FROM test_filter_lambda_off diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenAssertions.scala b/spark/src/test/scala/org/apache/comet/CometCodegenAssertions.scala index fcd81dc6925..bce8bfc5986 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenAssertions.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenAssertions.scala @@ -42,16 +42,6 @@ trait CometCodegenAssertions { s"expected codegen dispatcher activity, got $after") } - /** Asserts the dispatcher did not run during `f`, guarding against a silent fallback to it. */ - protected def assertNoCodegen(f: => Unit): Unit = { - CometScalaUDFCodegen.resetStats() - f - val after = CometScalaUDFCodegen.stats() - assert( - after.compileCount + after.cacheHitCount == 0, - s"expected no codegen dispatcher activity, got $after") - } - /** * Asserts the composed subtree fused into one kernel signature, not N (one per sub-expression). * Uses the JVM-wide signature set rather than `compileCount` because per-task `boundExpr` diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala index 100ba0b6158..49fadd16e0e 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala @@ -123,63 +123,4 @@ class CometCodegenHOFSuite checkSparkAnswerAndOperator(sql(q)) } } - - private def withStringArrayTable(f: => Unit): Unit = { - withTable("t") { - sql("CREATE TABLE t (a ARRAY, c INT) USING parquet") - sql("INSERT INTO t VALUES (array('a', NULL, '', 'b'), 1), (array(), 2), (NULL, 3)") - f - } - } - - test("ArrayFilter with a natively supported lambda body does not use the dispatcher") { - withStringArrayTable { - Seq( - "SELECT filter(a, x -> nvl(x, '') != '') FROM t", - "SELECT filter(a, x -> x IS NOT NULL AND length(x) > c) FROM t").foreach { query => - assertNoCodegen { - checkSparkAnswerAndOperator(sql(query)) - } - withSQLConf(CometConf.COMET_NATIVE_LAMBDA_ENABLED.key -> "false") { - assertCodegenRan { - checkSparkAnswerAndOperator(sql(query)) - } - } - } - } - } - - test("ArrayFilter keeps a lambda the native side cannot evaluate on the dispatcher") { - withStringArrayTable { - Seq( - // The body needs the dispatcher, which cannot see the lambda parameter. - "SELECT filter(a, x -> x RLIKE '^a') FROM t", - // A nested lambda. - "SELECT filter(a, x -> exists(a, y -> y = x)) FROM t", - // A lambda of two parameters. - "SELECT filter(a, (x, i) -> i > 0) FROM t").foreach { query => - assertCodegenRan { - checkSparkAnswerAndOperator(sql(query)) - } - } - } - } - - test( - "ArrayFilter keeps a lambda that short-circuits an operation that can raise on the dispatcher") { - withArrayIntTable("(array(0, 1, 5)), (array()), (null)") { - val query = "SELECT filter(a, x -> x != 0 AND 100 / x > 1) FROM t" - // Spark skips `100 / x` for the element 0. The native side would raise on it. - withSQLConf("spark.sql.ansi.enabled" -> "true") { - assertCodegenRan { - checkSparkAnswerAndOperator(sql(query)) - } - } - withSQLConf("spark.sql.ansi.enabled" -> "false") { - assertNoCodegen { - checkSparkAnswerAndOperator(sql(query)) - } - } - } - } } From 7bc084ccc424ac142101fc908edeb07cd4275c19 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 1 Nov 2026 08:30:37 -0700 Subject: [PATCH 4/4] feat: run nvl(x, '') != '' on a numeric array as array_compact On a numeric, date or timestamp array Spark plans the customer's rule filter(arr, x -> nvl(x, '') != '') 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 lambda is x IS NOT NULL. Send it down the array_compact fast path instead of casting every element to a string. On 2M rows of ten-element int arrays (every fifth element NULL), Spark 3.5.5 local[4], best of 5: 0.453 s on the dispatcher, 0.281 s with the native lambda, 0.086 s with this change, against 0.220 s for Spark. --- .../user-guide/latest/tuning/expressions.md | 7 +- .../scala/org/apache/comet/serde/arrays.scala | 36 +++++- .../array/array_filter_numeric_nvl.sql | 121 ++++++++++++++++++ 3 files changed, 160 insertions(+), 4 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_filter_numeric_nvl.sql diff --git a/docs/source/user-guide/latest/tuning/expressions.md b/docs/source/user-guide/latest/tuning/expressions.md index 134e439cc82..1ee62b48d69 100644 --- a/docs/source/user-guide/latest/tuning/expressions.md +++ b/docs/source/user-guide/latest/tuning/expressions.md @@ -26,7 +26,7 @@ A higher-order function takes a lambda: `filter(arr, x -> ...)`, `transform`, `e | Lambda | How Comet runs it | | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `filter(arr, x -> x IS NOT NULL)` (what `array_compact` expands to) | Natively, as `array_compact`. | +| `filter(arr, x -> x IS NOT NULL)` (what `array_compact` expands to), and `filter(arr, x -> nvl(x, '') != '')` on an array of numbers, dates or timestamps | Natively, as `array_compact`. | | `filter` with a one-parameter lambda whose body uses only natively supported expressions, over an array of primitive, decimal, string, binary, date or timestamp elements | Natively. The lambda runs once over all the elements of a batch, vectorized, instead of once per element. | | Any other lambda, and every other higher-order function | Through the [JVM codegen dispatcher](../scala_java_udfs.md). Spark's interpreter evaluates the lambda for each element inside a Comet kernel, so expect Spark's speed, not more. | @@ -43,7 +43,9 @@ To get the most out of Comet when a query uses lambdas: strings with the default collation. - **Do not test a numeric array for nulls with a string comparison.** With ANSI mode off, Spark casts every element of `filter(int_arr, x -> nvl(x, '') != '')` to a string to compare it with `''`; with ANSI mode on the query does - not analyze. Write `x IS NOT NULL`, or call `array_compact`. + not analyze. Comet recognizes exactly this form, with either operand first, and runs it as `array_compact`, which + never casts. A variation, such as another default value or an extra condition, runs as a general lambda and does + cast every element. Write `x IS NOT NULL`, or call `array_compact`. - **Keep the lambda body to natively supported expressions.** A body that calls a regular expression function, a UDF or another lambda cannot be split from the dispatcher, so the whole lambda runs there. Use [`spark.comet.explain.codegen.enabled`](../understanding-comet-plans.md#sparkcometexplaincodegenenabled) to see which @@ -72,7 +74,6 @@ To get the most out of Comet when a query uses lambdas: | ----------------------------------------------------------- | ------ | ----------------------- | -------------- | | `filter(strs, x -> nvl(x, '') != '')` (string array) | 616 ms | 547 ms | 307 ms | | `array_remove(array_compact(strs), '')` (same result) | 587 ms | 252 ms | 252 ms | -| `filter(ints, x -> nvl(x, '') != '')` (integer array) | 508 ms | 456 ms | 331 ms | | `filter(ints, x -> x IS NOT NULL)` | 286 ms | 124 ms | 125 ms | | `filter(ints, x -> x > threshold)`, `threshold` is a column | 346 ms | 335 ms | 110 ms | diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index 88ca3089c6e..1790bd8eea6 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -23,11 +23,12 @@ import scala.annotation.tailrec import scala.jdk.CollectionConverters._ import org.apache.spark.SPARK_VERSION -import org.apache.spark.sql.catalyst.expressions.{Add, And, ArrayAggregate, ArrayAppend, ArrayContains, ArrayDistinct, ArrayExcept, ArrayExists, ArrayFilter, ArrayForAll, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayPosition, ArrayRemove, ArraySort, ArraysOverlap, ArraysZip, ArrayTransform, ArrayUnion, Attribute, AttributeReference, BinaryComparison, BoundReference, CaseWhen, Cast, Coalesce, Contains, CreateArray, Divide, ElementAt, EmptyRow, EndsWith, EvalMode, Expression, Flatten, GetArrayItem, GreaterThan, HigherOrderFunction, If, In, InSet, IntegralDivide, IsNaN, IsNotNull, IsNull, LambdaFunction, Length, LessThan, Literal, Lower, Multiply, NamedLambdaVariable, Not, Or, Remainder, Reverse, Sequence, Size, Slice, SortArray, StartsWith, Substring, Subtract, Upper, ZipWith} +import org.apache.spark.sql.catalyst.expressions.{Add, And, ArrayAggregate, ArrayAppend, ArrayContains, ArrayDistinct, ArrayExcept, ArrayExists, ArrayFilter, ArrayForAll, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayPosition, ArrayRemove, ArraySort, ArraysOverlap, ArraysZip, ArrayTransform, ArrayUnion, Attribute, AttributeReference, BinaryComparison, BoundReference, CaseWhen, Cast, Coalesce, Contains, CreateArray, Divide, ElementAt, EmptyRow, EndsWith, EqualTo, EvalMode, Expression, Flatten, GetArrayItem, GreaterThan, HigherOrderFunction, If, In, InSet, IntegralDivide, IsNaN, IsNotNull, IsNull, LambdaFunction, Length, LessThan, Literal, Lower, Multiply, NamedLambdaVariable, Not, Or, Remainder, Reverse, Sequence, Size, Slice, SortArray, StartsWith, Substring, Subtract, Upper, ZipWith} import org.apache.spark.sql.catalyst.util.GenericArrayData import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ +import org.apache.spark.unsafe.types.UTF8String import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.withFallbackReason @@ -1046,6 +1047,11 @@ object CometArrayFilter // 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 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, body) => // The body is made of expressions Comet evaluates natively, so run the lambda over the @@ -1056,6 +1062,34 @@ object CometArrayFilter } } + /** + * Whether `body` is `x IS NOT NULL` written the way a rule on a numeric array writes it, + * `nvl(x, '') != ''`. Spark plans that on an int, double, decimal, date or timestamp element as + * `NOT (coalesce(CAST(x AS STRING), '') = '')`, in either operand order. Such a cast is empty + * for no value, so the predicate holds exactly when `x` is not NULL. Evaluating it as written + * casts every element to a string. A string element has no cast and is left to the lambda. + */ + private def isNotNullInDisguise(body: Expression, lambdaVar: NamedLambdaVariable): Boolean = { + def isEmptyString(e: Expression): Boolean = e match { + case Literal(s: UTF8String, _: StringType) => s.numBytes() == 0 + case _ => false + } + def isCoalescedCast(e: Expression): Boolean = e match { + case Coalesce(Seq(c: Cast, empty)) if isEmptyString(empty) => + c.child match { + case v: NamedLambdaVariable => v.exprId == lambdaVar.exprId + case _ => false + } + case _ => false + } + isFixedWidth(lambdaVar.dataType) && (body match { + case Not(EqualTo(l, r)) => + (isCoalescedCast(l) && isEmptyString(r)) || + (isEmptyString(l) && isCoalescedCast(r)) + case _ => false + }) + } + /** * 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 diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_filter_numeric_nvl.sql b/spark/src/test/resources/sql-tests/expressions/array/array_filter_numeric_nvl.sql new file mode 100644 index 00000000000..b02a39d44f2 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_filter_numeric_nvl.sql @@ -0,0 +1,121 @@ +-- 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(arr, x -> nvl(x, '') != '')` on a numeric, boolean, date or timestamp array is the +-- customer's "drop the empty elements" rule. Spark plans it as +-- `NOT (coalesce(CAST(x AS STRING), '') = '')`, which no cast result can make false except a +-- NULL element, so Comet runs it as `array_compact` instead of casting every element to a string. +-- Each query must return what Spark returns, including for NULL arrays, empty arrays and +-- NULL elements. The string and non-matching shapes at the end must keep their lambda semantics. +-- There is no boolean query: Spark refuses `coalesce(boolean_value, '')`, so `nvl(x, '')` cannot be +-- written on a boolean array. + +-- Config: spark.sql.ansi.enabled=false + +statement +CREATE TABLE test_filter_numeric_nvl( + id int, + strs array, + ints array, + longs array, + dbls array, + flts array, + decs array, + dates array, + tss array, + bools array, + tinyints array, + smallints array +) USING parquet + +statement +INSERT INTO test_filter_numeric_nvl VALUES + (1, array('a', NULL, '', 'b'), array(1, NULL, 0, 5, -3), array(1L, NULL, 0L, 9000000000L), + array(1.5, NULL, CAST('NaN' AS DOUBLE), -0.0, 0.0, CAST('Infinity' AS DOUBLE)), + array(CAST(1.5 AS FLOAT), NULL, CAST('NaN' AS FLOAT), CAST(2.5 AS FLOAT)), + array(CAST(1.25 AS DECIMAL(10, 2)), NULL, CAST(-3.50 AS DECIMAL(10, 2)), CAST(0 AS DECIMAL(10, 2))), + array(DATE '2020-01-01', NULL, DATE '2024-02-29'), + array(TIMESTAMP '2020-01-01 00:00:00', NULL, TIMESTAMP '2024-02-29 12:34:56'), + array(true, NULL, false), + array(CAST(1 AS TINYINT), NULL, CAST(-128 AS TINYINT)), + array(CAST(1 AS SMALLINT), NULL, CAST(32767 AS SMALLINT))), + (2, array(), array(), array(), array(), array(), array(), array(), array(), array(), array(), array()), + (3, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL), + (4, array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), + array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL)), + (5, array('x', 'y'), array(7, 8, 9), array(7L, 8L), array(7.0, 8.0), array(CAST(7 AS FLOAT)), + array(CAST(7 AS DECIMAL(10, 2))), array(DATE '2000-01-01'), array(TIMESTAMP '2000-01-01 00:00:00'), + array(true), array(CAST(7 AS TINYINT)), array(CAST(7 AS SMALLINT))) + +query expect_native(filter) +SELECT id, filter(ints, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(ints, x -> nvl(x, '') <> '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(ints, x -> '' != nvl(x, '')) FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(ints, x -> coalesce(x, '') != '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(longs, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(dbls, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(flts, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(decs, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(dates, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(tss, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(tinyints, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT id, filter(smallints, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +-- The result is a usable array: nested in other functions, and in the shape of the rule. +query expect_native(filter) +SELECT id, size(filter(ints, x -> nvl(x, '') != '')), array_max(filter(longs, x -> nvl(x, '') != '')) +FROM test_filter_numeric_nvl + +query expect_native(filter) +SELECT sum(size(filter(ints, x -> nvl(x, '') != ''))) FROM test_filter_numeric_nvl + +-- A string element has no cast: empty strings are dropped as well as NULLs. +query expect_native(filter) +SELECT id, filter(strs, x -> nvl(x, '') != '') FROM test_filter_numeric_nvl + +-- Not the same predicate, so these keep their lambda semantics: another default value, the equal +-- operator, and a body that tests more than the element. +query +SELECT id, filter(ints, x -> nvl(x, '0') != '0') FROM test_filter_numeric_nvl + +query +SELECT id, filter(ints, x -> nvl(x, '') = '') FROM test_filter_numeric_nvl + +query +SELECT id, filter(ints, x -> nvl(x, '') != '' AND x > 0) FROM test_filter_numeric_nvl