From df0b1b9ff0213cd369a4b68e10b8cb0b9ec06a3d Mon Sep 17 00:00:00 2001 From: msaf Date: Sat, 17 Oct 2026 12:16:01 +0100 Subject: [PATCH 1/2] fix: ship the functions of a native plan's ScalaUDFs once, in one serialization stream The codegen dispatcher closure-serialized every ScalaUDF expression on its own, so a function captured by several expressions, and an object several functions capture, was written once per expression: ltv_flag's project of ten UDFs over one ~150 MB map built ~1.35 GB of bytes on the driver per plan conversion, again on every AQE re-plan, and failed with Java heap space. Each expression now carries its ScalaUDF functions as references. The native plan serialization appends one pool holding all functions its expressions reference, written in one Java serialization stream so shared captured objects are written once, and caches it per function set for AQE re-conversions. Each task registers the pools of its plans and deserializes a pool once, on its first use. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../org/apache/comet/CometExecIterator.scala | 3 + .../apache/comet/serde/CometScalaUDF.scala | 24 +- .../udf/codegen/CometScalaUDFCodegen.scala | 10 +- .../udf/codegen/CometSharedFunctionRef.scala | 555 ++++++++++++++++++ .../udf/codegen/CometUdfFunctionPool.scala | 288 +++++++++ .../apache/spark/sql/comet/operators.scala | 3 +- .../org/apache/comet/CometCodegenSuite.scala | 33 +- 7 files changed, 904 insertions(+), 12 deletions(-) create mode 100644 spark/src/main/scala/org/apache/comet/udf/codegen/CometSharedFunctionRef.scala create mode 100644 spark/src/main/scala/org/apache/comet/udf/codegen/CometUdfFunctionPool.scala diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 4a5e418c1ae..50cc3ae3f6c 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -44,6 +44,7 @@ import org.apache.comet.exceptions.CometQueryExecutionException import org.apache.comet.parquet.CometFileKeyUnwrapper import org.apache.comet.serde.Config.ConfigMap import org.apache.comet.shuffle.ShufflePartitionPusher +import org.apache.comet.udf.codegen.CometUdfFunctionPool import org.apache.comet.vector.NativeUtil /** @@ -123,6 +124,8 @@ class CometExecIterator( val memoryConfig = CometExecIterator.getMemoryConfig(conf) + CometUdfFunctionPool.registerPlan(taskContext, protobufQueryPlan) + val createdPlan = nativeLib.createPlan( id, inputObjects, 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 8fcab69cb48..e0abde27c9c 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala @@ -32,7 +32,7 @@ import org.apache.comet.CometSparkSessionExtensions.{withCodegenDispatchExpr, wi import org.apache.comet.codegen.CometBatchKernelCodegen import org.apache.comet.serde.ExprOuterClass.Expr import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, serializeDataType} -import org.apache.comet.udf.codegen.CometScalaUDFCodegen +import org.apache.comet.udf.codegen.{CometScalaUDFCodegen, CometUdfFunctionPool} /** * Routes scalar `ScalaUDF` (Scala and Java UDFs) through the codegen dispatcher. @@ -124,13 +124,25 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { // fails planning, which is a much worse outcome than falling the operator back to Spark. // `CometStaticInvoke` / `CometInvoke` route unrecognized nodes here as a catch-all, so the // trees reaching this point are arbitrary. + val functions = CometUdfFunctionPool.functionsOf(boundExpr) + functions.find(!CometUdfFunctionPool.isSerializable(_)).foreach { function => + withFallbackReason( + expr, + s"$exprName: codegen dispatch: UDF function ${function.getClass.getName} could not " + + "be closure-serialized") + return None + } val bytes = try { - val serializer = SparkEnv.get.closureSerializer.newInstance() - val buffer = serializer.serialize(boundExpr) - val serialized = new Array[Byte](buffer.remaining()) - buffer.get(serialized) - serialized + if (functions.nonEmpty) { + CometUdfFunctionPool.serializeExpression(boundExpr, functions) + } else { + val serializer = SparkEnv.get.closureSerializer.newInstance() + val buffer = serializer.serialize(boundExpr) + val serialized = new Array[Byte](buffer.remaining()) + buffer.get(serialized) + serialized + } } catch { // `NonFatal` rather than `NotSerializableException`: Java serialization reports an // unserializable object graph as one of several exception types depending on where in diff --git a/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala b/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala index ef541fb3e2d..f4702f5c971 100644 --- a/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala +++ b/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala @@ -155,9 +155,13 @@ class CometScalaUDFCodegen extends CometUDF with Logging { .getOrElse(classOf[Expression].getClassLoader) val boundExpr = try { - SparkEnv.get.closureSerializer - .newInstance() - .deserialize[Expression](ByteBuffer.wrap(bytes), loader) + if (CometUdfFunctionPool.isSharedExpression(bytes)) { + CometUdfFunctionPool.deserializeExpression(bytes, loader) + } else { + SparkEnv.get.closureSerializer + .newInstance() + .deserialize[Expression](ByteBuffer.wrap(bytes), loader) + } } catch { case NonFatal(t) => logError( diff --git a/spark/src/main/scala/org/apache/comet/udf/codegen/CometSharedFunctionRef.scala b/spark/src/main/scala/org/apache/comet/udf/codegen/CometSharedFunctionRef.scala new file mode 100644 index 00000000000..348f16516c1 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/udf/codegen/CometSharedFunctionRef.scala @@ -0,0 +1,555 @@ +/* + * 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.comet.udf.codegen + +// scalastyle:off argcount +sealed trait CometSharedFunctionRef extends Serializable { + def id: Long +} + +object CometSharedFunctionRef { + def apply(id: Long, arity: Int): CometSharedFunctionRef = arity match { + case 0 => CometSharedFunction0(id) + case 1 => CometSharedFunction1(id) + case 2 => CometSharedFunction2(id) + case 3 => CometSharedFunction3(id) + case 4 => CometSharedFunction4(id) + case 5 => CometSharedFunction5(id) + case 6 => CometSharedFunction6(id) + case 7 => CometSharedFunction7(id) + case 8 => CometSharedFunction8(id) + case 9 => CometSharedFunction9(id) + case 10 => CometSharedFunction10(id) + case 11 => CometSharedFunction11(id) + case 12 => CometSharedFunction12(id) + case 13 => CometSharedFunction13(id) + case 14 => CometSharedFunction14(id) + case 15 => CometSharedFunction15(id) + case 16 => CometSharedFunction16(id) + case 17 => CometSharedFunction17(id) + case 18 => CometSharedFunction18(id) + case 19 => CometSharedFunction19(id) + case 20 => CometSharedFunction20(id) + case 21 => CometSharedFunction21(id) + case 22 => CometSharedFunction22(id) + } + + def unresolved(id: Long): IllegalStateException = + new IllegalStateException(s"Comet shared UDF function $id was not resolved") +} + +final case class CometSharedFunction0(id: Long) extends (() => Any) with CometSharedFunctionRef { + override def apply(): Any = throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction1(id: Long) + extends ((Any) => Any) + with CometSharedFunctionRef { + override def apply(a1: Any): Any = throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction2(id: Long) + extends ((Any, Any) => Any) + with CometSharedFunctionRef { + override def apply(a1: Any, a2: Any): Any = throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction3(id: Long) + extends ((Any, Any, Any) => Any) + with CometSharedFunctionRef { + override def apply(a1: Any, a2: Any, a3: Any): Any = throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction4(id: Long) + extends ((Any, Any, Any, Any) => Any) + with CometSharedFunctionRef { + override def apply(a1: Any, a2: Any, a3: Any, a4: Any): Any = + throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction5(id: Long) + extends ((Any, Any, Any, Any, Any) => Any) + with CometSharedFunctionRef { + override def apply(a1: Any, a2: Any, a3: Any, a4: Any, a5: Any): Any = + throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction6(id: Long) + extends ((Any, Any, Any, Any, Any, Any) => Any) + with CometSharedFunctionRef { + override def apply(a1: Any, a2: Any, a3: Any, a4: Any, a5: Any, a6: Any): Any = + throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction7(id: Long) + extends ((Any, Any, Any, Any, Any, Any, Any) => Any) + with CometSharedFunctionRef { + override def apply(a1: Any, a2: Any, a3: Any, a4: Any, a5: Any, a6: Any, a7: Any): Any = + throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction8(id: Long) + extends ((Any, Any, Any, Any, Any, Any, Any, Any) => Any) + with CometSharedFunctionRef { + override def apply(a1: Any, a2: Any, a3: Any, a4: Any, a5: Any, a6: Any, a7: Any, a8: Any) + : Any = throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction9(id: Long) + extends ((Any, Any, Any, Any, Any, Any, Any, Any, Any) => Any) + with CometSharedFunctionRef { + override def apply( + a1: Any, + a2: Any, + a3: Any, + a4: Any, + a5: Any, + a6: Any, + a7: Any, + a8: Any, + a9: Any): Any = throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction10(id: Long) + extends ((Any, Any, Any, Any, Any, Any, Any, Any, Any, Any) => Any) + with CometSharedFunctionRef { + override def apply( + a1: Any, + a2: Any, + a3: Any, + a4: Any, + a5: Any, + a6: Any, + a7: Any, + a8: Any, + a9: Any, + a10: Any): Any = throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction11(id: Long) + extends ((Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any) => Any) + with CometSharedFunctionRef { + override def apply( + a1: Any, + a2: Any, + a3: Any, + a4: Any, + a5: Any, + a6: Any, + a7: Any, + a8: Any, + a9: Any, + a10: Any, + a11: Any): Any = throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction12(id: Long) + extends ((Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any) => Any) + with CometSharedFunctionRef { + override def apply( + a1: Any, + a2: Any, + a3: Any, + a4: Any, + a5: Any, + a6: Any, + a7: Any, + a8: Any, + a9: Any, + a10: Any, + a11: Any, + a12: Any): Any = throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction13(id: Long) + extends ((Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any) => Any) + with CometSharedFunctionRef { + override def apply( + a1: Any, + a2: Any, + a3: Any, + a4: Any, + a5: Any, + a6: Any, + a7: Any, + a8: Any, + a9: Any, + a10: Any, + a11: Any, + a12: Any, + a13: Any): Any = throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction14(id: Long) + extends ((Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any) => Any) + with CometSharedFunctionRef { + override def apply( + a1: Any, + a2: Any, + a3: Any, + a4: Any, + a5: Any, + a6: Any, + a7: Any, + a8: Any, + a9: Any, + a10: Any, + a11: Any, + a12: Any, + a13: Any, + a14: Any): Any = throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction15(id: Long) + extends ((Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any) => Any) + with CometSharedFunctionRef { + override def apply( + a1: Any, + a2: Any, + a3: Any, + a4: Any, + a5: Any, + a6: Any, + a7: Any, + a8: Any, + a9: Any, + a10: Any, + a11: Any, + a12: Any, + a13: Any, + a14: Any, + a15: Any): Any = throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction16(id: Long) + extends ( + ( + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any) => Any) + with CometSharedFunctionRef { + override def apply( + a1: Any, + a2: Any, + a3: Any, + a4: Any, + a5: Any, + a6: Any, + a7: Any, + a8: Any, + a9: Any, + a10: Any, + a11: Any, + a12: Any, + a13: Any, + a14: Any, + a15: Any, + a16: Any): Any = throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction17(id: Long) + extends ( + ( + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any) => Any) + with CometSharedFunctionRef { + override def apply( + a1: Any, + a2: Any, + a3: Any, + a4: Any, + a5: Any, + a6: Any, + a7: Any, + a8: Any, + a9: Any, + a10: Any, + a11: Any, + a12: Any, + a13: Any, + a14: Any, + a15: Any, + a16: Any, + a17: Any): Any = throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction18(id: Long) + extends ( + ( + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any) => Any) + with CometSharedFunctionRef { + override def apply( + a1: Any, + a2: Any, + a3: Any, + a4: Any, + a5: Any, + a6: Any, + a7: Any, + a8: Any, + a9: Any, + a10: Any, + a11: Any, + a12: Any, + a13: Any, + a14: Any, + a15: Any, + a16: Any, + a17: Any, + a18: Any): Any = throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction19(id: Long) + extends ( + ( + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any) => Any) + with CometSharedFunctionRef { + override def apply( + a1: Any, + a2: Any, + a3: Any, + a4: Any, + a5: Any, + a6: Any, + a7: Any, + a8: Any, + a9: Any, + a10: Any, + a11: Any, + a12: Any, + a13: Any, + a14: Any, + a15: Any, + a16: Any, + a17: Any, + a18: Any, + a19: Any): Any = throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction20(id: Long) + extends ( + ( + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any) => Any) + with CometSharedFunctionRef { + override def apply( + a1: Any, + a2: Any, + a3: Any, + a4: Any, + a5: Any, + a6: Any, + a7: Any, + a8: Any, + a9: Any, + a10: Any, + a11: Any, + a12: Any, + a13: Any, + a14: Any, + a15: Any, + a16: Any, + a17: Any, + a18: Any, + a19: Any, + a20: Any): Any = throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction21(id: Long) + extends ( + ( + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any) => Any) + with CometSharedFunctionRef { + override def apply( + a1: Any, + a2: Any, + a3: Any, + a4: Any, + a5: Any, + a6: Any, + a7: Any, + a8: Any, + a9: Any, + a10: Any, + a11: Any, + a12: Any, + a13: Any, + a14: Any, + a15: Any, + a16: Any, + a17: Any, + a18: Any, + a19: Any, + a20: Any, + a21: Any): Any = throw CometSharedFunctionRef.unresolved(id) +} + +final case class CometSharedFunction22(id: Long) + extends ( + ( + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any, + Any) => Any) + with CometSharedFunctionRef { + override def apply( + a1: Any, + a2: Any, + a3: Any, + a4: Any, + a5: Any, + a6: Any, + a7: Any, + a8: Any, + a9: Any, + a10: Any, + a11: Any, + a12: Any, + a13: Any, + a14: Any, + a15: Any, + a16: Any, + a17: Any, + a18: Any, + a19: Any, + a20: Any, + a21: Any, + a22: Any): Any = throw CometSharedFunctionRef.unresolved(id) +} +// scalastyle:on argcount diff --git a/spark/src/main/scala/org/apache/comet/udf/codegen/CometUdfFunctionPool.scala b/spark/src/main/scala/org/apache/comet/udf/codegen/CometUdfFunctionPool.scala new file mode 100644 index 00000000000..e84a0a10498 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/udf/codegen/CometUdfFunctionPool.scala @@ -0,0 +1,288 @@ +/* + * 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.comet.udf.codegen + +import java.io.{ByteArrayInputStream, ByteArrayOutputStream, DataInputStream, DataOutputStream} +import java.io.{ObjectOutputStream, OutputStream} +import java.nio.ByteBuffer +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.{AtomicLong, AtomicReference} + +import scala.collection.mutable +import scala.reflect.ClassTag + +import org.apache.spark.{SparkEnv, TaskContext} +import org.apache.spark.sql.catalyst.expressions.{Expression, ScalaUDF} + +import com.google.common.cache.{Cache, CacheBuilder} +import com.google.common.collect.MapMaker +import com.google.protobuf.{ByteString, CodedInputStream, CodedOutputStream, Message, WireFormat} +import com.google.protobuf.Descriptors.FieldDescriptor + +import org.apache.comet.serde.ExprOuterClass +import org.apache.comet.serde.OperatorOuterClass.Operator + +object CometUdfFunctionPool { + + val PoolFieldNumber: Int = 99 + + private val ExprMagic: Array[Byte] = "CMTUDFX1".getBytes("US-ASCII") + private val PoolMagic: Array[Byte] = "CMTUDFP1".getBytes("US-ASCII") + + private val nextId = new AtomicLong(0) + private val idsByFunction: java.util.concurrent.ConcurrentMap[AnyRef, java.lang.Long] = + new MapMaker().weakKeys().makeMap() + private val functionsById: java.util.concurrent.ConcurrentMap[java.lang.Long, AnyRef] = + new MapMaker().weakValues().makeMap() + private val serializable: java.util.concurrent.ConcurrentMap[AnyRef, java.lang.Boolean] = + new MapMaker().weakKeys().makeMap() + private val pools: Cache[Vector[Long], Array[Byte]] = + CacheBuilder.newBuilder().softValues().maximumSize(64).build() + + private val poolBuilds = new AtomicLong(0) + private val poolDeserializations = new AtomicLong(0) + + private val tasks = new ConcurrentHashMap[java.lang.Long, TaskPools]() + + def poolBuildCount: Long = poolBuilds.get() + + def poolDeserializationCount: Long = poolDeserializations.get() + + def functionsOf(expr: Expression): Seq[AnyRef] = { + val seen = java.util.Collections.newSetFromMap( + new java.util.IdentityHashMap[AnyRef, java.lang.Boolean]()) + expr.collect { case udf: ScalaUDF => udf.function }.filter(seen.add) + } + + def isSerializable(function: AnyRef): Boolean = + serializable + .computeIfAbsent( + function, + f => + try { + val out = new ObjectOutputStream(NullOutputStream) + out.writeObject(f) + out.close() + java.lang.Boolean.TRUE + } catch { + case scala.util.control.NonFatal(_) => java.lang.Boolean.FALSE + }) + .booleanValue() + + def serializeExpression(expr: Expression, functions: Seq[AnyRef]): Array[Byte] = { + val ids = new java.util.IdentityHashMap[AnyRef, java.lang.Long]() + functions.foreach(f => ids.put(f, idOf(f))) + val withRefs = expr.transform { + case udf: ScalaUDF if ids.containsKey(udf.function) => + udf.copy(function = CometSharedFunctionRef(ids.get(udf.function), udf.children.size)) + } + withHeader(ExprMagic, functions.map(f => ids.get(f).longValue()), withRefs) + } + + def isSharedExpression(bytes: Array[Byte]): Boolean = startsWith(bytes, ExprMagic) + + def deserializeExpression(bytes: Array[Byte], loader: ClassLoader): Expression = { + val in = new ByteArrayInputStream(bytes) + val ids = readIds(new DataInputStream(in), ExprMagic) + val functions = ids.map(id => id -> resolve(id, loader)).toMap + readObject[Expression](bytes, bytes.length - in.available(), loader).transform { + case udf: ScalaUDF if udf.function.isInstanceOf[CometSharedFunctionRef] => + udf.copy(function = functions(udf.function.asInstanceOf[CometSharedFunctionRef].id)) + } + } + + def attachPools(plan: Operator, planBytes: Array[Byte]): Array[Byte] = { + if (functionsById.isEmpty) return planBytes + val referenced = mutable.LinkedHashSet.empty[Long] + val pooled = mutable.HashSet.empty[Long] + collect(plan, referenced, pooled) + val missing = referenced.filterNot(pooled.contains).toVector.sorted + if (missing.isEmpty) return planBytes + val functions = missing.map(id => Option(functionsById.get(id))) + if (functions.exists(_.isEmpty)) return planBytes + val pool = pools.get( + missing, + () => { + poolBuilds.incrementAndGet() + withHeader(PoolMagic, missing, functions.map(_.get).toArray[AnyRef]) + }) + val size = planBytes.length + CodedOutputStream.computeByteArraySize(PoolFieldNumber, pool) + val result = java.util.Arrays.copyOf(planBytes, size) + val coded = CodedOutputStream.newInstance(result, planBytes.length, size - planBytes.length) + coded.writeByteArray(PoolFieldNumber, pool) + coded.checkNoSpaceLeft() + result + } + + def registerPlan(context: TaskContext, planBytes: Array[Byte]): Unit = { + if (context == null) return + val found = mutable.ArrayBuffer.empty[Array[Byte]] + scanPools(CodedInputStream.newInstance(planBytes), found) + if (found.isEmpty) return + val taskId = java.lang.Long.valueOf(context.taskAttemptId()) + val taskPools = tasks.computeIfAbsent( + taskId, + _ => { + val created = new TaskPools + context.addTaskCompletionListener[Unit](_ => tasks.remove(taskId)) + created + }) + found.foreach(taskPools.add) + } + + private def idOf(function: AnyRef): Long = { + val id = idsByFunction.computeIfAbsent(function, _ => nextId.incrementAndGet()) + functionsById.putIfAbsent(id, function) + id + } + + private def resolve(id: Long, loader: ClassLoader): AnyRef = { + val fromTask = Option(TaskContext.get()) + .flatMap(context => Option(tasks.get(context.taskAttemptId()))) + .flatMap(_.resolve(id, loader)) + fromTask + .orElse(Option(functionsById.get(id))) + .getOrElse(throw new IllegalStateException( + s"CometScalaUDFCodegen: UDF function $id is not in any function pool of this task")) + } + + private final class TaskPools { + private val slots = mutable.ArrayBuffer.empty[PoolSlot] + + def add(bytes: Array[Byte]): Unit = synchronized { + val ids = readIds(new DataInputStream(new ByteArrayInputStream(bytes)), PoolMagic) + if (!slots.exists(_.ids == ids)) slots += new PoolSlot(ids, bytes) + } + + def resolve(id: Long, loader: ClassLoader): Option[AnyRef] = { + val slot = synchronized(slots.find(_.ids.contains(id))) + slot.map(_.function(id, loader)) + } + } + + private final class PoolSlot(val ids: Vector[Long], bytes: Array[Byte]) { + private val functions = new AtomicReference[Map[Long, AnyRef]]() + + def function(id: Long, loader: ClassLoader): AnyRef = synchronized { + if (functions.get() == null) { + poolDeserializations.incrementAndGet() + val in = new ByteArrayInputStream(bytes) + readIds(new DataInputStream(in), PoolMagic) + val values = readObject[Array[AnyRef]](bytes, bytes.length - in.available(), loader) + functions.set(ids.zip(values).toMap) + } + functions.get()(id) + } + } + + private def readIds(in: DataInputStream, magic: Array[Byte]): Vector[Long] = { + val head = new Array[Byte](magic.length) + in.readFully(head) + require(java.util.Arrays.equals(head, magic), "not a Comet UDF function pool payload") + val count = in.readInt() + Vector.fill(count)(in.readLong()) + } + + private def startsWith(bytes: Array[Byte], magic: Array[Byte]): Boolean = + bytes.length >= magic.length && + java.util.Arrays.equals(java.util.Arrays.copyOf(bytes, magic.length), magic) + + private def collect( + message: Message, + referenced: mutable.LinkedHashSet[Long], + pooled: mutable.HashSet[Long]): Unit = { + message match { + case op: Operator => + Option(op.getUnknownFields.getField(PoolFieldNumber)).foreach { field => + field.getLengthDelimitedList.forEach { pool => + pooled ++= readIds(new DataInputStream(pool.newInput()), PoolMagic) + } + } + case udf: ExprOuterClass.JvmScalarUdf + if udf.getClassName == classOf[CometScalaUDFCodegen].getName && + udf.getArgsCount > 0 && udf.getArgs(0).hasLiteral => + val payload = udf.getArgs(0).getLiteral.getBytesVal + if (payload.startsWith(ByteString.copyFrom(ExprMagic))) { + referenced ++= readIds(new DataInputStream(payload.newInput()), ExprMagic) + } + case _ => + } + val fields = message.getDescriptorForType.getFields.iterator() + while (fields.hasNext) { + val field = fields.next() + if (field.getJavaType == FieldDescriptor.JavaType.MESSAGE && !field.isMapField) { + if (field.isRepeated) { + val count = message.getRepeatedFieldCount(field) + var i = 0 + while (i < count) { + collect(message.getRepeatedField(field, i).asInstanceOf[Message], referenced, pooled) + i += 1 + } + } else if (message.hasField(field)) { + collect(message.getField(field).asInstanceOf[Message], referenced, pooled) + } + } + } + } + + private def scanPools(in: CodedInputStream, found: mutable.ArrayBuffer[Array[Byte]]): Unit = { + in.setSizeLimit(Int.MaxValue) + var tag = in.readTag() + while (tag != 0) { + WireFormat.getTagFieldNumber(tag) match { + case PoolFieldNumber + if WireFormat.getTagWireType(tag) == WireFormat.WIRETYPE_LENGTH_DELIMITED => + found += in.readByteArray() + case Operator.CHILDREN_FIELD_NUMBER + if WireFormat.getTagWireType(tag) == WireFormat.WIRETYPE_LENGTH_DELIMITED => + val limit = in.pushLimit(in.readRawVarint32()) + scanPools(in, found) + in.popLimit(limit) + case _ => + in.skipField(tag) + } + tag = in.readTag() + } + } + + private object NullOutputStream extends OutputStream { + override def write(b: Int): Unit = () + override def write(b: Array[Byte], off: Int, len: Int): Unit = () + } + + private def withHeader(magic: Array[Byte], ids: Seq[Long], obj: AnyRef): Array[Byte] = { + val header = new ByteArrayOutputStream() + val out = new DataOutputStream(header) + out.write(magic) + out.writeInt(ids.size) + ids.foreach(id => out.writeLong(id)) + out.flush() + val buffer = SparkEnv.get.closureSerializer.newInstance().serialize(obj) + val result = new Array[Byte](header.size() + buffer.remaining()) + System.arraycopy(header.toByteArray, 0, result, 0, header.size()) + buffer.get(result, header.size(), buffer.remaining()) + result + } + + private def readObject[T: ClassTag](bytes: Array[Byte], offset: Int, loader: ClassLoader): T = + SparkEnv.get.closureSerializer + .newInstance() + .deserialize[T](ByteBuffer.wrap(bytes, offset, bytes.length - offset), loader) +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 8a93cd31730..309abfbf1fe 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -65,6 +65,7 @@ import org.apache.comet.serde.OperatorOuterClass.{AggregateMode => CometAggregat import org.apache.comet.serde.QueryPlanSerde import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, isStringCollationType, supportedSortType} import org.apache.comet.serde.operator.CometSink +import org.apache.comet.udf.codegen.CometUdfFunctionPool /** * Trait for injecting per-partition planning data into operator nodes. @@ -732,7 +733,7 @@ object CometExec { val codedOutput = CodedOutputStream.newInstance(bytes) nativePlan.writeTo(codedOutput) codedOutput.checkNoSpaceLeft() - bytes + CometUdfFunctionPool.attachPools(nativePlan, bytes) } def getCometIterator( diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 0b9a19864ea..5f7f35d6abd 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -29,7 +29,7 @@ import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.api.java.UDF1 import org.apache.spark.sql.catalyst.expressions.{Add, Alias, AttributeReference, BoundReference, Cast, CreateArray, CreateMap, CreateNamedStruct, Expression, Hypot, Literal, MapConcat} import org.apache.spark.sql.catalyst.expressions.objects.Invoke -import org.apache.spark.sql.comet.CometProjectExec +import org.apache.spark.sql.comet.{CometNativeExec, CometProjectExec} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ @@ -40,7 +40,7 @@ import org.apache.comet.codegen.CometBatchKernelCodegen import org.apache.comet.codegen.CometBatchKernelCodegen.ArrowColumnSpec import org.apache.comet.serde.{CometScalaUDF, QueryPlanSerde} import org.apache.comet.serde.ExprOuterClass.Expr.ExprStructCase -import org.apache.comet.udf.codegen.CometScalaUDFCodegen +import org.apache.comet.udf.codegen.{CometScalaUDFCodegen, CometUdfFunctionPool} import org.apache.comet.vector.CometVector /** @@ -2347,6 +2347,35 @@ class CometCodegenSuite Invoke(target, "twice", StringType, Seq(Literal(UTF8String.fromString("ab"), StringType))) assert(runKernel(folded, 1)(_.getUTF8String(0).toString) === "abab") } + + test("UDFs that capture one large object ship it once per serialized plan") { + val captured = Array.tabulate(1 << 19)(_.toLong) + val capturedBytes = captured.length * 8 + spark.udf.register("sharedA", (id: Long) => captured((id % captured.length).toInt) + 1) + spark.udf.register("sharedB", (id: Long) => captured((id % captured.length).toInt) * 2) + spark.udf.register("sharedC", (id: Long) => captured.length.toLong - id) + val query = "SELECT id, sharedA(id), sharedB(id), sharedC(id), sharedA(id + 1), " + + "sharedB(id + 2), sharedC(id + 3), sharedA(id + 4), sharedB(id + 5) FROM t" + withTable("t") { + spark.range(0, 2000, 1, 4).write.saveAsTable("t") + val builds = CometUdfFunctionPool.poolBuildCount + val deserializations = CometUdfFunctionPool.poolDeserializationCount + val df = sql(query) + checkSparkAnswerAndOperator(df) + val planBytes = collect(df.queryExecution.executedPlan) { + case op: CometNativeExec if op.serializedPlanOpt.isDefined => + op.serializedPlanOpt.plan.get.length + } + assert(planBytes.nonEmpty) + assert(planBytes.max > capturedBytes, s"plan sizes $planBytes") + assert(planBytes.max < 2 * capturedBytes, s"plan sizes $planBytes") + assert(CometUdfFunctionPool.poolBuildCount == builds + 1) + assert(CometUdfFunctionPool.poolDeserializationCount > deserializations) + + checkSparkAnswerAndOperator(sql(query)) + assert(CometUdfFunctionPool.poolBuildCount == builds + 1) + } + } } /** From ab179fa2ea83a541cccc42911a2d9ea5d15e1320 Mon Sep 17 00:00:00 2001 From: msaf Date: Sat, 17 Oct 2026 12:56:01 +0100 Subject: [PATCH 2/2] fix: cache a native plan's UDF function pools without Guava's Cache, which the shaded jar cannot load Co-Authored-By: Claude Opus 5.5 (1M context) --- .../udf/codegen/CometUdfFunctionPool.scala | 21 ++++++++++++------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/udf/codegen/CometUdfFunctionPool.scala b/spark/src/main/scala/org/apache/comet/udf/codegen/CometUdfFunctionPool.scala index e84a0a10498..f9af430f8e0 100644 --- a/spark/src/main/scala/org/apache/comet/udf/codegen/CometUdfFunctionPool.scala +++ b/spark/src/main/scala/org/apache/comet/udf/codegen/CometUdfFunctionPool.scala @@ -21,6 +21,7 @@ package org.apache.comet.udf.codegen import java.io.{ByteArrayInputStream, ByteArrayOutputStream, DataInputStream, DataOutputStream} import java.io.{ObjectOutputStream, OutputStream} +import java.lang.ref.SoftReference import java.nio.ByteBuffer import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.atomic.{AtomicLong, AtomicReference} @@ -31,7 +32,6 @@ import scala.reflect.ClassTag import org.apache.spark.{SparkEnv, TaskContext} import org.apache.spark.sql.catalyst.expressions.{Expression, ScalaUDF} -import com.google.common.cache.{Cache, CacheBuilder} import com.google.common.collect.MapMaker import com.google.protobuf.{ByteString, CodedInputStream, CodedOutputStream, Message, WireFormat} import com.google.protobuf.Descriptors.FieldDescriptor @@ -53,8 +53,12 @@ object CometUdfFunctionPool { new MapMaker().weakValues().makeMap() private val serializable: java.util.concurrent.ConcurrentMap[AnyRef, java.lang.Boolean] = new MapMaker().weakKeys().makeMap() - private val pools: Cache[Vector[Long], Array[Byte]] = - CacheBuilder.newBuilder().softValues().maximumSize(64).build() + private val pools = + new java.util.LinkedHashMap[Vector[Long], SoftReference[Array[Byte]]](16, 0.75f, true) { + override def removeEldestEntry( + eldest: java.util.Map.Entry[Vector[Long], SoftReference[Array[Byte]]]): Boolean = + size() > 64 + } private val poolBuilds = new AtomicLong(0) private val poolDeserializations = new AtomicLong(0) @@ -117,12 +121,13 @@ object CometUdfFunctionPool { if (missing.isEmpty) return planBytes val functions = missing.map(id => Option(functionsById.get(id))) if (functions.exists(_.isEmpty)) return planBytes - val pool = pools.get( - missing, - () => { + val pool = + pools.synchronized(Option(pools.get(missing)).flatMap(r => Option(r.get))).getOrElse { poolBuilds.incrementAndGet() - withHeader(PoolMagic, missing, functions.map(_.get).toArray[AnyRef]) - }) + val built = withHeader(PoolMagic, missing, functions.map(_.get).toArray[AnyRef]) + pools.synchronized(pools.put(missing, new SoftReference(built))) + built + } val size = planBytes.length + CodedOutputStream.computeByteArraySize(PoolFieldNumber, pool) val result = java.util.Arrays.copyOf(planBytes, size) val coded = CodedOutputStream.newInstance(result, planBytes.length, size - planBytes.length)