fix: ship lambda expression in proto instead of via process-local registry

The previous CometLambdaRegistry approach registered the ArrayExists
expression in a JVM-local ConcurrentHashMap on the driver and looked it
up by UUID on the executor. In a real cluster the executor JVM never
sees the driver's registry and the lookup fails. Local-mode tests
masked this because driver and executor share the same JVM.

Serialize the ArrayExists Catalyst expression to bytes in the serde
layer and embed it in the proto as a BinaryType literal arg. The UDF
deserializes on each invocation. CometLambdaRegistry is no longer
referenced and is removed.

Author: andygrove

spark/src/main/scala/org/apache/comet/serde/arrays.scala

@@ -31,7 +31,6 @@ import org.apache.comet.CometConf
 import org.apache.comet.CometSparkSessionExtensions.withInfo
 import org.apache.comet.serde.QueryPlanSerde._
 import org.apache.comet.shims.{CometExprShim, CometTypeShim}
-import org.apache.comet.udf.CometLambdaRegistry
 
 object CometArrayRemove
     extends CometExpressionSerde[ArrayRemove]
@@ -852,25 +851,31 @@ object CometArrayExists extends CometExpressionSerde[ArrayExists] {
       return None
     }
 
-    val registryKey = CometLambdaRegistry.register(expr)
-    val keyLiteral = Literal(registryKey)
-    val keyProto = exprToProtoInternal(keyLiteral, inputs, binding)
-    if (keyProto.isEmpty) {
-      CometLambdaRegistry.remove(registryKey)
-      withInfo(expr, "Failed to serialize registry key literal")
+    // The lambda body is evaluated on the executor inside ArrayExistsUDF, so the entire
+    // ArrayExists Catalyst expression must travel with the plan. The driver-side
+    // CometLambdaRegistry approach broke under real distributed execution because the
+    // executor JVM doesn't share the driver's map. Serializing the expression as bytes
+    // and shipping it through the proto as a Literal arg keeps the lambda self-contained.
+    val baos = new java.io.ByteArrayOutputStream()
+    val oos = new java.io.ObjectOutputStream(baos)
+    try {
+      oos.writeObject(expr)
+    } finally {
+      oos.close()
+    }
+    val payloadProto = exprToProtoInternal(Literal(baos.toByteArray, BinaryType), inputs, binding)
+    if (payloadProto.isEmpty) {
+      withInfo(expr, "Failed to serialize lambda expression payload")
       return None
     }
 
-    val returnType = serializeDataType(BooleanType).getOrElse {
-      CometLambdaRegistry.remove(registryKey)
-      return None
-    }
+    val returnType = serializeDataType(BooleanType).getOrElse(return None)
 
     val udfBuilder = ExprOuterClass.JvmScalarUdf
       .newBuilder()
       .setClassName("org.apache.comet.udf.ArrayExistsUDF")
       .addArgs(arrayProto.get)
-      .addArgs(keyProto.get)
+      .addArgs(payloadProto.get)
       .setReturnType(returnType)
       .setReturnNullable(expr.nullable)
 

spark/src/main/scala/org/apache/comet/udf/ArrayExistsUDF.scala

@@ -19,8 +19,6 @@
 
 package org.apache.comet.udf
 
-import java.nio.charset.StandardCharsets
-
 import org.apache.arrow.vector._
 import org.apache.arrow.vector.complex.ListVector
 import org.apache.spark.sql.catalyst.expressions.{ArrayExists, LambdaFunction, NamedLambdaVariable}
@@ -34,7 +32,9 @@ import org.apache.comet.CometArrowAllocator
  *
  * Inputs:
  *   - inputs(0): ListVector (the array column)
- *   - inputs(1): VarCharVector length-1 scalar (registry key for the lambda expression)
+ *   - inputs(1): VarBinaryVector length-1 scalar containing the Java-serialized [[ArrayExists]]
+ *     Catalyst expression. Shipping the expression in the proto avoids the driver-vs-executor
+ *     mismatch a process-local registry would suffer.
  *
  * Output: BitVector (nullable boolean), same length as the input array vector.
  *
@@ -48,13 +48,17 @@ class ArrayExistsUDF extends CometUDF {
   override def evaluate(inputs: Array[ValueVector], numRows: Int): ValueVector = {
     require(inputs.length == 2, s"ArrayExistsUDF expects 2 inputs, got ${inputs.length}")
     val listVec = inputs(0).asInstanceOf[ListVector]
-    val keyVec = inputs(1).asInstanceOf[VarCharVector]
+    val payloadVec = inputs(1).asInstanceOf[VarBinaryVector]
     require(
-      keyVec.getValueCount >= 1 && !keyVec.isNull(0),
-      "ArrayExistsUDF requires a non-null scalar registry key")
+      payloadVec.getValueCount >= 1 && !payloadVec.isNull(0),
+      "ArrayExistsUDF requires a non-null scalar payload")
 
-    val registryKey = new String(keyVec.get(0), StandardCharsets.UTF_8)
-    val arrayExistsExpr = CometLambdaRegistry.get(registryKey).asInstanceOf[ArrayExists]
+    val payloadBytes = payloadVec.get(0)
+    val bais = new java.io.ByteArrayInputStream(payloadBytes)
+    val ois = new java.io.ObjectInputStream(bais)
+    val arrayExistsExpr =
+      try ois.readObject().asInstanceOf[ArrayExists]
+      finally ois.close()
 
     val LambdaFunction(_, Seq(elementVar: NamedLambdaVariable), _) = arrayExistsExpr.function
     val body = arrayExistsExpr.functionForEval