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

[Tagar]

[POC / DO-NOT-MERGE] — a design-preview PoC to anchor discussion on SPARK-51705. Posted in the spirit of the RDD-on-Connect PoC (#55888); not proposed for merge in current form. Feedback on the wire-protocol shape is explicitly solicited before this is polished into a mergeable series.

What changes were proposed in this pull request?

Adds a broadcast-variable API to Spark Connect (Python-only in this PoC), closing the gap where a Connect client cannot create a broadcast variable and reference it from a Python UDF. Today spark.sparkContext raises JVM_ATTRIBUTE_NOT_SUPPORTED on a Connect client, and a UDF that captures a broadcast fails at the worker with BROADCAST_VARIABLE_NOT_LOADED.

New public API — portable, byte-for-byte, with classic PySpark:

bc = spark.broadcast({"a": "apple", "b": "banana"})

@udf("string")
def enrich(key):
    return bc.value.get(key, "unknown")

df.select(enrich(df.key)).show()

Design ("server-mediated pickle lane", reusing the classic Broadcast[T] / TorrentBroadcast machinery):

client:  spark.broadcast(v) -> cloudpickle(v) -> client.cache_artifact(bytes) -> sha256
         -> ExecutePlan(CreateBroadcastCommand{artifact_hash}) <- CreateBroadcastResult{broadcast_id}
         -> ConnectBroadcast(id, v)   # __reduce__ -> (_from_id,(id,)), identical to classic
server:  getCachedBlockId(hash) -> getLocalBytes(blockId).toInputStream()   (transformCachedLocalRelation idiom)
         -> stage temp file -> new PythonBroadcast(path) -> sc.broadcast(pb) on the driver SparkContext in SessionHolder
         -> register in a new per-session SessionHolder broadcasts registry
plan:    transformPythonFunction resolves PythonUDF.broadcast_ids from the registry into SimplePythonFunction.broadcastVars
worker:  PythonRunner.writeBroadcasts -> setup_broadcasts -> _broadcastRegistry -> Broadcast._from_id   (UNCHANGED)

Key properties:

• No executor/worker changes. Peer-to-peer block distribution, per-executor caching, and fetch-on-demand are reused as-is; the only server change is populating SimplePythonFunction.broadcastVars at the three SparkConnectPlanner sites that currently hard-code an empty list.
• broadcast_id == the driver-side Broadcast.id, which is exactly what the worker keys on in _broadcastRegistry, so the classic _from_id unpickle contract resolves unchanged and UDF source is portable between classic and Connect.
• Transport reuses the existing cache/<sha256> artifact channel (content-addressing, dedup, per-session CacheId isolation, ref-counted GC, transparent at-rest encryption) — no new artifact namespace; the typed handle comes from CreateBroadcastResult.

Why are the changes needed?

Broadcast variables are a core PySpark primitive with no Spark Connect equivalent. Because Connect is the forward direction of the platform (and any classic-less deployment runs on it), the current guidance — "use a non-Connect cluster, or pass the value as a UDF parameter" — is a feature-parity gap and a portability cliff for commonly-taught code. sc.broadcast re-serialized as a per-task parameter defeats the purpose of broadcasting. This implements the ask in SPARK-51705.

Does this PR introduce any user-facing change?

Yes — a new SparkSession.broadcast(value) method. In classic mode it is a thin alias for spark.sparkContext.broadcast(value) (no behavior change); in Connect mode it returns a ConnectBroadcast proxy whose .value and closure-capture semantics match the classic Broadcast. Additive proto fields only; gated behavior, no change to existing APIs.

Proto surface (additive)

• expressions.proto: PythonUDF.broadcast_ids (field 6).
• commands.proto: CreateBroadcastCommand / UnpersistBroadcastCommand on the Command oneof (fields 20/21).
• base.proto: CreateBroadcastResult on ExecutePlanResponse (field 24).

_pb2.py / _pb2.pyi regenerated with the repo's pinned toolchain (protobuf==6.33.5, mypy-protobuf==3.3.0, buf remote plugins).

Carrier note (re-verified against current master, 2026-07-20)

broadcast_ids rides the legacy PythonUDF message, resolved into SimplePythonFunction.broadcastVars. The Language-agnostic / Unified UDF Protocol (SPARK-55278) is not wired into Spark Connect, adds no proto, and has no broadcast handling, so the legacy carrier is correct today. If Connect UDFs later migrate onto that path, broadcast_ids would move with them — tracked, not a blocker.

How was this patch tested?

PoC. Locally: proto regenerated with the pinned toolchain and runtime-verified (field round-trip under protobuf 6.33.5); Python ruff check + ruff format --check clean. A new python/pyspark/sql/tests/connect/test_connect_broadcast.py covers the E2E dict-in-UDF, portability, unpersist/destroy, and the negative BROADCAST_NOT_FOUND case; full E2E + Scala compile run under CI. Definition-of-done includes removing broadcast from test_connect_compatibility.py::expected_missing_connect_methods.

Scope / follow-ups

• PoC scope = Python scalar/pandas UDF only. transformPythonTableFunction (UDTF) and transformPythonDataSource keep empty broadcasts for now.
• Scala (ScalarScalaUDF) is a separate follow-up. A companion spike explored it: unlike Python (id injected into broadcastVars without touching the closure), a Scala UDF is a Java-serialized closure carrying the Broadcast[T] object graph, so the id must be resolved during closure deserialization (client writeReplace id-token → server readResolve from the registry). Related to closure-deserialization on Connect (SPARK-46032).
• Coordination: the original SPARK-51705 sketch proposed a dedicated streaming RPC; this PoC deliberately reuses AddArtifacts + ExecutePlan commands instead. Wire-protocol feedback welcome — happy to align.

Related work

• SPARK-51705 — the feature request this implements.
• #55888 — [POC] RDD in Python Spark Connect (HyukjinKwon), which enumerated broadcast among the not-yet-supported SparkContext methods.
• SPARK-46032 — closure capture across the Connect boundary (related for the Scala follow-up).