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..f9af430f8e0 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/udf/codegen/CometUdfFunctionPool.scala @@ -0,0 +1,293 @@ +/* + * 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.lang.ref.SoftReference +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.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 = + 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) + + 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.synchronized(Option(pools.get(missing)).flatMap(r => Option(r.get))).getOrElse { + poolBuilds.incrementAndGet() + 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) + 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) + } + } } /**