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

[Tagar]

What changes were proposed in this pull request?

DRAFT / DO-NOT-MERGE — design preview to anchor the SPARK-51705 discussion. Not for merge in current form.

This PR adds a broadcast-variable API to Spark Connect (Python-only v1), closing the gap where a Connect/Serverless client cannot create a broadcast variable and reference it from a Python UDF (today this fails at the worker with BROADCAST_VARIABLE_NOT_LOADED).

Design: Candidate A — server-mediated pickle lane reusing classic Broadcast[T]/TorrentBroadcast.

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,)), mirrors classic contract
server:  getCachedBlockId(hash) -> getLocalBytes(blockId).toInputStream()  (transformCachedLocalRelation idiom)
         -> stage temp file -> new PythonBroadcast(path) -> sc.broadcast(pb) on the live driver SC in SessionHolder
         -> register in a new SessionHolder broadcasts registry (sibling to pythonAccumulator)
plan:    transformPythonFunction resolves PythonUDF.broadcast_ids from the registry into SimplePythonFunction.broadcastVars
worker:  PythonRunner.writeBroadcasts -> setup_broadcasts -> _broadcastRegistry -> Broadcast._from_id   (UNCHANGED)

The transport reuses the existing cache/<sha256> artifact channel (content-addressing, dedup, per-session CacheId(sessionUUID,hash) isolation, ref-counted GC, transparent at-rest encryption) rather than adding a new broadcasts/ prefix. The typed handle comes from CreateBroadcastResult, not the artifact prefix.

Why are the changes needed?

JIRA: https://issues.apache.org/jira/browse/SPARK-51705 ([CONNECT] Support sc.broadcast over Spark Connect, open, unassigned). Martin Grund endorsed "lift broadcast to the SparkSession" on-list. Broadcast is a core PySpark primitive with no Connect equivalent; Databricks Serverless / shared clusters run Connect, so the current guidance is to abandon shared/serverless for single-user clusters — a governance regression this closes.

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

broadcast_ids rides the legacy PythonUDF proto message (field 6), resolved in SparkConnectPlanner.transformPythonFunction into SimplePythonFunction.broadcastVars. This was the #1 gate to re-check:

• The Language-agnostic / Unified UDF Protocol (JIRA SPARK-55278, status Reopened, no fixVersion) did partially land, but on a separate layer: a new udf/worker/ module + ExternalUserDefinedFunction/ExternalUDFExec under a default-OFF experimental SQLConf (spark.sql.execution.udf.unified.execution.enabled, v4.2.0). It is not wired into Spark Connect (0 references under sql/connect), added no .proto, and has no broadcast handling at all. So the legacy PythonUDF carrier is correct and unobstructed for v1. Do NOT put broadcast_ids on any udf/worker/proto/* message.
• The JVM broadcast wire format is byte-identical on master (PythonWorkerUtils.writeBroadcasts unchanged; core/broadcast.py unchanged). Note: the Python worker decode was refactored (SPARK-56519 / SPARK-56324) — worker_util.setup_broadcasts now delegates to pyspark.worker_message.BroadcastInfo.from_stream — but this is Python-internal and format-preserving, so the Scala-side Connect changes need no worker edits.
• SPARK-55476 ("Refactor how broadcast variables are passed") is OPEN, not merged (PR [SPARK-55476][PYTHON] Refactor broadcast variable protocol apache/spark#54258 closed unmerged) — no wire-format change landed.

How was this patch tested?

(To be added; DRAFT skeleton.)

• Parity: flip broadcast out of test_connect_compatibility.py::expected_missing_connect_methods.
• E2E Python: spark.broadcast(dict_10k) referenced in a UDF over a >=1M-row DataFrame returns correct results (the exact case that fails today).
• Portability: UDF source byte-identical classic vs Connect; identical outputs.
• Encryption-on (spark.io.encryption.enabled=true): value re-encrypted at rest on server, worker reads via EncryptedPythonBroadcastServer.
• Lifecycle/leak: unpersist()/destroy() frees driver+executor blocks and releases the RefCountedCacheId; SessionHolder.close() sweeps.
• Negative: unknown/foreign broadcast_id -> BROADCAST_NOT_FOUND (loud), not a silent empty list.

Caveats / scope

• DRAFT / DO-NOT-MERGE. Coordinate with Jayadeep Jayaraman's dedicated-RPC sketch on SPARK-51705 (his approach may prefer a dedicated RPC over ExecutePlan commands) before finalizing.

Local verification status (2026-07-20)

• Proto stubs regenerated with Spark's pinned toolchain (protobuf==6.33.5, mypy-protobuf==3.3.0, buf remote plugins protocolbuffers/python:v33.5) via dev/connect-gen-protos.sh — the _pb2.py/_pb2.pyi diffs are limited to the new broadcast fields/messages; runtime import + field round-trip verified under protobuf 6.33.5.
• Python lint/format clean: ruff check (pinned ruff==0.14.8) passes; ruff format --check clean on all new/modified files.
• Not run locally (no JDK / full env here): Scala compile (build/sbt), full mypy (needs numpy/pandas-stubs), and the Connect E2E test suite (needs a running Connect server). These will run under Apache CI.

Scope / caveats

• v1 = Python scalar/pandas UDF only. UDTF (transformPythonTableFunction) and DataSource (transformPythonDataSource) paths intentionally keep empty broadcasts; Scala/ScalarScalaUDF deferred to a separate JIRA.
• Follow-up risk: if Connect Python UDFs ever migrate onto the unified ExternalUserDefinedFunction (opaque payload, no broadcast field) and the SQLConf flips to default-on, broadcasts need re-plumbing on that carrier. Tracked, not a v1 blocker.
• Re-verify proto field numbers at PR time — other in-flight PRs can claim numbers.

This pull request and its description were written by Isaac.

Proto diff

--- a/sql/connect/common/src/main/protobuf/spark/connect/expressions.proto
+++ b/sql/connect/common/src/main/protobuf/spark/connect/expressions.proto
@@ message PythonUDF (currently lines ~467-478)
   message PythonUDF {
     // Output type of the Python UDF
     DataType output_type = 1;
     // EvalType of the Python UDF
     int32 eval_type = 2;
     // The encoded commands of the Python UDF
     bytes command = 3;
     // The name of the Python UDF
     string python_ver = 4;
     // Additional includes for the Python UDF
     repeated string additional_includes = 5;
+    // (SPARK-51705) Server-assigned ids of broadcast variables referenced by this UDF.
+    // Resolved by SparkConnectPlanner against the per-session SessionHolder broadcasts
+    // registry into SimplePythonFunction.broadcastVars. Field 6 is next-free.
+    repeated int64 broadcast_ids = 6;
   }
   //  NOTE: ScalarScalaUDF (next-free is also 6) intentionally left unchanged for Python-only v1.

--- a/sql/connect/common/src/main/protobuf/spark/connect/commands.proto
+++ b/sql/connect/common/src/main/protobuf/spark/connect/commands.proto
@@ message Command.command_type oneof (currently lines ~36-61; 1..19 used, highest pipeline_command=19, extension=999)
   message Command {
     oneof command_type {
       // ... existing arms 1..19 ...
       PipelineCommand pipeline_command = 19;
+      CreateBroadcastCommand    create_broadcast_command    = 20;
+      UnpersistBroadcastCommand unpersist_broadcast_command = 21;
       // This field is used to mark extensions to the protocol.
       google.protobuf.Any extension = 999;
     }
   }
+
+// (SPARK-51705) Create a broadcast variable from an already-uploaded cache/<sha256> artifact.
+message CreateBroadcastCommand {
+  // sha256 hash returned by client.cache_artifact(cloudpickle(value)).
+  string artifact_hash = 1;
+  // Optional: uncompressed byte size, used to enforce BROADCAST_VALUE_TOO_LARGE quota.
+  int64 size_bytes = 2;
+}
+
+// (SPARK-51705) Release a broadcast variable created over Connect.
+message UnpersistBroadcastCommand {
+  int64 broadcast_id = 1;
+  bool  blocking     = 2;
+  bool  destroy      = 3;
+}

--- a/sql/connect/common/src/main/protobuf/spark/connect/base.proto
+++ b/sql/connect/common/src/main/protobuf/spark/connect/base.proto
@@ message ExecutePlanResponse.response_type oneof (currently arms up to 23; field 3 is an UNGUARDED gap - do NOT reuse)
   message ExecutePlanResponse {
     // ... existing arms ... (highest oneof arm = 23)
+    // (SPARK-51705) Server-assigned handle for a CreateBroadcastCommand. Field 24 is next-free.
+    CreateBroadcastResult create_broadcast_result = 24;
     // Support arbitrary result objects.
     google.protobuf.Any extension = 999;
   }
+
+// (SPARK-51705) Server-assigned broadcast handle (mirror CheckpointCommandResult).
+message CreateBroadcastResult {
+  int64 broadcast_id = 1;
+}

---

Companion: Scala UDF spike (broadcast-connect-scala-spike)

The anchor customer workload consumes its broadcast inside a Scala UDF, which Python-only v1 does not cover. The companion branch broadcast-connect-scala-spike is an exploratory spike (not for merge) documenting the path (see SCALA-SPIKE-FINDINGS.md):

• A Scala Connect UDF is a Java-serialized closure carrying the Broadcast[T] object graph, so a broadcast id must be resolved during closure deserialization — unlike Python, where the id rides PythonUDF.broadcast_ids and is injected into SimplePythonFunction.broadcastVars without touching the closure.
• Viable approach: client ConnectBroadcast[T].writeReplace emits an id-only token; server readResolve swaps the id for the registry Broadcast[_] via a thread-local around unpackScalaUDF. Stacks on this PR's SessionHolder registry.
• SPARK-46032 (closure deserialization on Connect) is a related reliability risk, not a strict blocker.
• Recommendation: land Python v1 (this PR) first; pursue Scala as a separate follow-up JIRA.

Provenance

Design of record and cited impact analysis were developed in a companion research effort; this PR implements the "Candidate A" pickle lane. Note: the fork's base history is apache/spark and contains Apache's own SSL test-fixture PEM keys (core/src/test/resources/key.pem, etc.) — those are upstream fixtures, not introduced by these commits.