perf: evaluate the child of size once when it needs no NULL guard

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),

native/spark-expr/src/array_funcs/size.rs

@@ -130,11 +130,9 @@ fn spark_size_list_like(array: &ArrayRef) -> Result<ArrayRef, DataFusionError> {
         }
     };
 
-    // 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);
     }

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 {

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"