pyspark-broadcast-lite

A user-land broadcast-variable workaround for PySpark on Spark Connect and Databricks Serverless, where SparkContext.broadcast(value) is unavailable because the Connect client has no SparkContext.

Uses the public SparkSession.addArtifact(..., pyfile=True) API to ship a Python module containing your serialized data. Inside UDFs, importlib.import_module(name) hits Python's normal import cache — giving once-per-worker-process load semantics, the closest this workaround gets to real broadcast behavior.

Not a replacement for real broadcast. No peer-to-peer block distribution, session-scoped lifecycle only, practical size ceiling in the low hundreds of MB. Intended as a stopgap until SPARK-51705 lands a proper SparkSession.broadcast(value) API.

Install

pip install git+https://github.com/Tagar/pyspark-broadcast-lite.git

Requires Python ≥ 3.9 and either:

• Spark Connect (OSS, ≥ 3.5.0), or
• Apache Spark Classic ≥ 4.0.0 (addArtifact was added to Classic in SPARK-50718), or
• Databricks Connect ≥ 14.0.1 (including Databricks Serverless).

Quick start

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, udf
from broadcast_lite import broadcast_lite

spark = SparkSession.builder.getOrCreate()

lookup = {"a": "apple", "b": "banana", "c": "cherry"}
mod_name = broadcast_lite(spark, lookup)

@udf("string")
def enrich(key):
    import importlib
    return importlib.import_module(mod_name).DATA.get(key, "unknown")

df = spark.createDataFrame([("a",), ("b",), ("z",)], ["key"])
df.select(enrich(col("key")).alias("fruit")).show()
# +-------+
# |  fruit|
# +-------+
# |  apple|
# | banana|
# |unknown|
# +-------+

For the once-per-partition setup pattern (recommended for high row counts), see examples/pandas_udf.py.

How it works

1. The value is pickled on the driver with pickle.HIGHEST_PROTOCOL.
2. For pickles ≤ 8 MB (default), the bytes are base64-encoded and inlined into a single .py source file whose import computes DATA. For larger values, a .zip package carries the pickle as package data, loaded via pkgutil.get_data over Python's built-in zipimport.
3. The file (or zip) is uploaded via spark.addArtifact(path, pyfile=True). Spark Connect's ArtifactManager inserts it into sys.path on Python workers and triggers importlib.invalidate_caches() for plain .py artifacts.
4. Inside the UDF, importlib.import_module(mod_name) executes the module body once per worker process and caches the result in sys.modules. Subsequent calls are O(1) dictionary lookups.

When to use — and when not to

Captured-value size	Recommendation
< ~100 KB	Don't bother. Let cloudpickle embed it in the UDF closure.
100 KB – 1 MB	Either works fine.
1 MB – 10 MB	broadcast_lite worth it — avoids closure bloat.
10 MB – 100 MB	Recommended. Cloudpickled closures become slow and large.
> 100 MB	Necessary. Raw cloudpickled closures hit spark.connect.grpc.maxInboundMessageSize (default 128 MiB) and spark.rpc.message.maxSize (default 128 MiB).
> 1 GB	Neither this nor sc.broadcast on managed services. Reshape the problem (join, structured store, per-partition cache).

Avoiding per-record overhead

A scalar @udf calls your function once per row. Even after the initial module import, every row pays a sys.modules lookup + attribute access. For millions of rows, use iterator-style pandas_udf or mapInPandas to get explicit once-per-partition setup:

from typing import Iterator
import pandas as pd
from pyspark.sql.functions import pandas_udf

mod_name = broadcast_lite(spark, big_dict)

@pandas_udf("string")
def enrich(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]:
    # runs ONCE per partition
    import importlib
    data = importlib.import_module(mod_name).DATA
    for batch in iterator:
        yield batch.map(lambda k: data.get(k, "?"))

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

Note: RDD.mapPartitions gives the same per-partition setup but is not available on Spark Connect (no RDD API). Use mapInPandas, mapInArrow, or iterator pandas_udf instead.

See examples/pandas_udf.py for a complete working example.

Caveats

• Not a real broadcast. No peer-to-peer block distribution — artifacts are fetched per executor from the Connect server's staging area. First-call latency scales linearly with executor count. Real sc.broadcast uses TorrentBroadcast for O(log N) fan-out.
• Size ceiling. Practical limit in the low hundreds of MB. Above that, upload time and Python-worker memory become problems.
• Session-scoped lifecycle only. There's no analog of Broadcast.unpersist(); uploaded modules live until the Connect session closes.
• Undocumented use on Databricks. The OSS mechanism (addArtifact(pyfile=True) → sys.path on executors) is stable and documented, but Databricks does not advertise using it for data. The wiring works; the use case is "unofficial."

The upstream fix

This package is a stopgap. The proper fix is tracked in SPARK-51705 ("Support sc.broadcast over Spark Connect"), which would add a native SparkSession.broadcast(value) API to both the DataFrame API and Spark Connect — making broadcast variables a first-class, typed, lifecycle-managed feature with the real TorrentBroadcast distribution topology under the hood.

If you want a clean, Databricks-supported broadcast API in a future Spark release, upvote the JIRA and follow the discussion.

Development

git clone https://github.com/Tagar/pyspark-broadcast-lite.git
cd pyspark-broadcast-lite
pip install -e ".[dev]"
pytest

The tests in tests/test_module_gen.py exercise the file-generation helpers directly and do not require a Spark session.

License

Apache-2.0. See LICENSE.