HyukjinKwon commented on PR #60120:
URL: https://github.com/apache/spark/pull/60120#issuecomment-6651121437
### Benchmark: bundled cloudpickle 3.1.2 (master) vs cloudpickle-rust 0.1.0
(this PR)
**Setup:** Apple M3 Max (16 cores), macOS, CPython 3.14.5, JDK 26.0.1,
classic `local[4]`. Both variants run on the same build of this PR and differ
only in `python/pyspark/cloudpickle` (on the driver and in `pyspark.zip`): the
master variant restores the bundled cloudpickle from the parent commit.
Functions and classes are defined in `__main__`, so they are pickled by value
as in a notebook. The PR numbers include this PR's `Pickler` subclass that
keeps `co_filename`; plain `cloudpickle_rust.dumps` is another ~2-8% faster.
**Serialization only**, through the same calls as PySpark:
`CloudPickleSerializer().dumps(...)` for UDF commands and RDD batches, and
`pickle.loads` as the worker does. Best of 9 repeats in each of 3 processes per
variant.
| Case | `dumps` master | `dumps` PR | Speedup | `loads` master | `loads` PR
| Speedup |
|---|---:|---:|---:|---:|---:|---:|
| lambda | 11.5 µs | 8.8 µs | 1.31x | 4.6 µs | 4.4 µs | 1.03x |
| closure using globals and a helper | 21.8 µs | 15.7 µs | 1.39x | 8.9 µs |
8.7 µs | 1.03x |
| 1,000 distinct lambdas | 5.0 ms | 3.4 ms | 1.49x | 2.0 ms | 1.8 ms | 1.13x
|
| UDF command, lambda | 14.6 µs | 11.1 µs | 1.32x | 5.5 µs | 5.5 µs | 1.00x |
| UDF command, closure over a dynamic class | 56.5 µs | 42.3 µs | 1.34x |
33.8 µs | 32.4 µs | 1.04x |
| RDD batch, 1,024 ints | 10.0 µs | 10.0 µs | 1.00x | 15.0 µs | 15.2 µs |
0.98x |
| RDD batch, 1,024 (k, v) tuples | 46.3 µs | 46.6 µs | 0.99x | 47.1 µs |
46.6 µs | 1.01x |
| RDD batch, 1,024 dicts | 102.3 µs | 101.7 µs | 1.01x | 121.1 µs | 120.1 µs
| 1.01x |
| RDD batch, 1,024 dynamic class instances | 1.1 ms | 719.4 µs | 1.48x |
330.5 µs | 336.9 µs | 0.98x |
**End to end**, median of 5 processes per variant (each the median of 7 runs
after 2 warm-ups), with the range of the 5 process medians in parentheses:
| Case | master | PR | Speedup |
|---|---:|---:|---:|
| Python UDF, Arrow-optimized (default), 2M rows | 243 ms (225-474) | 260 ms
(225-314) | 0.93x |
| Python UDF, pickled rows (useArrow=False), struct result, 500k rows | 168
ms (153-236) | 159 ms (151-266) | 1.05x |
| RDD map + reduceByKey of (k, v) tuples, 1M records | 146 ms (140-153) |
145 ms (135-150) | 1.00x |
| RDD collect of dynamic class instances, 200k records | 183 ms (181-186) |
161 ms (160-171) | 1.14x |
| RDD parallelize + collect of tuples, 500k records | 127 ms (125-143) | 129
ms (124-131) | 0.99x |
| RDD of 400 tasks, closure over a dynamic class | 729 ms (663-1038) | 696
ms (685-828) | 1.05x |
| Broadcast of a 1M-entry dict, read by 8 tasks | 272 ms (262-292) | 263 ms
(258-287) | 1.04x |
**Takeaways**
- Pickling functions, closures and UDF commands is 1.3x-1.5x faster, and so
is pickling instances of classes defined in `__main__`. Unpickling is about the
same (1.0x-1.1x): it runs in CPython's C unpickler, and only the reconstructors
moved to Rust.
- Batches of builtin values (ints, tuples, dicts) are unchanged: CPython's
pickler handles `dict`, `list` and `tuple` before it calls `reducer_override`,
so no cloudpickle code runs for them.
- End to end, only the job that ships many dynamic class instances to the
driver is clearly faster (1.14x, the ranges do not overlap). In the other jobs
serialization is a small part of the time, and the differences are within
run-to-run noise (the ranges overlap), including both Python UDF cases.
- Broadcast values are pickled with the standard `pickle`
(`Broadcast.dump`), not cloudpickle, so the broadcast case is a control; it is
within noise too.
<details>
<summary>Benchmark script (<code>python bench_cloudpickle.py micro|e2e
OUTPUT_JSON</code>, run alternately against the two variants)</summary>
```python
#
# Benchmarks PySpark's cloudpickle usage, to compare the bundled cloudpickle
with cloudpickle-rust.
#
# Usage: python bench_cloudpickle.py micro|e2e OUTPUT_JSON
#
# micro: serialization only, through the same calls as PySpark
(CloudPickleSerializer for UDF
# commands and RDD batches, pickle.loads on the worker side). No
JVM.
# e2e: end-to-end jobs on a classic local[4] SparkSession.
#
# Everything is defined in __main__, so functions and classes are pickled by
value, as in a shell
# or a notebook.
import gc
import json
import operator
import os
import pickle
import sys
import time
import timeit
GLOBAL_OFFSET = 7
def helper(x):
return x * 3 + GLOBAL_OFFSET
def make_closure(n):
def inner(x):
return helper(x) * n + GLOBAL_OFFSET
return inner
class Point:
def __init__(self, x, y):
self.x = x
self.y = y
def norm1(self):
return abs(self.x) + abs(self.y)
class Model:
"""A small dynamic class captured by a UDF, e.g. a tokenizer or a
model."""
def __init__(self, n):
self.weights = {"w%d" % i: i * 0.5 for i in range(n)}
self.bias = 1.5
def predict(self, x):
return x * self.weights["w1"] + self.bias
def micro():
from pyspark import cloudpickle
from pyspark.serializers import CloudPickleSerializer
from pyspark.sql.types import DoubleType, StringType
ser = CloudPickleSerializer()
model = Model(100)
def udf_with_model(x):
return model.predict(helper(x))
cases = {
"lambda": lambda x: x + 1,
"closure using globals and a helper": make_closure(3),
"1,000 distinct lambdas": [(lambda i: lambda x: x + i)(i) for i in
range(1000)],
"UDF command, lambda": (lambda s: None if s is None else str(s),
StringType()),
"UDF command, closure over a dynamic class": (udf_with_model,
DoubleType()),
"RDD batch, 1,024 ints": list(range(1024)),
"RDD batch, 1,024 (k, v) tuples": [(i % 100, i) for i in
range(1024)],
"RDD batch, 1,024 dicts": [{"id": i, "name": str(i)} for i in
range(1024)],
"RDD batch, 1,024 dynamic class instances": [Point(i, -i) for i in
range(1024)],
}
results = {}
for name, obj in cases.items():
payload = ser.dumps(obj)
results[name] = {
"dumps_us": best_of(lambda: ser.dumps(obj)),
"loads_us": best_of(lambda: pickle.loads(payload)),
"bytes": len(payload),
}
return {"cloudpickle": cloudpickle.__version__, "results": results}
def best_of(fn, repeat=9):
timer = timeit.Timer(fn)
loops, _ = timer.autorange()
gc.collect()
return min(timer.repeat(repeat=repeat, number=loops)) / loops * 1e6
def e2e():
from pyspark import cloudpickle
from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.functions import udf
spark = (
SparkSession.builder.master("local[4]")
.config("spark.ui.enabled", "false")
.config("spark.sql.shuffle.partitions", "8")
.config("spark.driver.memory", "4g")
.getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")
sc = spark.sparkContext
model = Model(100)
add_one = udf(lambda x: x + 1, "long")
to_struct = udf(lambda x: (x, str(x)), "struct<a:long,b:string>",
useArrow=False)
def predict(x):
return model.predict(helper(x))
tuples = [(i, str(i)) for i in range(500_000)]
big_dict = {i: str(i) for i in range(1_000_000)}
def broadcast_job():
bv = sc.broadcast(big_dict)
total = sc.parallelize(range(8), 8).map(lambda i: len(bv.value) +
i).sum()
bv.destroy()
return total
cases = {
"Python UDF, Arrow-optimized (default), 2M rows": lambda:
spark.range(
0, 2_000_000, 1, 8
)
.select(F.sum(add_one("id")))
.collect(),
"Python UDF, pickled rows (useArrow=False), struct result, 500k
rows": lambda: spark.range(
0, 500_000, 1, 8
)
.select(F.sum(F.length(to_struct("id").b)))
.collect(),
"RDD map + reduceByKey of (k, v) tuples, 1M records": lambda:
sc.parallelize(
range(1_000_000), 8
)
.map(lambda x: (x % 1000, x))
.reduceByKey(operator.add)
.count(),
"RDD collect of dynamic class instances, 200k records": lambda:
sc.parallelize(
range(200_000), 8
)
.map(lambda x: Point(x, -x))
.collect(),
"RDD parallelize + collect of tuples, 500k records": lambda:
sc.parallelize(
tuples, 8
).collect(),
"RDD of 400 tasks, closure over a dynamic class": lambda:
sc.parallelize(range(400), 400)
.map(predict)
.count(),
"Broadcast of a 1M-entry dict, read by 8 tasks": broadcast_job,
}
results = {}
for name, job in cases.items():
for _ in range(2):
job()
times = []
for _ in range(7):
gc.collect()
start = time.perf_counter()
job()
times.append((time.perf_counter() - start) * 1000)
results[name] = {"ms": times}
spark.stop()
return {"cloudpickle": cloudpickle.__version__, "results": results}
if __name__ == "__main__":
mode, output = sys.argv[1], sys.argv[2]
data = micro() if mode == "micro" else e2e()
data["python"] = sys.version.split()[0]
data["pid"] = os.getpid()
with open(output, "w") as f:
json.dump(data, f)
print("wrote", output, data["cloudpickle"])
```
</details>
---
<!-- GITHUB_MCP_FOOTER: This attribution is automatically appended by GitHub
MCP. -->
_This comment was generated with [GitHub MCP](http://go/mcps)._
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
