From a30b862fdc636d8c5166da849681c4636379c2ef Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 29 Oct 2026 14:12:07 -0600 Subject: [PATCH 1/8] 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/8] 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/8] 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/8] 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 From 165695918757d4a5a9e6e4df87867bb92d731458 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 5 Nov 2026 16:55:25 -0700 Subject: [PATCH 5/8] feat: run exists lambdas natively with DataFusion's array_any_match CometArrayExists now serializes a lambda whose body Comet supports natively as a HigherOrderFunc under array_any_match, as filter does with array_filter, when three-valued logic is on. The gate and the serialization move into a trait both serdes share. exists stays on the dispatcher with spark.sql.legacy.followThreeValuedLogicInArrayExists=false, since Spark returns false where array_any_match returns NULL, and when the array argument contains a cast to a map, whose TRY cast can produce a null key that Arrow cannot hold (#6172). Also fixes a higher-order function inside an expression that is evaluated only for some rows: as the right operand of AND or OR, or inside an argument of a function like array_contains. These failed with 'project index N out of bounds', because project() in null_short_circuit.rs collected the function's own lambda parameter as a column of the outer batch. It now ignores variables bound by a lambda inside the expression it projects. This affected filter as well as exists. --- .../spark_configs_support.md | 12 +- .../user-guide/latest/tuning/expressions.md | 15 +- native/core/src/execution/planner.rs | 120 +++++++++- .../conditional_funcs/null_short_circuit.rs | 22 +- .../scala/org/apache/comet/CometConf.scala | 6 +- .../scala/org/apache/comet/serde/arrays.scala | 215 +++++++++++------- .../array/array_exists_native_lambda.sql | 115 ++++++++++ .../array_exists_native_lambda_two_valued.sql | 31 +++ .../array/array_filter_native_lambda.sql | 13 +- .../array/array_filter_numeric_nvl.sql | 4 +- .../sql-tests/expressions/array/exists.sql | 3 +- .../apache/comet/CometCodegenHOFSuite.scala | 23 +- 12 files changed, 468 insertions(+), 111 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda_two_valued.sql diff --git a/docs/source/contributor-guide/spark_configs_support.md b/docs/source/contributor-guide/spark_configs_support.md index 92107975baf..7f667bbf714 100644 --- a/docs/source/contributor-guide/spark_configs_support.md +++ b/docs/source/contributor-guide/spark_configs_support.md @@ -62,11 +62,13 @@ least one result is `null`, `exists` returns `null` when the config is `true` an is `false`. A matching element still returns `true`, while a null input array still returns `null`, under either value. -**Comet status.** `CometArrayExists` routes Spark's `ArrayExists` expression through the JVM -codegen dispatcher. The dispatcher evaluates the Spark expression carrying the captured -`followThreeValuedLogic` value inside the Comet pipeline, so both config values preserve Spark's -semantics. If the dispatcher is disabled, the affected Comet operator falls back to Spark rather -than executing with different semantics. +**Comet status.** With the setting `true` (the default), `CometArrayExists` runs a lambda whose body Comet +supports natively under DataFusion's `array_any_match`, which has the same three-valued result: NULL when no +element matches but the predicate was NULL for some. With the setting `false`, and for every other lambda, +`CometArrayExists` routes Spark's `ArrayExists` expression through the JVM codegen dispatcher. The dispatcher +evaluates the Spark expression carrying the captured `followThreeValuedLogic` value inside the Comet pipeline, so +both config values preserve Spark's semantics. If the dispatcher is disabled, the affected Comet operator falls +back to Spark rather than executing with different semantics. **Test coverage.** `spark/src/test/resources/sql-tests/expressions/array/exists.sql` uses a `ConfigMatrix` to run both values and covers null arrays, empty arrays, matching predicates, diff --git a/docs/source/user-guide/latest/tuning/expressions.md b/docs/source/user-guide/latest/tuning/expressions.md index 1ee62b48d69..43f16225617 100644 --- a/docs/source/user-guide/latest/tuning/expressions.md +++ b/docs/source/user-guide/latest/tuning/expressions.md @@ -24,11 +24,16 @@ under the License. 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), 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. | +| Lambda | How Comet runs it | +| -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `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` and `exists` 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. | + +`exists` runs natively under DataFusion's `array_any_match`, which returns what Spark's three-valued logic returns: NULL +when no element matches but the lambda was NULL for some element. With +`spark.sql.legacy.followThreeValuedLogicInArrayExists=false`, Spark returns false there, and `exists` stays on the +dispatcher. 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. diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 31bf044f83f..f5acddd63ea 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -96,6 +96,7 @@ use crate::parquet::parquet_support::{ use datafusion::common::ScalarValue; use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::object_store::ObjectStoreUrl; +use datafusion::functions_nested::array_any_match::array_any_match_higher_order_function; 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}; @@ -1176,6 +1177,7 @@ impl PhysicalPlanner { 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(), + "array_any_match" => array_any_match_higher_order_function(), other => { return Err(GeneralError(format!( "Unsupported higher-order function '{other}'" @@ -5344,8 +5346,96 @@ mod tests { } } + /// A higher-order function as the right operand of OR or AND: DataFusion evaluates the right + /// operand over only the rows the left operand leaves open. + #[test] + fn test_higher_order_as_right_operand() -> Result<(), DataFusionError> { + use arrow::array::{BooleanArray, ListBuilder, 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 binary(left: spark_expression::Expr, right: spark_expression::Expr) -> Box { + Box::new(BinaryExpr { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + }) + } + // array_any_match(a, x#1 -> x#1 = c) + let any_match = expr(ExprStruct::HigherOrderFunc(Box::new(HigherOrderFunc { + func: "array_any_match".to_string(), + array: Some(Box::new(bound(0))), + param_name: "x#1".to_string(), + body: Some(Box::new(expr(ExprStruct::Eq(binary( + expr(ExprStruct::LambdaVariable(LambdaVariable { + name: "x#1".to_string(), + })), + bound(1), + ))))), + }))); + + let mut builder = ListBuilder::new(StringBuilder::new()); + for row in [vec!["a"], vec!["b"], vec!["a"], vec!["z"]] { + row.into_iter() + .for_each(|v| builder.values().append_value(v)); + builder.append(true); + } + let list = builder.finish(); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", list.data_type().clone(), true), + Field::new("c", DataType::Utf8, true), + Field::new("b", DataType::Boolean, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(list), + Arc::new(StringArray::from(vec!["a", "a", "x", "z"])), + Arc::new(BooleanArray::from(vec![true, false, false, false])), + ], + )?; + let session_ctx = SessionContext::new(); + let planner = PhysicalPlanner::new(Arc::from(session_ctx), 0); + + // b OR any_match(...): rows 0 and 3 are true, the others have no match. + let or = expr(ExprStruct::Or(binary(bound(2), any_match.clone()))); + let physical = planner.create_expr(&or, Arc::clone(&schema)).unwrap(); + let result = physical.evaluate(&batch)?.into_array(batch.num_rows())?; + let expected = BooleanArray::from(vec![true, false, false, true]); + assert_eq!(&result, &(Arc::new(expected) as ArrayRef)); + + // NOT b AND any_match(...): rows 1 and 2 are false, row 3 matches. + let and = expr(ExprStruct::And(binary( + expr(ExprStruct::Not(Box::new( + datafusion_comet_proto::spark_expression::UnaryExpr { + child: Some(Box::new(bound(2))), + }, + ))), + any_match, + ))); + let physical = planner.create_expr(&and, Arc::clone(&schema)).unwrap(); + let result = physical.evaluate(&batch)?.into_array(batch.num_rows())?; + let expected = BooleanArray::from(vec![false, false, false, true]); + assert_eq!(&result, &(Arc::new(expected) as ArrayRef)); + Ok(()) + } + #[test] fn test_higher_order_filter() -> Result<(), DataFusionError> { + use arrow::array::BooleanArray; use arrow::array::ListBuilder; use arrow::array::StringBuilder; use datafusion_comet_proto::spark_expression::{ @@ -5433,7 +5523,35 @@ mod tests { ]); assert_eq!(&result, &(Arc::new(expected) as ArrayRef)); - // The argument must be an array, and only array_filter is implemented. + // array_any_match(a, x#1 -> x#1 = c) is Spark's `exists`: true when an element matches, NULL + // when none does but the predicate was NULL for some element, false otherwise, and NULL for a + // NULL array. + let list = string_lists(vec![ + Some(vec![Some("a"), None, Some("b")]), + Some(vec![]), + None, + Some(vec![Some("b")]), + Some(vec![Some("b"), None]), + ]); + 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", "a"])), + ], + )?; + let physical = planner + .create_expr(&filter("array_any_match", bound(0)), Arc::clone(&schema)) + .unwrap(); + let result = physical.evaluate(&batch)?.into_array(batch.num_rows())?; + let expected = BooleanArray::from(vec![Some(true), Some(false), None, Some(false), None]); + assert_eq!(&result, &(Arc::new(expected) as ArrayRef)); + + // The argument must be an array, and only array_filter and array_any_match are implemented. let not_array = planner.create_expr(&filter("array_filter", bound(1)), Arc::clone(&schema)); assert!(not_array .unwrap_err() 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 4c540633803..64e8435a9f4 100644 --- a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs +++ b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs @@ -23,8 +23,9 @@ use arrow::record_batch::{RecordBatch, RecordBatchOptions}; 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, LambdaVariable, Literal}; +use datafusion::physical_expr::expressions::{Column, LambdaExpr, LambdaVariable, Literal}; use datafusion::physical_expr::PhysicalExpr; +use std::collections::HashSet; use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; use std::sync::Arc; @@ -293,12 +294,24 @@ pub(super) fn evaluate_projected( /// 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. pub(super) fn project(expr: &Arc) -> Result<(Vec, Arc)> { + // A lambda inside `expr` binds its parameters: they are not columns of the batch `expr` is + // evaluated on, so they are neither projected nor re-indexed. Lambda parameter names are unique + // per lambda, so no name is both bound inside `expr` and free in it. + let mut bound: HashSet = HashSet::new(); + expr.apply(|e| { + if let Some(lambda) = e.downcast_ref::() { + bound.extend(lambda.params().iter().cloned()); + } + Ok(TreeNodeRecursion::Continue) + })?; 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()); + if !bound.contains(variable.name()) { + projection.push(variable.index()); + } } Ok(TreeNodeRecursion::Continue) })?; @@ -318,7 +331,10 @@ pub(super) fn project(expr: &Arc) -> Result<(Vec, Arc )); } - } else if let Some(variable) = e.downcast_ref::() { + } else if let Some(variable) = e + .downcast_ref::() + .filter(|v| !bound.contains(v.name())) + { let new = position(variable.index())?; if new != variable.index() { return Ok(Transformed::yes(Arc::new(LambdaVariable::new( diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 57440dd9443..d09d94cd165 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -806,9 +806,9 @@ object CometConf extends ShimCometConf { "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.") + "Applies to `filter` and `exists`. 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/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index 1790bd8eea6..f52cffbab30 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -1029,93 +1029,46 @@ object CometFlatten extends CometExpressionSerde[Flatten] with ArraysBase { } } -object CometArrayFilter - extends CometExpressionSerde[ArrayFilter] - with CometTypeShim - with CometExprShim { - - override def getSupportLevel(expr: ArrayFilter): SupportLevel = Compatible() - - override def convert( - expr: ArrayFilter, - inputs: Seq[Attribute], - binding: Boolean): Option[ExprOuterClass.Expr] = { - expr.function match { - case LambdaFunction(IsNotNull(v: NamedLambdaVariable), Seq(lambdaVar), _) - if v.exprId == lambdaVar.exprId => - // Fast path: Catalyst desugars `array_compact` to `filter(arr, x -> x IS NOT NULL)`, so - // 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 - // 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) - } - } - - /** - * 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 - * 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) +/** + * What a serde needs to run the lambda of a higher-order function natively, under a DataFusion + * higher-order function: the gate on the lambda and its body, and the serialization. + */ +trait NativeLambdaSerde extends CometTypeShim with CometExprShim { - // The native lambda has the same semantics as Spark's `filter` for a one-parameter lambda - // whose body is deterministic, contains no nested lambda, and whose parameter is a type the - // native side can type-check against the list's element field. 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 = + // The native lambda has the same semantics as Spark's `filter` and `exists` 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. + protected def nativeLambdaSupported(argument: Expression, body: Expression): Boolean = CometConf.COMET_NATIVE_LAMBDA_ENABLED.get() && body.deterministic && !body.exists(_.isInstanceOf[HigherOrderFunction]) && capturesOnlyFixedWidthColumns(body) && shortCircuitSafe(body) && - (expr.argument.dataType match { + !castsToMap(argument) && + (argument.dataType match { case ArrayType(elementType, _) => isLambdaElementType(elementType) case _ => false }) + /** + * Whether `argument` casts to a type that holds a map. A TRY cast of a map whose key cast can + * fail makes a null key, which Arrow cannot hold (#6172), so such an argument stays on the + * dispatcher, which falls back to Spark for it. + */ + private def castsToMap(argument: Expression): Boolean = { + def holdsMap(dt: DataType): Boolean = dt match { + case _: MapType => true + case ArrayType(element, _) => holdsMap(element) + case StructType(fields) => fields.exists(f => holdsMap(f.dataType)) + case _ => false + } + argument.exists { + case c: Cast => holdsMap(c.dataType) + 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 @@ -1164,7 +1117,7 @@ object CometArrayFilter case _ => true } - private def isFixedWidth(dt: DataType): Boolean = dt match { + protected def isFixedWidth(dt: DataType): Boolean = dt match { case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType | DateType | TimestampType | TimestampNTZType | _: DecimalType => true @@ -1177,21 +1130,23 @@ object CometArrayFilter case _ => isFixedWidth(dt) } - private def convertLambda( - expr: ArrayFilter, + /** `func` is the DataFusion higher-order function the lambda runs under. */ + protected def convertLambda( + func: String, + argument: Expression, lambdaVar: NamedLambdaVariable, body: Expression, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { for { - arrayProto <- exprToProtoInternal(expr.argument, inputs, binding) + arrayProto <- exprToProtoInternal(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) } yield { val hof = ExprOuterClass.HigherOrderFunc .newBuilder() - .setFunc("array_filter") + .setFunc(func) .setArray(arrayProto) .setParamName(CometNamedLambdaVariable.uniqueName(lambdaVar)) .setBody(bodyProto) @@ -1200,6 +1155,76 @@ object CometArrayFilter } } +object CometArrayFilter extends CometExpressionSerde[ArrayFilter] with NativeLambdaSerde { + + override def getSupportLevel(expr: ArrayFilter): SupportLevel = Compatible() + + override def convert( + expr: ArrayFilter, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = { + expr.function match { + case LambdaFunction(IsNotNull(v: NamedLambdaVariable), Seq(lambdaVar), _) + if v.exprId == lambdaVar.exprId => + // Fast path: Catalyst desugars `array_compact` to `filter(arr, x -> x IS NOT NULL)`, so + // 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.argument, 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("array_filter", expr.argument, lambdaVar, body, inputs, binding).orElse( + dispatch(expr, inputs, binding)) + case _ => dispatch(expr, inputs, binding) + } + } + + /** + * 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 + * 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) +} + /** * 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. @@ -1420,7 +1445,31 @@ trait ArraysBase { object CometArrayTransform extends CometCodegenDispatch[ArrayTransform] -object CometArrayExists extends CometCodegenDispatch[ArrayExists] +/** + * `exists(arr, x -> body)`. When the body is made of natively supported expressions it runs under + * DataFusion's `array_any_match`, which returns what Spark returns with three-valued logic: true + * for a match, NULL when none matches but the predicate was NULL for some element, false + * otherwise. Spark returns false instead of NULL in that last case when + * `spark.sql.legacy.followThreeValuedLogicInArrayExists` is false, so `exists` stays on the + * dispatcher with that setting. + */ +object CometArrayExists extends CometCodegenDispatch[ArrayExists] with NativeLambdaSerde { + + override def convert( + expr: ArrayExists, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = { + expr.function match { + case LambdaFunction(body, Seq(lambdaVar: NamedLambdaVariable), _) + if binding && expr.followThreeValuedLogic && nativeLambdaSupported( + expr.argument, + body) => + convertLambda("array_any_match", expr.argument, lambdaVar, body, inputs, binding) + .orElse(super.convert(expr, inputs, binding)) + case _ => super.convert(expr, inputs, binding) + } + } +} object CometArrayForAll extends CometCodegenDispatch[ArrayForAll] diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda.sql b/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda.sql new file mode 100644 index 00000000000..6e6fb150124 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda.sql @@ -0,0 +1,115 @@ +-- 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. + +-- `exists` with a lambda whose body Comet evaluates natively runs through DataFusion's +-- `array_any_match`. Each query below must return what Spark returns, including for NULL arrays, +-- empty arrays, NULL elements, and predicates that are NULL for some elements and false for the +-- rest, where Spark's three-valued logic returns NULL. + +-- Config: spark.sql.ansi.enabled=false +-- Config: spark.sql.legacy.followThreeValuedLogicInArrayExists=true + +statement +CREATE TABLE test_exists_lambda( + id int, + strs array, + ints array, + dbls array, + flts array, + nested array>, + threshold int, + code string, + flag boolean +) USING parquet + +statement +INSERT INTO test_exists_lambda VALUES + (1, array('a', NULL, '', 'b'), array(1, NULL, 0, 5, -3), array(1.5, NULL, CAST('NaN' AS DOUBLE), -0.0, 0.0, 95.0), + array(CAST(1.5 AS FLOAT), -2.5, CAST(2.5 AS FLOAT)), array(array('x', ''), array('y')), 3, 'US', true), + (2, array(), array(), array(), array(), array(), 0, 'CN', false), + (3, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL), + (4, array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(NULL, NULL), array(array(NULL), NULL), 1, '', true), + (5, array('x', 'y'), array(7, 8, 9), array(7.0, 8.0), array(CAST(7 AS FLOAT)), array(array('x', 'y')), 8, 'IN', true), + (6, array('b'), array(2, 4), array(40.5, 41.0), array(CAST(1 AS FLOAT)), array(array('b')), 100, 'US', false) + +-- Range checks over an array of coordinates. +query expect_native(exists) +SELECT id, exists(dbls, x -> isnull(x) OR x < -90.0 OR x > 90.0 OR x = 0.0) FROM test_exists_lambda + +query expect_native(exists) +SELECT id, size(dbls) IN (0, -1) OR exists(dbls, x -> isnull(x) OR x = 0.0) FROM test_exists_lambda + +query expect_native(exists) +SELECT id, exists(flts, x -> x < 0.0 OR isnull(x)) FROM test_exists_lambda + +-- Strings, with a NULL or empty-string element. +query expect_native(exists) +SELECT id, exists(strs, x -> isnull(x) OR x = '') FROM test_exists_lambda + +query expect_native(exists) +SELECT id, exists(strs, x -> x IN ('a', 'b', 'c')) FROM test_exists_lambda + +query expect_native(exists) +SELECT id, exists(strs, x -> x = 'a') FROM test_exists_lambda + +-- Integers, and the NULL result when no element matches but the predicate was NULL for some. +query expect_native(exists) +SELECT id, exists(ints, x -> x > 4) FROM test_exists_lambda + +query expect_native(exists) +SELECT id, exists(ints, x -> x > 100) FROM test_exists_lambda + +query expect_native(exists) +SELECT id, exists(ints, x -> CAST(NULL AS BOOLEAN)) FROM test_exists_lambda + +query expect_native(exists) +SELECT id, exists(ints, x -> true), exists(ints, x -> false) FROM test_exists_lambda + +-- A captured fixed-width column. +query expect_native(exists) +SELECT id, exists(ints, x -> x > threshold) FROM test_exists_lambda + +-- The argument is a native expression of its own. +query expect_native(exists) +SELECT id, exists(flatten(nested), x -> isnull(x) OR x = '') FROM test_exists_lambda + +-- In a filter, in an aggregate, and combined with other conditions. +query expect_native(exists) +SELECT id FROM test_exists_lambda WHERE exists(ints, x -> x > 4) ORDER BY id + +-- Inside an aggregate the harness counts the expression as dispatched as well as native, so these only +-- compare with Spark. +query +SELECT count(CASE WHEN exists(dbls, x -> isnull(x) OR x = 0.0) AND flag THEN 1 END) +FROM test_exists_lambda + +query +SELECT count(CASE WHEN exists(strs, x -> x = '') AND code = 'US' THEN 1 END) +FROM test_exists_lambda + +-- A body that converts each element to a date, all of it natively supported expressions. +query expect_native(exists) +SELECT id, exists(dbls, x -> datediff(DATE '2024-03-15', CAST(timestamp_seconds(CAST(x AS BIGINT)) AS DATE)) >= 2) +FROM test_exists_lambda + +-- exists next to another lambda, as the right operand of OR, and inside an argument that is evaluated only +-- for some rows. +query expect_native(exists) +SELECT id, size(ints) < 1 OR exists(ints, x -> x > 4) FROM test_exists_lambda + +query expect_native(exists) +SELECT id, array_contains(ints, size(filter(ints, x -> x > 1))) OR exists(ints, x -> x > 100) FROM test_exists_lambda diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda_two_valued.sql b/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda_two_valued.sql new file mode 100644 index 00000000000..973db0f0919 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda_two_valued.sql @@ -0,0 +1,31 @@ +-- 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 `spark.sql.legacy.followThreeValuedLogicInArrayExists=false` Spark returns false where it +-- would return NULL because the predicate was NULL for some element and false for the rest. +-- DataFusion's `array_any_match` returns NULL there, so `exists` stays on the dispatcher. + +-- Config: spark.sql.legacy.followThreeValuedLogicInArrayExists=false + +statement +CREATE TABLE test_exists_two_valued(id int, ints array) USING parquet + +statement +INSERT INTO test_exists_two_valued VALUES (1, array(1, NULL, 3)), (2, array()), (3, NULL), (4, array(NULL, NULL)) + +query expect_dispatch(exists) +SELECT id, exists(ints, x -> x > 2), exists(ints, x -> x > 10) FROM test_exists_two_valued 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 bfe07fb050e..01296201079 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 @@ -65,7 +65,7 @@ INSERT INTO test_filter_lambda VALUES 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'), array(CAST(7 AS TINYINT)), array(CAST(7 AS SMALLINT)), 8, 7, 'y') --- The customer shape: drop NULL and empty-string elements. +-- A common cleaning shape: drop NULL and empty-string elements. query expect_native(filter) SELECT id, filter(strs, x -> nvl(x, '') != '') FROM test_filter_lambda @@ -198,3 +198,14 @@ SELECT id, filter(ints, v -> v > 0) FROM test_filter_lambda WHERE id < 0 -- 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 + +-- A lambda inside an argument that is evaluated only for the rows where an earlier argument or operand +-- leaves the result open. The lambda parameter belongs to the inner lambda, not to the outer batch. +query expect_native(filter) +SELECT id, array_contains(ints, size(filter(ints, x -> x > 1))) FROM test_filter_lambda + +query expect_native(filter) +SELECT id, size(ints) > 2 AND size(filter(ints, x -> x > 1)) > 1 FROM test_filter_lambda + +query expect_native(filter) +SELECT id, size(ints) < 1 OR size(filter(strs, x -> x IS NOT NULL AND x != '')) > 1 FROM test_filter_lambda 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 index b02a39d44f2..7412d469175 100644 --- 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 @@ -15,8 +15,8 @@ -- 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 +-- `filter(arr, x -> nvl(x, '') != '')` on a numeric, boolean, date or timestamp array is a common +-- "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 diff --git a/spark/src/test/resources/sql-tests/expressions/array/exists.sql b/spark/src/test/resources/sql-tests/expressions/array/exists.sql index 7e8aa389952..4ec25440615 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/exists.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/exists.sql @@ -15,7 +15,8 @@ -- specific language governing permissions and limitations -- under the License. --- Higher-order function exists. Runs through Comet's codegen dispatcher. +-- Higher-order function exists. Runs through Comet's codegen dispatcher, or natively when the lambda +-- allows it and three-valued logic is on (see array_exists_native_lambda.sql). -- ConfigMatrix: spark.sql.legacy.followThreeValuedLogicInArrayExists=false,true statement diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala index 49fadd16e0e..98ee299da38 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. * @@ -98,15 +100,22 @@ class CometCodegenHOFSuite test("ArrayExists respects three-valued logic through codegen dispatcher") { withArrayIntTable("(array(1, null, 3)), (array(1, 2, 3)), (array()), (null)") { - Seq(false, true).foreach { followThreeValuedLogic => - withSQLConf( - "spark.sql.legacy.followThreeValuedLogicInArrayExists" -> - followThreeValuedLogic.toString) { - assertCodegenRan { - checkSparkAnswerAndOperator(sql("SELECT exists(a, x -> x > 10) FROM t")) - } + // With the legacy setting off Spark returns false where DataFusion's array_any_match returns + // NULL, so only the dispatcher runs `exists`. With it on the lambda runs natively and the + // dispatcher stays idle, with the same answer as Spark. + withSQLConf("spark.sql.legacy.followThreeValuedLogicInArrayExists" -> "false") { + assertCodegenRan { + checkSparkAnswerAndOperator(sql("SELECT exists(a, x -> x > 10) FROM t")) } } + withSQLConf("spark.sql.legacy.followThreeValuedLogicInArrayExists" -> "true") { + CometScalaUDFCodegen.resetStats() + checkSparkAnswerAndOperator(sql("SELECT exists(a, x -> x > 10) FROM t")) + val after = CometScalaUDFCodegen.stats() + assert( + after.compileCount + after.cacheHitCount == 0, + s"expected exists to run natively, but the dispatcher ran: $after") + } } } From 66a1066c2cf1d0cf48af87fa8d772aabef3a8bd4 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 5 Nov 2026 17:26:21 -0700 Subject: [PATCH 6/8] bench: add exists cases to CometHigherOrderFunctionBenchmark and the tuning guide --- .../user-guide/latest/tuning/expressions.md | 13 ++++++- .../CometHigherOrderFunctionBenchmark.scala | 34 +++++++++++++++++-- 2 files changed, 44 insertions(+), 3 deletions(-) diff --git a/docs/source/user-guide/latest/tuning/expressions.md b/docs/source/user-guide/latest/tuning/expressions.md index 43f16225617..8112dfd8ab5 100644 --- a/docs/source/user-guide/latest/tuning/expressions.md +++ b/docs/source/user-guide/latest/tuning/expressions.md @@ -71,7 +71,7 @@ To get the most out of Comet when a query uses lambdas: ### Example -`CometHigherOrderFunctionBenchmark` filters 1 million rows of ten-element arrays, in which every fifth element is +`CometHigherOrderFunctionBenchmark` runs `filter` over 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: @@ -82,4 +82,15 @@ To get the most out of Comet when a query uses lambdas: | `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 | +The same benchmark runs `exists` over 1 million rows of ten-element arrays, in which every fifth element is `NULL` +(Spark 3.5, same laptop): + +| Query | Spark | Comet, nativeLambda off | Comet, default | +| ----------------------------------------------------------- | ------ | ----------------------- | -------------- | +| `exists(ints, x -> x > 90)` | 266 ms | 234 ms | 108 ms | +| `exists(dbls, x -> x IS NULL OR x < 0 OR x > 90 OR x = 0)` | 266 ms | 225 ms | 94 ms | +| `exists(ints, x -> x > threshold)`, `threshold` is a column | 251 ms | 196 ms | 105 ms | +| `size(dbls) < 1 OR exists(dbls, x -> x > 90)` | 355 ms | 403 ms | 110 ms | +| `exists(strs, x -> x IS NULL OR x = '')` | 357 ms | 183 ms | 168 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/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHigherOrderFunctionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHigherOrderFunctionBenchmark.scala index 741b624cc8f..a44b1dbcb3f 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 @@ -23,8 +23,8 @@ 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: + * Benchmark to measure `filter` and `exists` with a lambda, and the `array_compact` / + * `array_remove` forms a `filter` lambda can often be rewritten to. To run this benchmark: * {{{ * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometHigherOrderFunctionBenchmark * }}} @@ -48,11 +48,15 @@ object CometHigherOrderFunctionBenchmark extends CometBenchmarkBase { val intElems = (0 until 10).map { i => s"CASE WHEN (value + $i) % 5 = 0 THEN NULL ELSE cast((value + $i) % 100 as int) END" } + val dblElems = (0 until 10).map { i => + s"CASE WHEN (value + $i) % 5 = 0 THEN NULL ELSE cast((value + $i) % 97 as double) END" + } prepareTable( dir, spark.sql(s"""SELECT | array(${strElems.mkString(",\n")}) AS str_arr, | array(${intElems.mkString(",\n")}) AS int_arr, + | array(${dblElems.mkString(",\n")}) AS dbl_arr, | cast(value % 3 as int) AS threshold |FROM $tbl""".stripMargin)) f @@ -89,7 +93,33 @@ object CometHigherOrderFunctionBenchmark extends CometBenchmarkBase { } } + def existsLambdaBenchmark(values: Int): Unit = { + prepareFilterTable { + Seq( + "string: exists(x -> x IS NULL OR x = '')" -> + "SELECT exists(str_arr, x -> x IS NULL OR x = '') FROM parquetV1Table", + "int: exists(x -> x > 90)" -> + "SELECT exists(int_arr, x -> x > 90) FROM parquetV1Table", + "double: exists(x -> x IS NULL OR x < 0 OR x > 90 OR x = 0)" -> + "SELECT exists(dbl_arr, x -> x IS NULL OR x < 0.0 OR x > 90.0 OR x = 0.0) FROM parquetV1Table", + "int: exists(x -> x > threshold)" -> + "SELECT exists(int_arr, x -> x > threshold) FROM parquetV1Table", + // exists as the right operand of OR, so it runs only for the rows the left operand leaves open. + "double: size(arr) < 1 OR exists(x -> x > 90)" -> + "SELECT size(dbl_arr) < 1 OR exists(dbl_arr, x -> x > 90.0) FROM parquetV1Table") + .foreach { case (name, query) => + runExpressionBenchmark(name, values.toLong, query) + 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)(filterLambdaBenchmark) + runBenchmarkWithTable("ExistsLambda", 1024 * 1024)(existsLambdaBenchmark) } } From 3bb85cfc1b6e1e53d5433e653d40db1249c95dfe Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 5 Nov 2026 20:37:50 -0700 Subject: [PATCH 7/8] fix: keep an exists lambda that can raise on the dispatcher Spark's exists stops at the first element that matches, so it never evaluates the lambda on the elements after it, while array_any_match evaluates the lambda on every element. With ANSI on, exists(ints, x -> 100 / x > 1) over array(1, 0) raised DIVIDE_BY_ZERO natively where Spark returns true, and with ANSI off element_at with the index 0 raised in the same way. exists now also requires that no expression in the lambda can raise, whatever operators it uses. Also from review: - Reuse CometCast.canProduceNullMapKey for the map-cast check, so a widening key cast in the array argument runs natively. - Make CometArrayExists a plain serde that calls the dispatcher itself, like CometArrayFilter, so the generated docs say Native, and give its row the same notes as filter. - Collect the projection in one walk in NullShortCircuit's project(). - Drop tests that repeat the SQL files: the exists test in CometCodegenHOFSuite, the array_any_match half of the planner test and the numeric nvl query in array_filter_native_lambda.sql. --- docs/source/user-guide/latest/expressions.md | 2 +- .../user-guide/latest/tuning/expressions.md | 12 ++- native/core/src/execution/planner.rs | 74 +++++-------------- .../conditional_funcs/null_short_circuit.rs | 11 +-- .../scala/org/apache/comet/serde/arrays.scala | 58 ++++++++------- .../array/array_exists_native_lambda.sql | 12 ++- .../array/array_exists_native_lambda_ansi.sql | 57 ++++++++++++++ .../array/array_filter_native_lambda.sql | 4 - .../misc/codegen_dispatch_null_map_key.sql | 5 ++ .../apache/comet/CometCodegenHOFSuite.scala | 23 ------ 10 files changed, 135 insertions(+), 123 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda_ansi.sql diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index e9ecb63ef7a..760199fd227 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -387,7 +387,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 | | +| `exists` | ✅ | Native | A one-parameter lambda whose body is natively supported and cannot raise runs natively; other lambdas, and `exists` with `spark.sql.legacy.followThreeValuedLogicInArrayExists=false`, use the JVM codegen dispatcher. See [Expression Tuning](tuning/expressions.md#lambdas-and-higher-order-functions) | | `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 | | diff --git a/docs/source/user-guide/latest/tuning/expressions.md b/docs/source/user-guide/latest/tuning/expressions.md index 8112dfd8ab5..401ba77d4c8 100644 --- a/docs/source/user-guide/latest/tuning/expressions.md +++ b/docs/source/user-guide/latest/tuning/expressions.md @@ -60,10 +60,14 @@ To get the most out of Comet when a query uses lambdas: 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. + `filter(arr, x -> CASE WHEN x = 0 THEN false ELSE 100 / x > 1 END)` runs natively. +- **Keep operations that can raise out of an `exists` lambda.** Spark stops at the first element that matches, so + under ANSI mode `exists(array(1, 0), x -> 100 / x > 1)` returns true without dividing by zero, but the native lambda + evaluates every element. An `exists` lambda that contains any operation that can raise stays on the dispatcher, + whether or not it uses `AND`, `OR`, `IN` or `CASE WHEN`. +- **Capture only numeric, decimal, 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/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index f5acddd63ea..ec237e025bd 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -5346,28 +5346,30 @@ mod tests { } } + 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(spark_expression::BoundReference { + index, + datatype: None, + })) + } + /// A higher-order function as the right operand of OR or AND: DataFusion evaluates the right /// operand over only the rows the left operand leaves open. #[test] fn test_higher_order_as_right_operand() -> Result<(), DataFusionError> { use arrow::array::{BooleanArray, ListBuilder, StringBuilder}; use datafusion_comet_proto::spark_expression::{ - BinaryExpr, BoundReference, HigherOrderFunc, LambdaVariable, + BinaryExpr, 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 binary(left: spark_expression::Expr, right: spark_expression::Expr) -> Box { Box::new(BinaryExpr { left: Some(Box::new(left)), @@ -5435,26 +5437,12 @@ mod tests { #[test] fn test_higher_order_filter() -> Result<(), DataFusionError> { - use arrow::array::BooleanArray; use arrow::array::ListBuilder; use arrow::array::StringBuilder; use datafusion_comet_proto::spark_expression::{ - BinaryExpr, BoundReference, HigherOrderFunc, LambdaVariable, + BinaryExpr, 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(), @@ -5523,34 +5511,6 @@ mod tests { ]); assert_eq!(&result, &(Arc::new(expected) as ArrayRef)); - // array_any_match(a, x#1 -> x#1 = c) is Spark's `exists`: true when an element matches, NULL - // when none does but the predicate was NULL for some element, false otherwise, and NULL for a - // NULL array. - let list = string_lists(vec![ - Some(vec![Some("a"), None, Some("b")]), - Some(vec![]), - None, - Some(vec![Some("b")]), - Some(vec![Some("b"), None]), - ]); - 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", "a"])), - ], - )?; - let physical = planner - .create_expr(&filter("array_any_match", bound(0)), Arc::clone(&schema)) - .unwrap(); - let result = physical.evaluate(&batch)?.into_array(batch.num_rows())?; - let expected = BooleanArray::from(vec![Some(true), Some(false), None, Some(false), None]); - assert_eq!(&result, &(Arc::new(expected) as ArrayRef)); - // The argument must be an array, and only array_filter and array_any_match are implemented. let not_array = planner.create_expr(&filter("array_filter", bound(1)), Arc::clone(&schema)); assert!(not_array 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 64e8435a9f4..28b0fc57622 100644 --- a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs +++ b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs @@ -296,17 +296,14 @@ pub(super) fn evaluate_projected( pub(super) fn project(expr: &Arc) -> Result<(Vec, Arc)> { // A lambda inside `expr` binds its parameters: they are not columns of the batch `expr` is // evaluated on, so they are neither projected nor re-indexed. Lambda parameter names are unique - // per lambda, so no name is both bound inside `expr` and free in it. + // per lambda, so no name is both bound inside `expr` and free in it. `apply` visits each lambda + // before the variables it binds, so `bound` already holds a variable's name when it is reached. let mut bound: HashSet = HashSet::new(); + let mut projection = vec![]; expr.apply(|e| { if let Some(lambda) = e.downcast_ref::() { bound.extend(lambda.params().iter().cloned()); - } - Ok(TreeNodeRecursion::Continue) - })?; - let mut projection = vec![]; - expr.apply(|e| { - if let Some(column) = e.downcast_ref::() { + } else if let Some(column) = e.downcast_ref::() { projection.push(column.index()); } else if let Some(variable) = e.downcast_ref::() { if !bound.contains(variable.name()) { 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 f52cffbab30..f38a003db66 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -33,7 +33,7 @@ import org.apache.spark.unsafe.types.UTF8String 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.expressions.{CometCast, CometEvalMode} import org.apache.comet.serde.QueryPlanSerde._ import org.apache.comet.shims.{CometExprShim, CometTypeShim} @@ -1035,39 +1035,34 @@ object CometFlatten extends CometExpressionSerde[Flatten] with ArraysBase { */ trait NativeLambdaSerde extends CometTypeShim with CometExprShim { - // The native lambda has the same semantics as Spark's `filter` and `exists` 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. + // The native lambda has the same semantics as Spark's `filter` for a one-parameter lambda whose + // body is deterministic, contains no nested lambda, and whose parameter is a type the native + // side can type-check against the list's element field. `exists` adds a condition of its own, + // since Spark stops at the first match. It needs bound references: the native side re-indexes + // the columns the body captures, which it cannot do for unbound ones. protected def nativeLambdaSupported(argument: Expression, body: Expression): Boolean = CometConf.COMET_NATIVE_LAMBDA_ENABLED.get() && body.deterministic && !body.exists(_.isInstanceOf[HigherOrderFunction]) && capturesOnlyFixedWidthColumns(body) && shortCircuitSafe(body) && - !castsToMap(argument) && + !canProduceNullMapKey(argument) && (argument.dataType match { case ArrayType(elementType, _) => isLambdaElementType(elementType) case _ => false }) /** - * Whether `argument` casts to a type that holds a map. A TRY cast of a map whose key cast can - * fail makes a null key, which Arrow cannot hold (#6172), so such an argument stays on the - * dispatcher, which falls back to Spark for it. + * Whether `argument` has a TRY cast whose map key cast can fail, which makes a null key that + * Arrow cannot hold (#6172). The native side would have to build that map, so such an argument + * stays on the dispatcher, which runs it in Spark's own code, where the map never reaches + * Arrow. */ - private def castsToMap(argument: Expression): Boolean = { - def holdsMap(dt: DataType): Boolean = dt match { - case _: MapType => true - case ArrayType(element, _) => holdsMap(element) - case StructType(fields) => fields.exists(f => holdsMap(f.dataType)) - case _ => false - } + private def canProduceNullMapKey(argument: Expression): Boolean = argument.exists { - case c: Cast => holdsMap(c.dataType) + case c: Cast => CometCast.canProduceNullMapKey(c) case _ => false } - } /** * DataFusion copies each column the body captures once per element of its row's array before it @@ -1101,7 +1096,7 @@ trait NativeLambdaSerde extends CometTypeShim with CometExprShim { * 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 { + protected 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 | @@ -1452,8 +1447,15 @@ object CometArrayTransform extends CometCodegenDispatch[ArrayTransform] * otherwise. Spark returns false instead of NULL in that last case when * `spark.sql.legacy.followThreeValuedLogicInArrayExists` is false, so `exists` stays on the * dispatcher with that setting. + * + * Spark also stops at the first element that matches, so it never evaluates the body on the + * elements after it, while `array_any_match` evaluates the body on every element of the batch. A + * body in which any expression can raise would raise on an element Spark never reaches, so it + * stays on the dispatcher whatever operators it uses. */ -object CometArrayExists extends CometCodegenDispatch[ArrayExists] with NativeLambdaSerde { +object CometArrayExists extends CometExpressionSerde[ArrayExists] with NativeLambdaSerde { + + override def getSupportLevel(expr: ArrayExists): SupportLevel = Compatible() override def convert( expr: ArrayExists, @@ -1461,14 +1463,20 @@ object CometArrayExists extends CometCodegenDispatch[ArrayExists] with NativeLam binding: Boolean): Option[ExprOuterClass.Expr] = { expr.function match { case LambdaFunction(body, Seq(lambdaVar: NamedLambdaVariable), _) - if binding && expr.followThreeValuedLogic && nativeLambdaSupported( - expr.argument, - body) => + if binding && expr.followThreeValuedLogic && + nativeLambdaSupported(expr.argument, body) && !body.exists(mayRaise) => convertLambda("array_any_match", expr.argument, lambdaVar, body, inputs, binding) - .orElse(super.convert(expr, inputs, binding)) - case _ => super.convert(expr, inputs, binding) + .orElse(dispatch(expr, inputs, binding)) + case _ => dispatch(expr, inputs, binding) } } + + /** Spark's own evaluation through the codegen dispatcher, as for `transform` and `forall`. */ + private def dispatch( + expr: ArrayExists, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = + CometScalaUDF.emitJvmCodegenDispatch(expr, inputs, binding) } object CometArrayForAll extends CometCodegenDispatch[ArrayForAll] diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda.sql b/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda.sql index 6e6fb150124..9d6534a0322 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda.sql @@ -101,8 +101,16 @@ query SELECT count(CASE WHEN exists(strs, x -> x = '') AND code = 'US' THEN 1 END) FROM test_exists_lambda --- A body that converts each element to a date, all of it natively supported expressions. -query expect_native(exists) +-- Spark stops at the first element that matches, so it never evaluates the lambda on the elements after +-- it, while the native lambda evaluates every element. A lambda in which any expression can raise stays +-- on the dispatcher. `element_at` raises for the index 0 whatever the ANSI setting, and Spark returns +-- true for id 1 before it reaches the 0. +query expect_dispatch(exists) +SELECT id, exists(ints, x -> element_at(array(threshold, 5), x) > 1) FROM test_exists_lambda + +-- `timestamp_seconds` of a BIGINT raises on overflow whatever the ANSI setting, so a body that converts +-- each element to a date this way stays on the dispatcher too. +query expect_dispatch(exists) SELECT id, exists(dbls, x -> datediff(DATE '2024-03-15', CAST(timestamp_seconds(CAST(x AS BIGINT)) AS DATE)) >= 2) FROM test_exists_lambda diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda_ansi.sql b/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda_ansi.sql new file mode 100644 index 00000000000..50dea2c8c4b --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_exists_native_lambda_ansi.sql @@ -0,0 +1,57 @@ +-- 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's `exists` stops at the first element that matches, so it never evaluates the lambda on the +-- elements after it. The native lambda evaluates every element, so under ANSI mode a lambda in which any +-- expression can raise stays on the dispatcher, whatever operators it uses. In every row below an +-- element that matches comes before any element on which the lambda raises, so none of these queries +-- may raise. + +-- Config: spark.sql.ansi.enabled=true +-- Config: spark.sql.legacy.followThreeValuedLogicInArrayExists=true + +statement +CREATE TABLE test_exists_lambda_ansi(id int, ints array, strs array) USING parquet + +statement +INSERT INTO test_exists_lambda_ansi VALUES + (1, array(1, 0), array('5', 'a')), + (2, array(), array()), + (3, NULL, NULL), + (4, array(NULL, 5, 0), array(NULL, '7', '')), + (5, array(5, 2147483647), array('9', '2147483648')), + (6, array(NULL, 1000), array(NULL, '0')) + +query expect_dispatch(exists) +SELECT id, exists(ints, x -> 100 / x > 1) FROM test_exists_lambda_ansi + +query expect_dispatch(exists) +SELECT id, exists(ints, x -> x + 1 > 1) FROM test_exists_lambda_ansi + +query expect_dispatch(exists) +SELECT id, exists(strs, x -> CAST(x AS INT) > 1) FROM test_exists_lambda_ansi + +-- CASE is lazy natively too, but that does not help when the lambda raises on an element after the match. +query expect_dispatch(exists) +SELECT id, exists(ints, x -> CASE WHEN x IS NULL THEN false ELSE 10 / x > 1 END) FROM test_exists_lambda_ansi + +-- Nothing here can raise, so these stay native. +query expect_native(exists) +SELECT id, exists(ints, x -> x > 1) FROM test_exists_lambda_ansi + +query expect_native(exists) +SELECT id, exists(strs, x -> x IS NULL OR x = '') FROM test_exists_lambda_ansi 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 01296201079..7cde3143491 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 @@ -75,10 +75,6 @@ SELECT id, filter(strs, x -> x IS NOT NULL AND x != '') FROM test_filter_lambda 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 expect_native(filter) -SELECT id, filter(ints, x -> nvl(x, '') != '') FROM test_filter_lambda - query expect_native(filter) SELECT id, filter(ints, x -> x > 0) FROM test_filter_lambda diff --git a/spark/src/test/resources/sql-tests/expressions/misc/codegen_dispatch_null_map_key.sql b/spark/src/test/resources/sql-tests/expressions/misc/codegen_dispatch_null_map_key.sql index b6242f5c5a6..cb304234d71 100644 --- a/spark/src/test/resources/sql-tests/expressions/misc/codegen_dispatch_null_map_key.sql +++ b/spark/src/test/resources/sql-tests/expressions/misc/codegen_dispatch_null_map_key.sql @@ -45,6 +45,11 @@ query expect_dispatch(exists) SELECT id, exists(map_keys(try_cast(m AS map)), k -> k IS NULL) FROM test_null_map_key ORDER BY id +-- A widening key cast cannot make a null key, so the lambda runs natively +query expect_native(exists) +SELECT id, exists(map_keys(try_cast(m AS map)), k -> k IS NULL) +FROM test_null_map_key ORDER BY id + -- A widening key cast cannot fail, so the tree stays in the dispatcher query expect_dispatch(transform_values) SELECT id, map_keys(transform_values(try_cast(m AS map), (k, v) -> v)) diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenHOFSuite.scala index 98ee299da38..9b2511ce0d7 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. * @@ -98,27 +96,6 @@ class CometCodegenHOFSuite } } - test("ArrayExists respects three-valued logic through codegen dispatcher") { - withArrayIntTable("(array(1, null, 3)), (array(1, 2, 3)), (array()), (null)") { - // With the legacy setting off Spark returns false where DataFusion's array_any_match returns - // NULL, so only the dispatcher runs `exists`. With it on the lambda runs natively and the - // dispatcher stays idle, with the same answer as Spark. - withSQLConf("spark.sql.legacy.followThreeValuedLogicInArrayExists" -> "false") { - assertCodegenRan { - checkSparkAnswerAndOperator(sql("SELECT exists(a, x -> x > 10) FROM t")) - } - } - withSQLConf("spark.sql.legacy.followThreeValuedLogicInArrayExists" -> "true") { - CometScalaUDFCodegen.resetStats() - checkSparkAnswerAndOperator(sql("SELECT exists(a, x -> x > 10) FROM t")) - val after = CometScalaUDFCodegen.stats() - assert( - after.compileCount + after.cacheHitCount == 0, - s"expected exists to run natively, but the dispatcher ran: $after") - } - } - } - test("HOF query produces correct results across two collects (per-task isolation regression)") { // Regresses the per-task `boundExpr` isolation. When the dispatcher's compile cache lived on // the companion object, multiple tasks shared one `boundExpr` and concurrent partitions From 940918ed3127c63e97de5d6c058c3a152fa36eef Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 10 Nov 2026 08:38:57 -0700 Subject: [PATCH 8/8] fix: keep a native lambda body with a null map key or a date/time cast on the dispatcher A body that holds a TRY cast making a null map key cannot be built natively, as an argument with one cannot (#6172), and a cast between a date or timestamp and a string, date or timestamp fails natively for years beyond chrono's range, which Spark supports (#7471). Both now stay on the dispatcher, with a fixture that pins the path and the results. --- .../user-guide/latest/tuning/expressions.md | 4 ++ .../scala/org/apache/comet/serde/arrays.scala | 33 +++++++++-- ..._filter_native_lambda_unsupported_body.sql | 59 +++++++++++++++++++ 3 files changed, 90 insertions(+), 6 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_unsupported_body.sql diff --git a/docs/source/user-guide/latest/tuning/expressions.md b/docs/source/user-guide/latest/tuning/expressions.md index 401ba77d4c8..69af0d0a989 100644 --- a/docs/source/user-guide/latest/tuning/expressions.md +++ b/docs/source/user-guide/latest/tuning/expressions.md @@ -65,6 +65,10 @@ To get the most out of Comet when a query uses lambdas: under ANSI mode `exists(array(1, 0), x -> 100 / x > 1)` returns true without dividing by zero, but the native lambda evaluates every element. An `exists` lambda that contains any operation that can raise stays on the dispatcher, whether or not it uses `AND`, `OR`, `IN` or `CASE WHEN`. +- **Keep date and timestamp conversions out of the lambda.** A cast between a date or timestamp and a string, date or + timestamp fails natively for years beyond about 262,143, which Spark supports up to year 294,247 + ([#7471](https://github.com/apache/datafusion-comet/issues/7471)), so a lambda that contains one stays on the + dispatcher. So does a lambda with a `try_cast` that can make a `NULL` map key. - **Capture only numeric, decimal, 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. 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 f38a003db66..9c9f7d2d10d 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -1047,23 +1047,44 @@ trait NativeLambdaSerde extends CometTypeShim with CometExprShim { capturesOnlyFixedWidthColumns(body) && shortCircuitSafe(body) && !canProduceNullMapKey(argument) && + !canProduceNullMapKey(body) && + !convertsDateTime(body) && (argument.dataType match { case ArrayType(elementType, _) => isLambdaElementType(elementType) case _ => false }) /** - * Whether `argument` has a TRY cast whose map key cast can fail, which makes a null key that - * Arrow cannot hold (#6172). The native side would have to build that map, so such an argument - * stays on the dispatcher, which runs it in Spark's own code, where the map never reaches - * Arrow. + * Whether `e` has a TRY cast whose map key cast can fail, which makes a null key that Arrow + * cannot hold (#6172). The native side would have to build that map, so an argument or a body + * with one stays on the dispatcher, which runs it in Spark's own code, where the map never + * reaches Arrow. */ - private def canProduceNullMapKey(argument: Expression): Boolean = - argument.exists { + private def canProduceNullMapKey(e: Expression): Boolean = + e.exists { case c: Cast => CometCast.canProduceNullMapKey(c) case _ => false } + /** + * Whether `body` casts between a date or timestamp and a string, date or timestamp. The native + * cast fails for years beyond chrono's range (about 262,143), which Spark supports up to year + * 294,247 (#7471), so such a body stays on the dispatcher. + */ + private def convertsDateTime(body: Expression): Boolean = { + def isDateTime(dt: DataType): Boolean = dt match { + case DateType | TimestampType | TimestampNTZType => true + case _ => false + } + def isTextOrDateTime(dt: DataType): Boolean = dt.isInstanceOf[StringType] || isDateTime(dt) + body.exists { + case c: Cast if c.child.dataType != c.dataType => + (isDateTime(c.child.dataType) && isTextOrDateTime(c.dataType)) || + (isDateTime(c.dataType) && isTextOrDateTime(c.child.dataType)) + 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 diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_unsupported_body.sql b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_unsupported_body.sql new file mode 100644 index 00000000000..e1e3d8d2460 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_filter_native_lambda_unsupported_body.sql @@ -0,0 +1,59 @@ +-- 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 + +-- A lambda body with an expression that the native side cannot evaluate like Spark stays on the +-- dispatcher, which runs it in Spark's own code. An array argument or a body can hold a TRY cast +-- that makes a null map key (#6172), which Arrow cannot represent, and a cast between a date or +-- timestamp and another type fails natively for years beyond chrono's range (#7471). + +-- ConfigMatrix: spark.sql.ansi.enabled=true,false + +statement +CREATE TABLE test_filter_lambda_unsupported(id int, longs array, dates array, tss array) USING parquet + +statement +INSERT INTO test_filter_lambda_unsupported VALUES + (1, array(1L, 9999999999L), array(date_from_unix_date(100000000), DATE '2020-01-01', NULL), + array(timestamp_micros(9000000000000000000), TIMESTAMP '2020-01-01 00:00:00', NULL)), + (2, NULL, NULL, NULL), + (3, array(2L), array(DATE '2024-02-29'), array(TIMESTAMP '2024-02-29 12:34:56')) + +-- The key cast makes a null key for 9999999999 under TRY, which Spark's filter drops. +query expect_dispatch(filter) +SELECT id, filter(longs, x -> array_contains(map_keys(try_cast(map_from_arrays(array(x), array(1)) AS MAP)), 1)) FROM test_filter_lambda_unsupported + +query expect_dispatch(exists) +SELECT id, exists(longs, x -> array_contains(map_keys(try_cast(map_from_arrays(array(x), array(1)) AS MAP)), 1)) FROM test_filter_lambda_unsupported + +-- A widening key cast cannot make a null key. +query expect_native(filter) +SELECT id, filter(longs, x -> array_contains(map_keys(cast(map_from_arrays(array(x), array(1)) AS MAP)), 1L)) FROM test_filter_lambda_unsupported + +query expect_dispatch(filter) +SELECT id, filter(dates, x -> CAST(x AS STRING) <> '') FROM test_filter_lambda_unsupported + +query expect_dispatch(filter) +SELECT id, filter(dates, x -> year(CAST(x AS TIMESTAMP)) > 2000) FROM test_filter_lambda_unsupported + +query expect_dispatch(filter) +SELECT id, filter(tss, x -> year(CAST(x AS DATE)) > 2000) FROM test_filter_lambda_unsupported + +query expect_dispatch(exists) +SELECT id, exists(tss, x -> CAST(x AS STRING) <> '') FROM test_filter_lambda_unsupported + +-- A cast between a timestamp and a number is not affected. +query expect_native(filter) +SELECT id, filter(tss, x -> CAST(x AS BIGINT) > 0) FROM test_filter_lambda_unsupported