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..d2c5e360e21 100644 --- a/native/spark-expr/src/array_funcs/size.rs +++ b/native/spark-expr/src/array_funcs/size.rs @@ -130,11 +130,9 @@ 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, 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 07edb45b741..54fc5d00605 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -1066,16 +1066,24 @@ object CometSize extends CometExpressionSerde[Size] { inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { val arrayExprProto = exprToProtoInternal(expr.child, inputs, binding) + + // 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) + } + 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() @@ -1095,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 761bc954df7..13d82e0151c 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,45 @@ 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") { + // 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, + 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"