[DRAFT][SPARK-51705][CONNECT][PYTHON] Support SparkSession.broadcast() (broadcast variables over Spark Connect)

common/utils/src/main/resources/error/error-conditions.json

@@ -1254,6 +1254,11 @@
           "<message>"
         ]
       },
+      "BROADCAST_NOT_FOUND" : {
+        "message" : [
+          "Cannot find a broadcast variable with id: <id>. It may never have been created on this session, may belong to another session, or may have been unpersisted or destroyed."
+        ]
+      },
       "CANNOT_FIND_CACHED_LOCAL_RELATION" : {
         "message" : [
           "Cannot find a cached local relation for hash: <hash>"

dev/sparktestsupport/modules.py

@@ -1167,6 +1167,7 @@ def __hash__(self):
         # sql unittests
         "pyspark.sql.tests.connect.test_connect_plan",
         "pyspark.sql.tests.connect.test_connect_basic",
+        "pyspark.sql.tests.connect.test_connect_broadcast",
         "pyspark.sql.tests.connect.test_connect_dataframe_property",
         "pyspark.sql.tests.connect.test_connect_channel",
         "pyspark.sql.tests.connect.test_connect_clone_session",

python/docs/source/reference/pyspark.sql/spark_session.rst

@@ -47,6 +47,7 @@ See also :class:`SparkSession`.
     SparkSession.addArtifact
     SparkSession.addArtifacts
     SparkSession.addTag
+    SparkSession.broadcast
     SparkSession.catalog
     SparkSession.clearTags
     SparkSession.conf

python/pyspark/errors/error-conditions.json

@@ -30,6 +30,16 @@
       "Length mismatch: Expected axis has <expected_length> elements, new values have <actual_length> elements."
     ]
   },
+  "BROADCAST_NOT_FOUND": {
+    "message": [
+      "Cannot find a broadcast variable with id: <id>. It may never have been created on this session, may belong to another session, or may have been unpersisted or destroyed."
+    ]
+  },
+  "BROADCAST_VALUE_TOO_LARGE": {
+    "message": [
+      "The broadcast value of <size> bytes exceeds the maximum allowed size of <max_size> bytes."
+    ]
+  },
   "BROADCAST_VARIABLE_NOT_LOADED": {
     "message": [
       "Broadcast variable `<variable>` not loaded."

python/pyspark/sql/connect/broadcast.py

@@ -0,0 +1,114 @@
+#
+# 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.
+#
+"""
+(SPARK-51705) Client-side broadcast proxy for Spark Connect.
+
+``ConnectBroadcast`` mirrors the classic :class:`pyspark.core.broadcast.Broadcast` pickling
+contract so that a UDF closure can reference a broadcast variable transparently. The value was
+already uploaded to the server (cloudpickle -> cache artifact -> CreateBroadcastCommand) and a
+driver-side broadcast id was returned; this proxy keeps a local copy of the value for driver-side
+``.value`` reads, and on pickling emits ``(_from_id, (broadcast_id,))`` -- exactly what the classic
+:class:`Broadcast` emits -- so the executor/worker resolves it from its ``_broadcastRegistry``
+keyed by that same id. No JVM is available on the Connect client, so all lifecycle operations are
+routed to the server as commands.
+"""
+
+import threading
+from typing import Any, Generic, Tuple, TYPE_CHECKING, TypeVar
+
+# NOTE: This reuses the classic ``pyspark.core.broadcast`` machinery so the pickled UDF closure is
+# byte-identical to classic PySpark (``_from_id`` + the driver-side broadcast id). ``pyspark.core``
+# is not available in a Spark Connect-only install (``pyspark-client``); this module is therefore
+# imported lazily (from ``SparkSession.broadcast`` and ``PythonUDF.to_plan``), never at package
+# import time, so the base Connect package still imports without the classic core present.
+from pyspark.core.broadcast import Broadcast, _from_id
+
+if TYPE_CHECKING:
+    from pyspark.sql.connect.session import SparkSession
+
+T = TypeVar("T")
+
+
+class _BroadcastCaptureRegistry(threading.local):
+    """Thread-local registry of broadcast ids captured while pickling a UDF closure.
+
+    Thread-local is correct because ``CloudPickleSerializer().dumps(...)`` in
+    ``PythonUDF.to_plan`` and the subsequent drain into ``PythonUDF.broadcast_ids`` happen
+    synchronously on the same plan-building thread. This mirrors the classic
+    :class:`pyspark.core.broadcast.BroadcastPickleRegistry`.
+    """
+
+    def __init__(self) -> None:
+        self.__dict__.setdefault("_registry", set())
+
+    def add(self, bid: int) -> None:
+        self._registry.add(bid)
+
+    def drain(self) -> "list[int]":
+        """Return the captured ids and clear the registry (dumps-then-drain-then-clear)."""
+        captured = list(self._registry)
+        self._registry.clear()
+        return captured
+
+
+# Module-level thread-local registry, drained by ``PythonUDF.to_plan``.
+_broadcast_capture_registry = _BroadcastCaptureRegistry()
+
+
+class ConnectBroadcast(Broadcast, Generic[T]):
+    """A broadcast variable created with :meth:`SparkSession.broadcast` over Spark Connect.
+
+    See Also
+    --------
+    pyspark.core.broadcast.Broadcast
+    """
+
+    def __init__(self, session: "SparkSession", broadcast_id: int, value: T) -> None:
+        # Intentionally does NOT call Broadcast.__init__ -- there is no live SparkContext on the
+        # Connect client. We populate only the attributes needed for driver reads and pickling.
+        self._session = session
+        self._broadcast_id = broadcast_id
+        self._value = value
+        # Mirror the classic executor-side shape: no JVM broadcast / SparkContext handle.
+        self._jbroadcast = None
+        self._sc = None
+        self._python_broadcast = None
+
+    @property
+    def value(self) -> T:
+        """Return the broadcasted value (kept locally on the client for driver-side reads)."""
+        return self._value
+
+    def unpersist(self, blocking: bool = False) -> None:
+        """Delete cached copies of this broadcast on the executors via the server."""
+        self._session._client._unpersist_broadcast(
+            self._broadcast_id, blocking=blocking, destroy=False
+        )
+
+    def destroy(self, blocking: bool = False) -> None:
+        """Destroy all data and metadata related to this broadcast via the server."""
+        self._session._client._unpersist_broadcast(
+            self._broadcast_id, blocking=blocking, destroy=True
+        )
+
+    def __reduce__(self) -> Tuple[Any, Tuple[int]]:
+        # Verbatim classic contract: reference the shared classic ``_from_id`` and the driver-side
+        # broadcast id, so the worker resolves it from ``pyspark.core.broadcast._broadcastRegistry``
+        # keyed by the same id. Also side-register the id so ``PythonUDF.to_plan`` can attach it to
+        # the proto ``broadcast_ids`` field.
+        _broadcast_capture_registry.add(self._broadcast_id)
+        return _from_id, (self._broadcast_id,)

python/pyspark/sql/connect/client/core.py

@@ -1887,6 +1887,8 @@ def handle_response(
             if b.HasField("create_resource_profile_command_result"):
                 profile_id = b.create_resource_profile_command_result.profile_id
                 yield {"create_resource_profile_command_result": profile_id}
+            if b.HasField("create_broadcast_result"):
+                yield {"create_broadcast_result": b.create_broadcast_result.broadcast_id}
             if b.HasField("checkpoint_command_result"):
                 yield {
                     "checkpoint_command_result": proto_to_remote_cached_dataframe(
@@ -2516,6 +2518,27 @@ def _create_profile(self, profile: pb2.ResourceProfile) -> int:
         profile_id = properties["create_resource_profile_command_result"]
         return profile_id
 
+    def _create_broadcast(self, artifact_hash: str, size_bytes: int) -> int:
+        """(SPARK-51705) Create a broadcast variable from an already-uploaded cache artifact and
+        return the server-assigned (driver-side) broadcast id."""
+        logger.debug("Creating a broadcast variable")
+        cmd = pb2.Command()
+        cmd.create_broadcast_command.artifact_hash = artifact_hash
+        cmd.create_broadcast_command.size_bytes = size_bytes
+        _, properties, _ = self.execute_command(cmd)
+        return properties["create_broadcast_result"]
+
+    def _unpersist_broadcast(
+        self, broadcast_id: int, blocking: bool = False, destroy: bool = False
+    ) -> None:
+        """(SPARK-51705) Unpersist (or destroy) a broadcast variable created over Connect."""
+        logger.debug("Unpersisting a broadcast variable")
+        cmd = pb2.Command()
+        cmd.unpersist_broadcast_command.broadcast_id = broadcast_id
+        cmd.unpersist_broadcast_command.blocking = blocking
+        cmd.unpersist_broadcast_command.destroy = destroy
+        self.execute_command(cmd)
+
     def _delete_ml_cache(self, cache_ids: List[str], evict_only: bool = False) -> List[str]:
         # try best to delete the cache
         try:

python/pyspark/sql/connect/expressions.py

@@ -763,6 +763,15 @@ def to_plan(self, session: "SparkConnectClient") -> proto.PythonUDF:
         expr.eval_type = self._eval_type
         expr.command = CloudPickleSerializer().dumps((self._func, output_type))
         expr.python_ver = self._python_ver
+        # (SPARK-51705) Drain the thread-local broadcast capture registry populated while pickling
+        # the command above (ConnectBroadcast.__reduce__ side-registers its id). This mirrors the
+        # classic dumps-then-drain-then-clear idiom in rdd._prepare_for_python_RDD, and covers both
+        # inline UDFs and spark.udf.register since both funnel through this to_plan.
+        from pyspark.sql.connect.broadcast import _broadcast_capture_registry
+
+        broadcast_ids = _broadcast_capture_registry.drain()
+        if broadcast_ids:
+            expr.broadcast_ids.extend(broadcast_ids)
         return expr
 
     def __repr__(self) -> str:

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

@@ -1759,6 +1759,7 @@ class ExecutePlanResponse(google.protobuf.message.Message):
     PIPELINE_EVENT_RESULT_FIELD_NUMBER: builtins.int
     PIPELINE_COMMAND_RESULT_FIELD_NUMBER: builtins.int
     PIPELINE_QUERY_FUNCTION_EXECUTION_SIGNAL_FIELD_NUMBER: builtins.int
+    CREATE_BROADCAST_RESULT_FIELD_NUMBER: builtins.int
     EXTENSION_FIELD_NUMBER: builtins.int
     METRICS_FIELD_NUMBER: builtins.int
     OBSERVED_METRICS_FIELD_NUMBER: builtins.int
@@ -1841,6 +1842,9 @@ class ExecutePlanResponse(google.protobuf.message.Message):
         register its result with the server.
         """
     @property
+    def create_broadcast_result(self) -> global___CreateBroadcastResult:
+        """(SPARK-51705) Server-assigned handle for a CreateBroadcastCommand."""
+    @property
     def extension(self) -> google.protobuf.any_pb2.Any:
         """Support arbitrary result objects."""
     @property
@@ -1889,6 +1893,7 @@ class ExecutePlanResponse(google.protobuf.message.Message):
         | None = ...,
         pipeline_query_function_execution_signal: pyspark.sql.connect.proto.pipelines_pb2.PipelineQueryFunctionExecutionSignal
         | None = ...,
+        create_broadcast_result: global___CreateBroadcastResult | None = ...,
         extension: google.protobuf.any_pb2.Any | None = ...,
         metrics: global___ExecutePlanResponse.Metrics | None = ...,
         observed_metrics: collections.abc.Iterable[global___ExecutePlanResponse.ObservedMetrics]
@@ -1902,6 +1907,8 @@ class ExecutePlanResponse(google.protobuf.message.Message):
             b"arrow_batch",
             "checkpoint_command_result",
             b"checkpoint_command_result",
+            "create_broadcast_result",
+            b"create_broadcast_result",
             "create_resource_profile_command_result",
             b"create_resource_profile_command_result",
             "execution_progress",
@@ -1945,6 +1952,8 @@ class ExecutePlanResponse(google.protobuf.message.Message):
             b"arrow_batch",
             "checkpoint_command_result",
             b"checkpoint_command_result",
+            "create_broadcast_result",
+            b"create_broadcast_result",
             "create_resource_profile_command_result",
             b"create_resource_profile_command_result",
             "execution_progress",
@@ -2010,6 +2019,7 @@ class ExecutePlanResponse(google.protobuf.message.Message):
             "pipeline_event_result",
             "pipeline_command_result",
             "pipeline_query_function_execution_signal",
+            "create_broadcast_result",
             "extension",
         ]
         | None
@@ -4203,6 +4213,28 @@ class CheckpointCommandResult(google.protobuf.message.Message):
 
 global___CheckpointCommandResult = CheckpointCommandResult
 
+class CreateBroadcastResult(google.protobuf.message.Message):
+    """(SPARK-51705) Server-assigned broadcast handle (mirrors CheckpointCommandResult)."""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    BROADCAST_ID_FIELD_NUMBER: builtins.int
+    broadcast_id: builtins.int
+    """(Required) The driver-side broadcast id. The client embeds this id in the pickled
+    Broadcast closure (Broadcast._from_id), which the executor/worker resolves against its
+    _broadcastRegistry keyed by the same id.
+    """
+    def __init__(
+        self,
+        *,
+        broadcast_id: builtins.int = ...,
+    ) -> None: ...
+    def ClearField(
+        self, field_name: typing_extensions.Literal["broadcast_id", b"broadcast_id"]
+    ) -> None: ...
+
+global___CreateBroadcastResult = CreateBroadcastResult
+
 class CloneSessionRequest(google.protobuf.message.Message):
     DESCRIPTOR: google.protobuf.descriptor.Descriptor