[POC][SPARK-51705][CONNECT][SCALA] Support SparkSession.broadcast() for Scala UDFs over Spark Connect - #57573
Open
Tagar wants to merge 2 commits into
Open
[POC][SPARK-51705][CONNECT][SCALA] Support SparkSession.broadcast() for Scala UDFs over Spark Connect#57573Tagar wants to merge 2 commits into
Tagar wants to merge 2 commits into
Conversation
…) over Spark Connect Candidate-A pickle lane, Python-only v1: client cloudpickles the value and uploads it via the existing cache/<sha256> artifact channel; new CreateBroadcastCommand/UnpersistBroadcastCommand on ExecutePlan materialize a server-side Broadcast[PythonBroadcast] on the live driver SparkContext in SessionHolder; a new PythonUDF.broadcast_ids field (6) is resolved by SparkConnectPlanner into SimplePythonFunction.broadcastVars. Executor/worker path reused unchanged. broadcast_id == driver-side Broadcast.id so the worker's _broadcastRegistry/_from_id resolves it. DRAFT / DO-NOT-MERGE: proto _pb2.py stubs still need regen via dev/connect-gen-protos.sh (buf); BROADCAST_VALUE_TOO_LARGE quota not yet enforced; Scala/ScalarScalaUDF deferred (see companion spike branch). Co-authored-by: Isaac
…ala UDFs over Spark Connect
Stacks on the Python v1 (broadcast-connect-python-v1). Extends broadcast variables over Spark
Connect to the Scala ScalarScalaUDF path, which the Vizio "Cooker" anchor case needs (its broadcast
is consumed by a Scala UDF).
Design (Candidate A, writeReplace/readResolve):
- Client SparkSession.broadcast[T](v): SparkSerDeUtils.serialize(v) -> cacheArtifact (same cache/
channel as Python) -> CreateBroadcastCommand(value_type=JVM) -> CreateBroadcastResult(id) ->
ConnectBroadcast[T](id, v). The client is JVM-less so it cannot construct a real Broadcast[T];
ConnectBroadcast stands in and, when captured in a UDF closure, writeReplace() emits an id-only
ConnectBroadcastRef token (the value never hits the wire -- it already travelled via the artifact).
- Capture: writeReplace records the id in ConnectBroadcastCapture (thread-local); UdfToProtoUtils
.toProto drains it into ScalarScalaUDF.broadcast_ids (field 6) after serializing the closure.
- Server: SparkConnectPlanner.unpackScalaUDF resolves broadcast_ids against the per-session
SessionHolder registry (unknown id -> BROADCAST_NOT_FOUND) and binds them via a thread-local
around closure deserialization; ConnectBroadcastRef.readResolve swaps each token for the real
driver-side Broadcast[_]. No registry bound (client validation round-trip) -> a safe placeholder.
- CreateBroadcastCommand gains value_type {PYTHON(default), JVM}: JVM Java-deserializes the bytes
and sc.broadcast(value) directly so the executor reads Broadcast[T]. SessionHolder.broadcasts
widened to Broadcast[_] to hold both Broadcast[PythonBroadcast] (Python) and Broadcast[T] (Scala).
- Executor path unchanged (classic TorrentBroadcast fan-out).
Server unit tests added (registry holds Broadcast[T]; readResolve swap under withRegistry; client
round-trip placeholder; capture round-trip). Proto regenerated with the pinned toolchain; scalafmt
3.8.6 clean.
DRAFT / DO-NOT-MERGE: stacks on Python v1; SPARK-46032 (Scala closure deserialization on Connect)
is a related reliability risk; end-to-end Scala UDF test + UDAF coverage + BROADCAST_VALUE_TOO_LARGE
enforcement are follow-ups. See SCALA-PR-PLAN.md / SCALA-SPIKE-FINDINGS.md.
Co-authored-by: Isaac
Tagar
marked this pull request as ready for review
July 27, 2026 22:30
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
This PR adds
SparkSession.broadcast()support for Scala UDFs over Spark Connect (SPARK-51705), building on the Python-side support in #57385.Spark Connect's thin client has no
SparkContext, so the classicsc.broadcast(value)API is unavailable — attempting to close over a broadcast in a Connect UDF fails. This change introduces a server-mediated broadcast lane so Scala closures can capture and reuse a broadcast value on the server side, reusing the classicBroadcast[T]/TorrentBroadcastmachinery once the value reaches the driver.Key pieces:
SparkSession.broadcast[T](value: T): Broadcast[T]on the Connect client — serializes the value, uploads it via the artifact/cache lane, and registers it through a newCreateBroadcastCommand(value_type = JVM).Broadcast[T]shim inspark-connect-shimssoconnect-commoncan reference the real FQN without depending onspark-core.ConnectBroadcast/ConnectBroadcastRef/ thread-local resolver so a Java-serialized Scala UDF closure (UdfPacket) can resolve broadcast ids during deserialization on the server (writeReplace/readResolve).SparkConnectPlanner) resolves registered broadcasts and rebinds them into the deserialized closure;SessionHolderowns the per-session broadcast registry and its lifecycle.This is marked POC and stacks on #57385 (the first commit here is that PR's Python change; the second is the Scala work). It is intended to accompany the internal design discussion referenced in SPARK-51705.
Why are the changes needed?
Broadcast variables are the single largest feature-parity gap for migrating large Scala/UDF pipelines from classic Spark to Spark Connect / serverless. Without them, a large lookup table cannot be shared once per executor and must be re-materialized per record, causing OOM and cost blow-ups. See SPARK-51705.
Does this PR introduce any user-facing change?
Yes — a new
SparkSession.broadcast()method usable from the Scala Connect client, plusunpersistBroadcast. Behavior mirrors the classic broadcast API from the user's perspective.How was this patch tested?
SparkConnectSessionHolderSuitecoverage for the broadcast registry and lifecycle.build_and_testmatrix green on the author's fork (all connect / avro unit tests, linters, protobuf/Python codegen check, precompile).Was this patch authored or co-authored using generative AI tooling?
Yes, drafted with assistance from generative AI tooling and reviewed by the author.