PySpark DataFrame Doctest: localCheckpoint

[shehabgamin]

pyspark/sql/dataframe.py::pyspark.sql.dataframe.DataFrame.localCheckpoint

https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.localCheckpoint.html

[benzsevern]

Sharing a benchmarked use case for localCheckpoint in case it helps prioritize this.

Iterative DataFrame algorithms (connected components, ML loops) grow the logical plan one layer per round. On Spark, DataFrame.localCheckpoint() truncates that lineage so each round costs O(round-work). On Sail 0.6.4 cache() is a no-op (per the compatibility matrix) and persist/localCheckpoint/checkpoint aren't implemented yet, so round k re-executes rounds 1..k-1 and wall time grows roughly quadratically until it stalls.

Concrete numbers: weakly-connected-components via min-label propagation (pure DataFrame joins to a fixpoint), Sail 0.6.4 in kubernetes-cluster mode, 3 workers, timing the iterative stage alone:

• 1,500 input rows: 2.9s
• 4,000 rows: 22.4s
• 12,000 rows: did not converge in over 2 minutes (one task pegged, memory climbing)

The non-iterative stages (join, groupBy, the block self-join) are fast and scale fine. Only the iterative loop degrades, which is the signature of un-truncated lineage.

Minimal repro (label propagation; no checkpoint available to break lineage):

from pyspark.sql import SparkSession, functions as F
import pandas as pd, time
spark = SparkSession.builder.remote("sc://<sail>:50051").getOrCreate()
n = 4000
edges = spark.createDataFrame(pd.DataFrame({"a": range(n - 1), "b": range(1, n)}))
labels = spark.range(n).select(F.col("id").alias("node"), F.col("id").alias("label"))
und = edges.select(F.col("a").alias("src"), F.col("b").alias("dst")) \
    .unionByName(edges.select(F.col("b").alias("src"), F.col("a").alias("dst")))
for r in range(50):
    t = time.perf_counter()
    nbr = und.join(labels.select(F.col("node").alias("dst"), F.col("label").alias("dl")), on="dst") \
             .groupBy("src").agg(F.min("dl").alias("m")).select(F.col("src").alias("node"), "m")
    new = labels.join(nbr, on="node", how="left") \
        .select("node", F.least("label", F.coalesce("m", "label")).alias("label"))
    # new = new.localCheckpoint()   # this is what keeps it flat on Spark
    changed = new.join(labels.withColumnRenamed("label", "old"), on="node") \
        .where(F.col("label") != F.col("old")).limit(1).count()
    labels = new
    print(f"round {r}: {time.perf_counter() - t:.1f}s changed={changed}", flush=True)
    if changed == 0:
        break

Per-round time climbs each iteration; a working localCheckpoint() is what keeps it flat.

Is localCheckpoint (materialize partitions to executor storage, replace the plan with a scan) a reasonable first target for this family of ops? We hit this building a distributed entity-resolution pipeline on Sail and are happy to help with a repro test or a PR if it's useful and you can point us at where materialize-and-truncate would slot into the plan/execution layer.

Thanks!

[linhr]

Sharing a benchmarked use case for localCheckpoint in case it helps prioritize this.

Thanks for sharing the use case! This is very valuable.

We've been discussing the feature internally and had a few prototypes (e.g., #1882, #1969). We will share an update when we make more progress.

Sail is an engine built for the cloud-native environment with the assumption of decoupled compute and storage. What we've been thinking so far is that all the cached or persisted data must be stored externally, either in an object storage or a shared file system (e.g., a network file system shared by the cluster). This means that we may implement cache, persist, checkpoint, and localCheckpoint in a similar manner, in which case we probably won't actually have "local" persisted data in the worker's memory or local disk, since each worker is ephemeral (auto-scaled based on the compute demand). We can see that it's not a one-to-one mapping to Spark features, but we believe it simplifies the architecture (since we don't need to extend the worker's lifetime for cached data and track cache affinity in the driver) and is scalable for broad use cases.

Does this design make sense for your setup? Happy to hear your thoughts!

[benzsevern]

Yeah, this works for us. The property I actually need is the lineage cut, replace the growing plan with a scan of materialized data, and external-backed materialization gives me exactly that. I never needed the bytes to physically sit on a worker; I latched onto localCheckpoint because on Spark it's the cheapest way to get the truncation, not because of where the data lands.

One nice side effect: an externally-backed localCheckpoint is more robust than Spark's, where the checkpoint is lost if the executor dies. Decoupled storage means it survives worker churn, which is strictly better for the auto-scaled case. No objection to all four sharing that mechanism either; for my use the in-memory-reuse semantics of cache don't matter, the truncation does.

The one thing that does matter is per-round cost. WCC does a materialize every round for ~log(n) rounds (pointer-jumping) or up to ~diameter rounds (plain label-prop), so the checkpoint write+read sits on the hot path, not a one-shot. The reason Spark splits localCheckpoint from reliable checkpoint is to keep that per-round write cheap when you don't need durability of the intermediate. So the question that matters more to me than local-vs-external: will there be a low-latency backing option (local SSD / shared NVMe scratch, or an object store tuned for small-object latency), or a way to pick the tier per call? As long as the per-round checkpoint cost stays well under the per-round compute, the external design fully covers this family of ops. If the only path is a high-latency object store with a full commit protocol, a 30-round loop could spend more time committing than computing.

Offer still stands on a repro test or a PR. I'm blocked at scale without this today, cache() is a no-op on 0.6.4 and round-tripping through collect() isn't viable past a few thousand rows, so there's no clean workaround for the 100M run. Happy to run my WCC bench against whichever of #1882 / #1969 lands and report per-round timings.

Thanks for the detailed reply.

[linhr]

Thanks for the response!

Yeah the checkpoint location would be configurable (hopefully even changeable at runtime within the session) so that the user can pick the most suitable storage for checkpointing/caching needs.

Technically, we rely on the object_store library for all storage operations in Sail, but it doesn't have to be a remote object storage service. For example, for a local Sail server, it can even be memory://, which mimics Spark's cache behavior. For distributed settings, a user may choose to connect to a low-latency network file system so that data persistence is not a bottleneck in iterative processing.

We appreciate your interest in contributing and helping with tests! We will likely have a polished proof-of-concept in the next few weeks (cc @lonless9 who will be working on this). We will let you know when it's ready and we'd love to see if it works for your use case.