From ad2587dc50cf6ab49e29227466bfa77600243ca8 Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 3 Nov 2026 19:13:16 -0800 Subject: [PATCH 1/2] perf: evaluate the child of `size` once when it needs no NULL guard `CometSize` serialized its child twice, as `CASE WHEN isnotnull(child) THEN size(child) ELSE END`, so DataFusion evaluated it once for the guard and again for the size. Native `size` already gives -1 for a NULL array or map, the legacy answer, so emit it alone when `legacySizeOfNull` is on or the child cannot be NULL. The guard stays for a nullable child when a NULL has to give NULL. --- native/spark-expr/benches/array_size.rs | 7 ++- native/spark-expr/src/array_funcs/size.rs | 12 +++-- .../scala/org/apache/comet/serde/arrays.scala | 10 ++++ .../comet/CometArrayExpressionSuite.scala | 47 +++++++++++++++++++ 4 files changed, 67 insertions(+), 9 deletions(-) diff --git a/native/spark-expr/benches/array_size.rs b/native/spark-expr/benches/array_size.rs index c82468f462a..70f8f937214 100644 --- a/native/spark-expr/benches/array_size.rs +++ b/native/spark-expr/benches/array_size.rs @@ -136,10 +136,9 @@ fn criterion_benchmark(c: &mut Criterion) { &create_list_array(rows, 64, true), ); - // No-null shape: `CometSize.convert` wraps size() in a `CASE WHEN isnotnull(child)` - // that filters null rows out before the THEN branch runs, so in a real Comet plan - // spark_size_list_like only ever sees a null-free array. This shape measures that - // path. + // No-null shape: an array without nulls, such as the one `CometSize.convert`'s + // `CASE WHEN isnotnull(child)` guard passes on when `legacySizeOfNull` is off, + // takes spark_size_list_like's fast path. This shape measures that path. bench( "spark_size: list, no nulls", &create_list_array(rows, 5, false), diff --git a/native/spark-expr/src/array_funcs/size.rs b/native/spark-expr/src/array_funcs/size.rs index 2311da5fa9e..472e91c4cc7 100644 --- a/native/spark-expr/src/array_funcs/size.rs +++ b/native/spark-expr/src/array_funcs/size.rs @@ -130,11 +130,13 @@ fn spark_size_list_like(array: &ArrayRef) -> Result { } }; - // Fast path for the production shape: `CometSize.convert` wraps size() in a - // `CASE WHEN isnotnull(child)` that filters null rows out before the THEN - // branch runs, so this function only ever sees a null-free array in a real - // Comet plan. Return the length kernel output as-is; skip the downcast and - // `Int32Array::clone` that `rewrite_nulls_to_minus_one` would otherwise pay. + // Fast path for an array without nulls. `CometSize.convert` keeps its + // `CASE WHEN isnotnull(child)` guard, which filters null rows out before + // the THEN branch runs, only when `legacySizeOfNull` is off and the child + // is nullable. Otherwise nulls reach this function and the rewrite below + // turns them into -1. Return the length kernel output as-is; skip the + // downcast and `Int32Array::clone` that `rewrite_nulls_to_minus_one` would + // otherwise pay. if array.null_count() == 0 { return Ok(lengths); } 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..09001f747bc 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -1066,6 +1066,16 @@ object CometSize extends CometExpressionSerde[Size] { inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { val arrayExprProto = exprToProtoInternal(expr.child, inputs, binding) + + // Native size already gives -1 for a NULL array or map, the legacy answer, and a child that + // cannot be NULL never takes the NULL branch, so the guard only matters when a NULL has to + // give NULL. Without it the child is evaluated once over the whole batch rather than once + // for the guard and again for the size, so an expensive child such as a Scala UDF is not + // paid for twice and a nondeterministic one stays in step. + if (expr.legacySizeOfNull || !expr.child.nullable) { + return scalarFunctionExprToProto("size", arrayExprProto) + } + for { isNotNullExprProto <- createIsNotNullExprProto(expr, inputs, binding) sizeScalarExprProto <- scalarFunctionExprToProto("size", arrayExprProto) diff --git a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala index 761bc954df7..5c72afd414c 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -19,6 +19,8 @@ package org.apache.comet +import java.util.concurrent.atomic.AtomicLong + import scala.util.Random import org.apache.hadoop.fs.Path @@ -37,6 +39,11 @@ import org.apache.comet.DataTypeSupport.isComplexType import org.apache.comet.serde.{ArrayElementEqualitySupport, ArraySetSupport, CometArrayDistinct, CometArrayExcept, CometArrayJoin, CometArrayRemove, CometArrayReverse, CometArraysOverlap, CometArrayUnion, CometFlatten, CometSequence, Compatible, ExprOuterClass, Incompatible, SparkFixes} import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions} +/** Counts calls to a test UDF. Comet tests run in local mode, so tasks see this object. */ +object SizeChildCounter { + val calls = new AtomicLong(0) +} + class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { test("array set signed-zero Spark patch versions") { @@ -1346,6 +1353,46 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp } } + test("size evaluates its child once") { + // A NULL guard around size serializes the child a second time, so a UDF child would be + // called once for the guard and again for every row whose array is not NULL. + spark.udf.register( + "splitOrNull", + (s: String) => { + SizeChildCounter.calls.incrementAndGet() + if (s == null) null else s.split(",", -1).toSeq + }) + // `udf` alone names the org.apache.comet.udf package here on Scala 2.12. + val splitOrEmpty = org.apache.spark.sql.functions.udf { (s: String) => + SizeChildCounter.calls.incrementAndGet() + if (s == null) Seq.empty[String] else s.split(",", -1).toSeq + } + spark.udf.register("splitOrEmpty", splitOrEmpty.asNonNullable()) + // Each call plans the query anew, so it runs with the configuration in effect. + def calls(query: String): Long = { + SizeChildCounter.calls.set(0) + sql(query).collect() + SizeChildCounter.calls.get() + } + withParquetTable(Seq("a,b,c", null, "", "x,y").map(Tuple1(_)), "tbl") { + // Native size already gives -1 for NULL, the legacy answer, and a child that cannot be + // NULL needs no guard. + Seq("true" -> "splitOrNull", "false" -> "splitOrEmpty").foreach { case (legacy, split) => + withSQLConf( + SQLConf.LEGACY_SIZE_OF_NULL.key -> legacy, + SQLConf.ANSI_ENABLED.key -> "false") { + val query = s"SELECT size($split(_1)) FROM tbl" + checkSparkAnswerAndOperator(query) + var sparkCalls = 0L + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkCalls = calls(query) + } + assert(calls(query) == sparkCalls) + } + } + } + } + // https://github.com/apache/datafusion-comet/issues/4560 test("array_size returns null for null input") { val table = "t1" From 01e01608510e040809549913c63b4d1c1b9f45c1 Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 5 Nov 2026 08:52:19 -0800 Subject: [PATCH 2/2] address review: build the NULL literal directly, trim comments --- native/spark-expr/src/array_funcs/size.rs | 10 +++------- .../scala/org/apache/comet/serde/arrays.scala | 17 +++++------------ .../comet/CometArrayExpressionSuite.scala | 3 +-- 3 files changed, 9 insertions(+), 21 deletions(-) diff --git a/native/spark-expr/src/array_funcs/size.rs b/native/spark-expr/src/array_funcs/size.rs index 472e91c4cc7..d2c5e360e21 100644 --- a/native/spark-expr/src/array_funcs/size.rs +++ b/native/spark-expr/src/array_funcs/size.rs @@ -130,13 +130,9 @@ fn spark_size_list_like(array: &ArrayRef) -> Result { } }; - // Fast path for an array without nulls. `CometSize.convert` keeps its - // `CASE WHEN isnotnull(child)` guard, which filters null rows out before - // the THEN branch runs, only when `legacySizeOfNull` is off and the child - // is nullable. Otherwise nulls reach this function and the rewrite below - // turns them into -1. Return the length kernel output as-is; skip the - // downcast and `Int32Array::clone` that `rewrite_nulls_to_minus_one` would - // otherwise pay. + // Fast path for an array without nulls, such as the one `CometSize.convert`'s + // `CASE WHEN isnotnull(child)` guard passes on. Skip the downcast and + // `Int32Array::clone` that `rewrite_nulls_to_minus_one` would pay. if array.null_count() == 0 { return Ok(lengths); } 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 09001f747bc..54fc5d00605 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -1067,11 +1067,9 @@ object CometSize extends CometExpressionSerde[Size] { binding: Boolean): Option[ExprOuterClass.Expr] = { val arrayExprProto = exprToProtoInternal(expr.child, inputs, binding) - // Native size already gives -1 for a NULL array or map, the legacy answer, and a child that - // cannot be NULL never takes the NULL branch, so the guard only matters when a NULL has to - // give NULL. Without it the child is evaluated once over the whole batch rather than once - // for the guard and again for the size, so an expensive child such as a Scala UDF is not - // paid for twice and a nondeterministic one stays in step. + // Native size gives -1 for NULL, the legacy answer, so the NULL guard is needed only when + // a NULL child must give NULL. The guard serializes the child twice, which would run an + // expensive child such as a Scala UDF twice per row. if (expr.legacySizeOfNull || !expr.child.nullable) { return scalarFunctionExprToProto("size", arrayExprProto) } @@ -1079,13 +1077,13 @@ object CometSize extends CometExpressionSerde[Size] { for { isNotNullExprProto <- createIsNotNullExprProto(expr, inputs, binding) sizeScalarExprProto <- scalarFunctionExprToProto("size", arrayExprProto) - emptyLiteralExprProto <- createLiteralExprProto(expr.legacySizeOfNull) + nullLiteralExprProto <- exprToProtoInternal(Literal(null, IntegerType), inputs, binding) } yield { val caseWhenExpr = ExprOuterClass.CaseWhen .newBuilder() .addWhen(isNotNullExprProto) .addThen(sizeScalarExprProto) - .setElseExpr(emptyLiteralExprProto) + .setElseExpr(nullLiteralExprProto) .build() ExprOuterClass.Expr .newBuilder() @@ -1105,11 +1103,6 @@ object CometSize extends CometExpressionSerde[Size] { (builder, unaryExpr) => builder.setIsNotNull(unaryExpr)) } - private def createLiteralExprProto(legacySizeOfNull: Boolean): Option[ExprOuterClass.Expr] = { - val value = if (legacySizeOfNull) -1 else null - exprToProtoInternal(Literal(value, IntegerType), Seq.empty, binding = true) - } - } object CometArrayPosition extends CometExpressionSerde[ArrayPosition] with ArraysBase { diff --git a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala index 5c72afd414c..13d82e0151c 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -1375,8 +1375,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp SizeChildCounter.calls.get() } withParquetTable(Seq("a,b,c", null, "", "x,y").map(Tuple1(_)), "tbl") { - // Native size already gives -1 for NULL, the legacy answer, and a child that cannot be - // NULL needs no guard. + // The two shapes that need no NULL guard: legacy size of NULL and a non-nullable child. Seq("true" -> "splitOrNull", "false" -> "splitOrEmpty").foreach { case (legacy, split) => withSQLConf( SQLConf.LEGACY_SIZE_OF_NULL.key -> legacy,