[DRAFT][SPARK-51705][CONNECT][SCALA] Support SparkSession.broadcast() for Scala UDFs over Spark Connect

python/pyspark/sql/connect/proto/commands_pb2.pyi

@@ -56,6 +56,33 @@ else:
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
+class _BroadcastValueType:
+    ValueType = typing.NewType("ValueType", builtins.int)
+    V: typing_extensions.TypeAlias = ValueType
+
+class _BroadcastValueTypeEnumTypeWrapper(
+    google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_BroadcastValueType.ValueType],
+    builtins.type,
+):  # noqa: F821
+    DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+    BROADCAST_VALUE_TYPE_UNSPECIFIED: _BroadcastValueType.ValueType  # 0
+    """Unset -> treated as PYTHON for backward compatibility."""
+    BROADCAST_VALUE_TYPE_PYTHON: _BroadcastValueType.ValueType  # 1
+    """Wrap the staged bytes as a PythonBroadcast file; the Python worker reads it. Broadcast[PythonBroadcast]."""
+    BROADCAST_VALUE_TYPE_JVM: _BroadcastValueType.ValueType  # 2
+    """Java-deserialize the bytes into a JVM value T and sc.broadcast(T). Broadcast[T] for Scala UDFs."""
+
+class BroadcastValueType(_BroadcastValueType, metaclass=_BroadcastValueTypeEnumTypeWrapper):
+    """(SPARK-51705) Discriminates how CreateBroadcastCommand materializes the broadcast value."""
+
+BROADCAST_VALUE_TYPE_UNSPECIFIED: BroadcastValueType.ValueType  # 0
+"""Unset -> treated as PYTHON for backward compatibility."""
+BROADCAST_VALUE_TYPE_PYTHON: BroadcastValueType.ValueType  # 1
+"""Wrap the staged bytes as a PythonBroadcast file; the Python worker reads it. Broadcast[PythonBroadcast]."""
+BROADCAST_VALUE_TYPE_JVM: BroadcastValueType.ValueType  # 2
+"""Java-deserialize the bytes into a JVM value T and sc.broadcast(T). Broadcast[T] for Scala UDFs."""
+global___BroadcastValueType = BroadcastValueType
+
 class _StreamingQueryEventType:
     ValueType = typing.NewType("ValueType", builtins.int)
     V: typing_extensions.TypeAlias = ValueType
@@ -344,20 +371,33 @@ class CreateBroadcastCommand(google.protobuf.message.Message):
 
     ARTIFACT_HASH_FIELD_NUMBER: builtins.int
     SIZE_BYTES_FIELD_NUMBER: builtins.int
+    VALUE_TYPE_FIELD_NUMBER: builtins.int
     artifact_hash: builtins.str
-    """(Required) sha256 hash returned by client.cache_artifact(cloudpickle(value))."""
+    """(Required) sha256 hash of the previously-uploaded value bytes in the cache/ artifact channel.
+    For PYTHON the bytes are cloudpickle(value); for JVM they are SparkSerDeUtils.serialize(value).
+    """
     size_bytes: builtins.int
     """(Optional) Uncompressed byte size, used to enforce the BROADCAST_VALUE_TOO_LARGE quota."""
+    value_type: global___BroadcastValueType.ValueType
+    """(Optional) (SPARK-51705) How the server should materialize the value from the uploaded bytes.
+    Defaults to PYTHON when unset for backward compatibility with the Python-only v1.
+    """
     def __init__(
         self,
         *,
         artifact_hash: builtins.str = ...,
         size_bytes: builtins.int = ...,
+        value_type: global___BroadcastValueType.ValueType = ...,
     ) -> None: ...
     def ClearField(
         self,
         field_name: typing_extensions.Literal[
-            "artifact_hash", b"artifact_hash", "size_bytes", b"size_bytes"
+            "artifact_hash",
+            b"artifact_hash",
+            "size_bytes",
+            b"size_bytes",
+            "value_type",
+            b"value_type",
         ],
     ) -> None: ...
 

python/pyspark/sql/connect/proto/expressions_pb2.pyi

@@ -1891,6 +1891,7 @@ class ScalarScalaUDF(google.protobuf.message.Message):
     OUTPUTTYPE_FIELD_NUMBER: builtins.int
     NULLABLE_FIELD_NUMBER: builtins.int
     AGGREGATE_FIELD_NUMBER: builtins.int
+    BROADCAST_IDS_FIELD_NUMBER: builtins.int
     payload: builtins.bytes
     """(Required) Serialized JVM object containing UDF definition, input encoders and output encoder"""
     @property
@@ -1907,6 +1908,17 @@ class ScalarScalaUDF(google.protobuf.message.Message):
     """(Required) True if the UDF can return null value"""
     aggregate: builtins.bool
     """(Required) Indicate if the UDF is an aggregate function"""
+    @property
+    def broadcast_ids(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """(Optional) (SPARK-51705) Server-assigned ids of broadcast variables captured by this UDF's
+        closure. The captured broadcasts serialize as id-only tokens (ConnectBroadcast.writeReplace);
+        this out-of-band id list lets SparkConnectPlanner bind the per-session SessionHolder broadcast
+        registry (via a thread-local) before deserializing the payload, so ConnectBroadcastRef
+        .readResolve can swap each id for the real driver-side Broadcast[T]. An id that is unknown to
+        this session fails loudly with BROADCAST_NOT_FOUND.
+        """
     def __init__(
         self,
         *,
@@ -1916,6 +1928,7 @@ class ScalarScalaUDF(google.protobuf.message.Message):
         outputType: pyspark.sql.connect.proto.types_pb2.DataType | None = ...,
         nullable: builtins.bool = ...,
         aggregate: builtins.bool = ...,
+        broadcast_ids: collections.abc.Iterable[builtins.int] | None = ...,
     ) -> None: ...
     def HasField(
         self, field_name: typing_extensions.Literal["outputType", b"outputType"]
@@ -1925,6 +1938,8 @@ class ScalarScalaUDF(google.protobuf.message.Message):
         field_name: typing_extensions.Literal[
             "aggregate",
             b"aggregate",
+            "broadcast_ids",
+            b"broadcast_ids",
             "inputTypes",
             b"inputTypes",
             "nullable",

sql/connect/common/src/main/protobuf/spark/connect/commands.proto

@@ -68,10 +68,24 @@ message Command {
 // sends this command with the returned hash. The server materializes a PythonBroadcast on the
 // live driver SparkContext and returns a CreateBroadcastResult with the driver-side broadcast id.
 message CreateBroadcastCommand {
-  // (Required) sha256 hash returned by client.cache_artifact(cloudpickle(value)).
+  // (Required) sha256 hash of the previously-uploaded value bytes in the cache/ artifact channel.
+  // For PYTHON the bytes are cloudpickle(value); for JVM they are SparkSerDeUtils.serialize(value).
   string artifact_hash = 1;
   // (Optional) Uncompressed byte size, used to enforce the BROADCAST_VALUE_TOO_LARGE quota.
   int64 size_bytes = 2;
+  // (Optional) (SPARK-51705) How the server should materialize the value from the uploaded bytes.
+  // Defaults to PYTHON when unset for backward compatibility with the Python-only v1.
+  BroadcastValueType value_type = 3;
+}
+
+// (SPARK-51705) Discriminates how CreateBroadcastCommand materializes the broadcast value.
+enum BroadcastValueType {
+  // Unset -> treated as PYTHON for backward compatibility.
+  BROADCAST_VALUE_TYPE_UNSPECIFIED = 0;
+  // Wrap the staged bytes as a PythonBroadcast file; the Python worker reads it. Broadcast[PythonBroadcast].
+  BROADCAST_VALUE_TYPE_PYTHON = 1;
+  // Java-deserialize the bytes into a JVM value T and sc.broadcast(T). Broadcast[T] for Scala UDFs.
+  BROADCAST_VALUE_TYPE_JVM = 2;
 }
 
 // (SPARK-51705) Release a broadcast variable created over Spark Connect.

sql/connect/common/src/main/protobuf/spark/connect/expressions.proto

@@ -493,6 +493,13 @@ message ScalarScalaUDF {
   bool nullable = 4;
   // (Required) Indicate if the UDF is an aggregate function
   bool aggregate = 5;
+  // (Optional) (SPARK-51705) Server-assigned ids of broadcast variables captured by this UDF's
+  // closure. The captured broadcasts serialize as id-only tokens (ConnectBroadcast.writeReplace);
+  // this out-of-band id list lets SparkConnectPlanner bind the per-session SessionHolder broadcast
+  // registry (via a thread-local) before deserializing the payload, so ConnectBroadcastRef
+  // .readResolve can swap each id for the real driver-side Broadcast[T]. An id that is unknown to
+  // this session fails loudly with BROADCAST_NOT_FOUND.
+  repeated int64 broadcast_ids = 6;
 }
 
 message JavaUDF {

sql/connect/common/src/main/scala/org/apache/spark/sql/connect/ConnectBroadcast.scala

@@ -0,0 +1,162 @@
+/*
+ * 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.connect
+
+import scala.collection.mutable
+import scala.reflect.ClassTag
+
+import org.apache.spark.broadcast.Broadcast
+
+/**
+ * (SPARK-51705) A client-side stand-in for a driver-side [[Broadcast]] that a Spark Connect Scala
+ * client can capture inside a UDF closure even though the client has no SparkContext (so it can
+ * never call `sc.broadcast(v)` to obtain a real [[org.apache.spark.broadcast.TorrentBroadcast]]).
+ *
+ * This is the Scala analogue of the Python `ConnectBroadcast` proxy
+ * (`python/pyspark/sql/connect/broadcast.py`). It solves the same gap that motivated
+ * `SparkSession.broadcast()` for Python: the Connect client is JVM-less with respect to the
+ * cluster, so the user cannot construct the `Broadcast[T]` object that classic Spark expects them
+ * to capture.
+ *
+ * End to end (Scala): `SparkSession.broadcast(v)` serializes `v`, uploads it through the cache
+ * artifact channel, sends a `CreateBroadcastCommand(value_type = JVM)`, and returns a
+ * `ConnectBroadcast[T]` holding the server-assigned `Broadcast.id` (plus the local value for
+ * driver-side `.value` reads). When the user captures it in a UDF, [[writeReplace]] substitutes
+ * an id-only [[ConnectBroadcastRef]] token into the serialized closure; the captured value itself
+ * is never written to the wire (it already travels once, out of band, via the cache artifact). On
+ * the server, [[ConnectBroadcastRef.readResolve]] swaps the token for the real driver-side
+ * `Broadcast[T]` while the closure is deserialized (see `SparkConnectPlanner.unpackScalaUDF`).
+ *
+ * @param bid
+ *   the server-assigned driver-side `Broadcast.id`
+ * @param value_
+ *   the local value, retained only so client-side `.value` reads work (parity with classic
+ *   driver-side reads). Marked `@transient` so it is never serialized even if [[writeReplace]]
+ *   were somehow bypassed.
+ */
+private[sql] class ConnectBroadcast[T: ClassTag](
+    bid: Long,
+    @transient private val value_ : T,
+    @transient private val unpersistFn: (Long, Boolean, Boolean) => Unit)
+    extends Broadcast[T](bid) {
+
+  override protected def getValue(): T = value_
+
+  override protected def doUnpersist(blocking: Boolean): Unit =
+    unpersistFn(id, blocking, false)
+
+  override protected def doDestroy(blocking: Boolean): Unit =
+    unpersistFn(id, blocking, true)
+
+  /**
+   * When the enclosing UDF closure is Java-serialized, substitute an id-only token for this
+   * object. Mirrors the Python `ConnectBroadcast.__reduce__ -> (_from_id, (bid,))` contract.
+   * Because this returns the token, neither `value_` (already `@transient`) nor any
+   * SparkContext-bound state is ever serialized -- which is essential, since none exists on the
+   * client. Recording the id here (serialization is what proves the broadcast was captured) lets
+   * `UdfToProtoUtils.toProto` drain the exact set of captured ids into
+   * `ScalarScalaUDF.broadcast_ids`.
+   */
+  private def writeReplace(): AnyRef = {
+    ConnectBroadcastCapture.record(id)
+    new ConnectBroadcastRef(id)
+  }
+}
+
+/**
+ * The wire token that stands in for a [[ConnectBroadcast]] inside a serialized Scala UDF closure.
+ * Carries only the broadcast id.
+ *
+ * [[readResolve]] runs during closure deserialization. On the server, `SparkConnectPlanner`
+ * installs the per-session broadcast registry into [[ConnectBroadcastResolver]] immediately
+ * before deserializing, so this resolves the id to the real driver-side `Broadcast[_]`. When no
+ * registry is bound (for example the client-side `checkDeserializable` round-trip that
+ * `UdfToProtoUtils.toUdfPacketBytes` performs to validate the closure), it returns a lightweight
+ * unresolved placeholder rather than failing -- that placeholder is only ever produced during the
+ * client's own validation round-trip and is never executed.
+ */
+private[sql] class ConnectBroadcastRef(val id: Long) extends Serializable {
+  private def readResolve(): AnyRef =
+    ConnectBroadcastResolver.resolve(id).getOrElse(new UnresolvedConnectBroadcast(id))
+}
+
+/**
+ * Placeholder produced by [[ConnectBroadcastRef.readResolve]] when no registry is bound to the
+ * deserializing thread (client-side validation round-trip only). It intentionally throws if its
+ * value is ever read, so a misuse cannot silently return wrong data.
+ */
+private[sql] class UnresolvedConnectBroadcast(bid: Long) extends Broadcast[Any](bid) {
+  private def fail(): Nothing = throw new IllegalStateException(
+    s"ConnectBroadcast($bid) was not resolved to a driver-side broadcast. This placeholder is " +
+      "only expected during client-side closure validation and must never be executed.")
+  override protected def getValue(): Any = fail()
+  override protected def doUnpersist(blocking: Boolean): Unit = fail()
+  override protected def doDestroy(blocking: Boolean): Unit = fail()
+}
+
+/**
+ * Thread-local bridge that lets [[ConnectBroadcastRef.readResolve]] (which receives no context
+ * from `ObjectInputStream`) reach the per-session broadcast registry that only the server's
+ * `SparkConnectPlanner` holds.
+ *
+ * This type lives in `sql/connect/common` so both the client (which defines the token) and the
+ * server (which depends on common) see the identical class -- required for Java deserialization
+ * to bind the token. A thread-local is safe because `unpackScalaUDF` deserializes synchronously
+ * on the request-handling thread; the value is always cleared in a `finally`.
+ */
+private[sql] object ConnectBroadcastResolver {
+  private val bound = new ThreadLocal[Map[Long, Broadcast[_]]]()
+
+  def withRegistry[R](registry: Map[Long, Broadcast[_]])(body: => R): R = {
+    if (registry.isEmpty) {
+      // Nothing to resolve; avoid touching the thread-local at all.
+      body
+    } else {
+      val prev = bound.get()
+      bound.set(registry)
+      try body
+      finally {
+        if (prev == null) bound.remove() else bound.set(prev)
+      }
+    }
+  }
+
+  def resolve(id: Long): Option[Broadcast[_]] =
+    Option(bound.get()).flatMap(_.get(id))
+}
+
+/**
+ * Client-side thread-local that records the ids of [[ConnectBroadcast]]s captured while building
+ * a single UDF proto. `UdfToProtoUtils.toProto` drains this into `ScalarScalaUDF.broadcast_ids`
+ * right after serializing the closure (serialization is what triggers
+ * [[ConnectBroadcast.writeReplace]]), mirroring the Python `PythonUDF.to_plan` drain of its
+ * `threading.local` registry. Drain-then- clear on the same plan-build thread.
+ */
+private[sql] object ConnectBroadcastCapture {
+  private val captured = new ThreadLocal[mutable.LinkedHashSet[Long]] {
+    override def initialValue(): mutable.LinkedHashSet[Long] = mutable.LinkedHashSet.empty
+  }
+
+  def record(id: Long): Unit = captured.get() += id
+
+  /** Return the captured ids and clear the thread-local. */
+  def drain(): Seq[Long] = {
+    val ids = captured.get().toSeq
+    captured.remove()
+    ids
+  }
+}

sql/connect/common/src/main/scala/org/apache/spark/sql/connect/SparkSession.scala

@@ -38,6 +38,7 @@ import org.apache.arrow.memory.RootAllocator
 import org.apache.spark.{SparkConf, SparkContext}
 import org.apache.spark.annotation.{DeveloperApi, Experimental, Since}
 import org.apache.spark.api.java.JavaRDD
+import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.connect.proto
 import org.apache.spark.connect.proto.ExecutePlanResponse
 import org.apache.spark.connect.proto.ExecutePlanResponse.ObservedMetrics
@@ -59,6 +60,7 @@ import org.apache.spark.sql.sources.BaseRelation
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.util.{CloseableIterator, ExecutionListenerManager}
 import org.apache.spark.util.ArrayImplicits._
+import org.apache.spark.util.SparkSerDeUtils
 
 /**
  * The entry point to programming Spark with the Dataset and DataFrame API.
@@ -106,6 +108,42 @@ class SparkSession private[sql] (
   override def sparkContext: SparkContext =
     throw ConnectClientUnsupportedErrors.sparkContext()
 
+  /**
+   * (SPARK-51705) Create a broadcast variable usable inside a Scala UDF on Spark Connect.
+   *
+   * The Connect client has no `SparkContext`, so `sparkContext.broadcast(...)` is unavailable.
+   * This serializes `value`, uploads it once through the artifact cache channel, asks the server
+   * to materialize a driver-side `Broadcast[T]`, and returns a [[ConnectBroadcast]] proxy
+   * carrying the server-assigned id (plus the local value for driver-side `.value` reads). When
+   * the proxy is captured inside a UDF closure it serializes as an id-only token that the server
+   * swaps for the real broadcast; the executor path is then identical to classic Spark.
+   */
+  def broadcast[T: scala.reflect.ClassTag](value: T): Broadcast[T] = {
+    val bytes = SparkSerDeUtils.serialize(value)
+    val hash = client.artifactManager.cacheArtifact(bytes)
+    val command = newCommand { builder =>
+      builder.getCreateBroadcastCommandBuilder
+        .setArtifactHash(hash)
+        .setSizeBytes(bytes.length.toLong)
+        .setValueType(proto.BroadcastValueType.BROADCAST_VALUE_TYPE_JVM)
+    }
+    val response = execute(command)
+      .find(_.hasCreateBroadcastResult)
+      .getOrElse(throw new RuntimeException("CreateBroadcastResult must be present"))
+    val broadcastId = response.getCreateBroadcastResult.getBroadcastId
+    new ConnectBroadcast[T](broadcastId, value, unpersistBroadcast)
+  }
+
+  /** Route a [[ConnectBroadcast]] unpersist/destroy through an UnpersistBroadcastCommand RPC. */
+  private[connect] def unpersistBroadcast(id: Long, blocking: Boolean, destroy: Boolean): Unit = {
+    execute(newCommand { builder =>
+      builder.getUnpersistBroadcastCommandBuilder
+        .setBroadcastId(id)
+        .setBlocking(blocking)
+        .setDestroy(destroy)
+    })
+  }
+
   /** @inheritdoc */
   val conf: RuntimeConfig = new RuntimeConfig(client)
 

sql/connect/common/src/main/scala/org/apache/spark/sql/connect/UdfToProtoUtils.scala

@@ -107,6 +107,12 @@ private[sql] object UdfToProtoUtils {
           .setAggregate(true)
         f.givenName.foreach(invokeUdf.setFunctionName)
     }
+    // (SPARK-51705) Drain the ids of any ConnectBroadcast captured by this closure into
+    // ScalarScalaUDF.broadcast_ids. This must run after toUdfPacketBytes above, because serializing
+    // the closure is what triggers ConnectBroadcast.writeReplace (and thus the id capture). Mirrors
+    // the Python side, where PythonUDF.to_plan drains a threading.local registry into
+    // `python_udf.broadcast_ids` right after CloudPickleSerializer().dumps(...).
+    ConnectBroadcastCapture.drain().foreach(id => protoUdf.addBroadcastIds(id))
     invokeUdf.build()
   }
 }