perf: generate the in-memory cache reader's class once per executor

.github/workflows/pr_build_linux.yml

@@ -560,6 +560,7 @@ jobs:
               org.apache.spark.sql.comet.CometTPCDSV1_4_PlanStabilitySuite
               org.apache.spark.sql.comet.CometTPCDSV2_7_PlanStabilitySuite
               org.apache.spark.sql.comet.CometTaskMetricsSuite
+              org.apache.spark.sql.comet.CometUnsafeProjectionSuite
               org.apache.spark.sql.comet.CometDppFallbackRepro3949Suite
               org.apache.spark.sql.comet.CometShuffleFallbackStickinessSuite
               org.apache.spark.sql.comet.PlanDataInjectorSuite

.github/workflows/pr_build_macos.yml

@@ -264,6 +264,7 @@ jobs:
               org.apache.spark.sql.comet.CometTPCDSV1_4_PlanStabilitySuite
               org.apache.spark.sql.comet.CometTPCDSV2_7_PlanStabilitySuite
               org.apache.spark.sql.comet.CometTaskMetricsSuite
+              org.apache.spark.sql.comet.CometUnsafeProjectionSuite
               org.apache.spark.sql.comet.CometDppFallbackRepro3949Suite
               org.apache.spark.sql.comet.CometShuffleFallbackStickinessSuite
               org.apache.spark.sql.comet.PlanDataInjectorSuite

spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala

@@ -77,7 +77,8 @@ case class CometColumnarToRowExec(child: SparkPlan)
     // plan (this) in the closure.
     val localOutput = this.output
     child.executeColumnar().mapPartitionsInternal { batches =>
-      val toUnsafe = UnsafeProjection.create(localOutput, localOutput)
+      // Outside whole-stage codegen this runs once per partition, so reuse the generated class.
+      val toUnsafe = CometUnsafeProjection.create(localOutput)
       batches.flatMap { batch =>
         numInputBatches += 1
         numOutputRows += batch.numRows().toLong

spark/src/main/scala/org/apache/spark/sql/comet/CometUnsafeProjection.scala

@@ -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.
+ */
+
+package org.apache.spark.sql.comet
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, CodeGeneratorWithInterpretedFallback, InterpretedUnsafeProjection, UnsafeProjection}
+import org.apache.spark.sql.catalyst.expressions.codegen.{CodeAndComment, CodeFormatter, CodegenContext, CodeGenerator, GeneratedClass, GenerateUnsafeProjection}
+
+/**
+ * Creates the projection that copies each row of a batch into an UnsafeRow, as
+ * `UnsafeProjection.create(output, output)` does, but generates and compiles its class only once
+ * per executor for each column layout.
+ *
+ * `UnsafeProjection.create` generates the projection's Java source on every call, and only the
+ * compiled class is cached. `CometColumnarToRowExec` calls it once per partition whenever it runs
+ * outside whole-stage codegen, which Spark skips for a schema with more fields than
+ * `spark.sql.codegen.maxFields`, nested fields included. Spark's own scans fall back to rows for
+ * such schemas, so its `ColumnarToRowExec` rarely meets one, but Comet's operators always produce
+ * batches. The source grows with the number of nested fields, and each level of nesting repeats
+ * the work of splitting its children into methods: for 100 columns nested three levels deep,
+ * generating it took about 19 ms in every partition.
+ *
+ * Every call returns a new instance of the shared class, so callers own their projection as they
+ * do one from `UnsafeProjection.create`.
+ */
+private[comet] object CometUnsafeProjection
+    extends CodeGeneratorWithInterpretedFallback[Seq[BoundReference], UnsafeProjection] {
+
+  /** A compiled projection class and the objects its instances reference. */
+  private case class Generated(generatedClass: GeneratedClass, references: Array[Any]) {
+    def newProjection(): UnsafeProjection =
+      generatedClass.generate(references).asInstanceOf[UnsafeProjection]
+  }
+
+  private val classes = new GeneratedClassCache[ColumnLayout, Generated]()
+
+  /** How many projection classes this executor has generated, for tests. */
+  private[comet] def generatedClassCount: Long = classes.generatedCount
+
+  /** A projection of rows with the columns of `output` to UnsafeRows. */
+  def create(output: Seq[Attribute]): UnsafeProjection =
+    createObject(output.zipWithIndex.map { case (attr, ordinal) =>
+      BoundReference(ordinal, attr.dataType, attr.nullable)
+    })
+
+  override protected def createCodeGeneratedObject(
+      columns: Seq[BoundReference]): UnsafeProjection =
+    classes
+      .getOrGenerate(ColumnLayout.of(columns)) {
+        val generated = generate(columns)
+        // Instances of a shared class share its references, so share only a class that has
+        // none. GenerateUnsafeProjection references no objects for bound columns.
+        (generated, generated.references.isEmpty)
+      }
+      .newProjection()
+
+  override protected def createInterpretedObject(columns: Seq[BoundReference]): UnsafeProjection =
+    InterpretedUnsafeProjection.createProjection(columns)
+
+  /**
+   * Generates and compiles the class that `GenerateUnsafeProjection.create` does, from the same
+   * template, which is repeated here because that method returns only an instance.
+   */
+  private def generate(columns: Seq[BoundReference]): Generated = {
+    val ctx = new CodegenContext
+    val eval = GenerateUnsafeProjection.createCode(ctx, columns)
+    val body =
+      s"""
+         |public java.lang.Object generate(Object[] references) {
+         |  return new SpecificUnsafeProjection(references);
+         |}
+         |
+         |class SpecificUnsafeProjection extends ${classOf[UnsafeProjection].getName} {
+         |
+         |  private Object[] references;
+         |  ${ctx.declareMutableStates()}
+         |
+         |  public SpecificUnsafeProjection(Object[] references) {
+         |    this.references = references;
+         |    ${ctx.initMutableStates()}
+         |  }
+         |
+         |  public void initialize(int partitionIndex) {
+         |    ${ctx.initPartition()}
+         |  }
+         |
+         |  // Scala.Function1 need this
+         |  public java.lang.Object apply(java.lang.Object row) {
+         |    return apply((InternalRow) row);
+         |  }
+         |
+         |  public UnsafeRow apply(InternalRow ${ctx.INPUT_ROW}) {
+         |    ${eval.code}
+         |    return ${eval.value};
+         |  }
+         |
+         |  ${ctx.declareAddedFunctions()}
+         |}
+       """.stripMargin
+    val code = CodeFormatter.stripOverlappingComments(
+      new CodeAndComment(body, ctx.getPlaceHolderToComments()))
+    val (generatedClass, _) = CodeGenerator.compile(code)
+    Generated(generatedClass, ctx.references.toArray)
+  }
+}

spark/src/main/scala/org/apache/spark/sql/comet/GeneratedClassCache.scala

@@ -0,0 +1,90 @@
+/*
+ * 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.comet
+
+import java.util.{LinkedHashMap => JLinkedHashMap, Map => JMap}
+import java.util.concurrent.atomic.AtomicLong
+
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.DataType
+
+/**
+ * What code generated over a list of columns depends on: the type and nullability of each column,
+ * and the method size at which `CodegenContext.splitExpressions` splits the code for them.
+ */
+private[comet] case class ColumnLayout(
+    columns: Seq[(DataType, Boolean)],
+    methodSplitThreshold: Int)
+
+private[comet] object ColumnLayout {
+
+  /** The layout of `columns` under the current `spark.sql.codegen.methodSplitThreshold`. */
+  def of(columns: Seq[Expression]): ColumnLayout =
+    ColumnLayout(columns.map(c => (c.dataType, c.nullable)), SQLConf.get.methodSplitThreshold)
+}
+
+/**
+ * Generated classes kept for the life of an executor, under a key for what their source depends
+ * on.
+ *
+ * Spark caches a compiled class under its source, so an operator that generates code in every
+ * partition still generates the source every time, which for a wide or deeply nested schema can
+ * take longer than the partition's rows. Keeping the class under a cheaper key generates it once.
+ * The cache holds the `maxEntries` most recently used values.
+ */
+private[comet] class GeneratedClassCache[K, V <: AnyRef](
+    maxEntries: Int = GeneratedClassCache.MaxEntries) {
+
+  /** Least recently used first. Guarded by `entries.synchronized`. */
+  private val entries = new JLinkedHashMap[K, V](16, 0.75f, true) {
+    override def removeEldestEntry(eldest: JMap.Entry[K, V]): Boolean = size() > maxEntries
+  }
+
+  private val generated = new AtomicLong(0)
+
+  /** How many values this cache has generated, for tests. */
+  def generatedCount: Long = generated.get()
+
+  /**
+   * The value for `key`, generated if absent. `generate` returns the value and whether other
+   * callers may share it. It runs outside the lock, so callers that miss on the same key at once
+   * each generate the value, as they would without the cache.
+   */
+  def getOrGenerate(key: K)(generate: => (V, Boolean)): V = {
+    val cached = entries.synchronized(entries.get(key))
+    if (cached != null) {
+      cached
+    } else {
+      val (value, shareable) = generate
+      generated.incrementAndGet()
+      if (shareable) {
+        val _ = entries.synchronized(entries.putIfAbsent(key, value))
+      }
+      value
+    }
+  }
+}
+
+private[comet] object GeneratedClassCache {
+
+  /** The default bound of Spark's compiled class cache, `spark.sql.codegen.cache.maxEntries`. */
+  val MaxEntries = 100
+}

spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIterator.scala

@@ -23,6 +23,7 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, CodeGeneratorWithInterpretedFallback, InterpretedUnsafeProjection, LeafExpression, UnsafeProjection}
 import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.catalyst.expressions.codegen.Block._
+import org.apache.spark.sql.comet.{ColumnLayout, CometUnsafeProjection, GeneratedClassCache}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.DataType
 import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
@@ -37,6 +38,12 @@ import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
  * splits the field writes of a wide projection into methods of bounded size, as it does for any
  * Spark projection. If a generated method still exceeds the huge-method limit, the reader backs
  * off the way WholeStageCodegenExec does, here to Spark's projection of each batch row.
+ *
+ * Spark reads a cache through this path, once per partition, when it does not read the cache as
+ * batches, which is always the case for a schema with more fields than
+ * `spark.sql.codegen.maxFields`, nested fields included. For such a schema generating the source
+ * can take longer than reading a partition's rows, so the reader's class is generated once per
+ * executor for each column layout.
  */
 private[arrow] class CachedBatchRowIterator(attributes: Seq[Attribute])
     extends CodeGeneratorWithInterpretedFallback[Iterator[ColumnarBatch], Iterator[InternalRow]] {
@@ -47,6 +54,35 @@ private[arrow] class CachedBatchRowIterator(attributes: Seq[Attribute])
 
   override protected def createCodeGeneratedObject(
       batches: Iterator[ColumnarBatch]): Iterator[InternalRow] = {
+    val reader = CachedBatchRowIterator.readers.getOrGenerate(ColumnLayout.of(attributes)) {
+      val generated = generate()
+      // Instances of a shared class share its references other than the batches, so share only
+      // a class that has no others.
+      (generated, generated.references.length == 1)
+    }
+    // Honor spark.sql.codegen.hugeMethodLimit as whole-stage codegen does, but never go above
+    // HotSpot's own limit: the config defaults to the largest method the JVM accepts, while this
+    // runs once per row and HotSpot never JIT-compiles a method longer than
+    // DEFAULT_JVM_HUGE_METHOD_LIMIT bytes.
+    val limit =
+      math.min(SQLConf.get.hugeMethodLimit, CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT)
+    if (reader.maxMethodCodeSize > limit) {
+      logInfo(
+        s"Generated cache reader for ${attributes.length} columns has a " +
+          s"${reader.maxMethodCodeSize}-byte method, above the $limit-byte limit; " +
+          "projecting cached rows with UnsafeProjection instead")
+      new ProjectedRows(batches, CometUnsafeProjection.create(attributes))
+    } else {
+      reader.newIterator(batches)
+    }
+  }
+
+  override protected def createInterpretedObject(
+      batches: Iterator[ColumnarBatch]): Iterator[InternalRow] =
+    new ProjectedRows(batches, InterpretedUnsafeProjection.createProjection(fields))
+
+  /** Generates and compiles the reader's class, which takes its batches as a reference. */
+  private def generate(): CachedBatchRowIterator.GeneratedReader = {
     val ctx = new CodegenContext
     val vectorClass = classOf[ColumnVector].getName
     val batchClass = classOf[ColumnarBatch].getName
@@ -62,7 +98,9 @@ private[arrow] class CachedBatchRowIterator(attributes: Seq[Attribute])
     // With ctx.currentVars unset, GenerateUnsafeProjection splits the field writes into methods
     // that take the input row as their argument. The reads above ignore it.
     val projection = GenerateUnsafeProjection.createCode(ctx, reads)
-    val batchesRef = ctx.addReferenceObj("batches", batches, "scala.collection.Iterator")
+    // Each instance gets its own batches, set in a copy of the references.
+    val batchesIndex = ctx.references.length
+    val batchesRef = ctx.addReferenceObj("batches", null, "scala.collection.Iterator")
     val code = s"""
       public Object generate(Object[] references) {
         return new SpecificCachedBatchRowIterator(references);
@@ -105,26 +143,37 @@ private[arrow] class CachedBatchRowIterator(attributes: Seq[Attribute])
     """
     val (compiled, stats) =
       CodeGenerator.compile(new CodeAndComment(code, ctx.getPlaceHolderToComments()))
-    // Honor spark.sql.codegen.hugeMethodLimit as whole-stage codegen does, but never go above
-    // HotSpot's own limit: the config defaults to the largest method the JVM accepts, while this
-    // runs once per row and HotSpot never JIT-compiles a method longer than
-    // DEFAULT_JVM_HUGE_METHOD_LIMIT bytes.
-    val limit =
-      math.min(SQLConf.get.hugeMethodLimit, CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT)
-    if (stats.maxMethodCodeSize > limit) {
-      logInfo(
-        s"Generated cache reader for ${attributes.length} columns has a " +
-          s"${stats.maxMethodCodeSize}-byte method, above the $limit-byte limit; " +
-          "projecting cached rows with UnsafeProjection instead")
-      new ProjectedRows(batches, UnsafeProjection.create(fields))
-    } else {
-      compiled.generate(ctx.references.toArray).asInstanceOf[Iterator[InternalRow]]
+    CachedBatchRowIterator.GeneratedReader(
+      compiled,
+      stats.maxMethodCodeSize,
+      ctx.references.toArray,
+      batchesIndex)
+  }
+}
+
+private[arrow] object CachedBatchRowIterator {
+
+  /**
+   * A compiled reader class, the size of its largest method, and the objects its instances
+   * reference, with a null slot at `batchesIndex` for each instance's batches.
+   */
+  private case class GeneratedReader(
+      generatedClass: GeneratedClass,
+      maxMethodCodeSize: Int,
+      references: Array[Any],
+      batchesIndex: Int) {
+
+    def newIterator(batches: Iterator[ColumnarBatch]): Iterator[InternalRow] = {
+      val instanceReferences = references.clone()
+      instanceReferences(batchesIndex) = batches
+      generatedClass.generate(instanceReferences).asInstanceOf[Iterator[InternalRow]]
     }
   }
 
-  override protected def createInterpretedObject(
-      batches: Iterator[ColumnarBatch]): Iterator[InternalRow] =
-    new ProjectedRows(batches, InterpretedUnsafeProjection.createProjection(fields))
+  private val readers = new GeneratedClassCache[ColumnLayout, GeneratedReader]()
+
+  /** How many reader classes this executor has generated, for tests. */
+  private[arrow] def generatedClassCount: Long = readers.generatedCount
 }
 
 /**