Fix review follow-ups and run the quickstart notebook in CI

[hamed]

Corrects the items left open by the 2026-08-23 source review, and adds a CI job
that executes examples/quickstart.ipynb.

Behavior

• Spark checkpoint directory is restored. bootstrap used to leave its own
  directory set on the session, which silently redirected every later
  checkpoint() the caller made. One detail surfaced here: setCheckpointDir
  appends a UUID and getCheckpointDir reports the combined path, so a naive
  save-and-restore deepens the path once per call. _checkpoint_root strips the
  UUID, and the caller's root stays stable across repeated calls.
• The local checkpoint fallback is per user. A fixed /tmp/replicas belongs
  to whoever creates it first; every other user on the machine then fails to
  write into it. The fallback is now <tempdir>/replicas-<user>.
• box_plot explains the replica column. It always required one. Without
  it, pandas raised a bare KeyError from melt that named neither the reason
  nor the fix.
• plot_pr validates ci. A value outside (0, 1] used to fail later,
  inside percentile_approx, far from the caller.
• replica is int32 on all three backends. pandas produced int64 and
  Polars Int64, so the three schemas disagreed on a column that holds a
  replica index.

Documentation

• at() states that its lowest-qualifying-threshold rule suits a metric that
  does not increase as the threshold falls, such as precision. A recall
  target degenerates to the group's minimum threshold, because recall does not
  decrease as the threshold falls. Every current caller passes precision, so
  this was latent.
• confusion_table states that the non-null and mutually-exclusive conditions
  on its indicator columns belong to the caller, and why they are not checked:
  validating them costs a full pass, which on Spark means an eager job before
  the lazy plan the caller asked for.
• The metric functions state the single-partition cost of an empty group_by.
  bootstrap states the memory cost of an empty by on Spark, and points at
  the README for the NaN exception to the cross-backend parity guarantee.

CI

A notebook job runs scripts/check_example_notebook.py examples/quickstart.ipynb.
The script executes the notebook in memory and compares every output with the
committed one, so a change that breaks the documented workflow — or silently
moves its numbers — turns the job red. Memory addresses are normalized and
Spark's stderr log lines are ignored; the committed file is never rewritten.

I verified both failure paths: a tampered expectation produces a diff, and an
injected defect in the Spark metric backend (recall = TP / (positives + 1))
fails at cell 16. The defect was reverted.

examples/precision_recall.ipynb is deliberately excluded. It is the reference
design: it defines its own copies of the functions and needs kagglehub,
scikit-learn, and CatBoost.

Not included

An earlier draft of the review reported a crash in box_plot and plot_pr when
hue, row, or col is "replica". Both helpers build their distribution
across replicas, so no caller passes that value. The crash is real but
unreachable, and it is not a finding.

Verification

• ruff check . — All checks passed!; ruff format --check . — 30 files already formatted.
• pytest -q — 104 passed (95 before, 9 new tests).
• scripts/check_example_notebook.py examples/quickstart.ipynb — clean, matches committed outputs, ~35s.

🤖 Generated with Claude Code

[hamed]

Thanks — points 1, 2 and 3 were right, and b95af81 reworks the checkpoint handling. Reply to each below.

1. Concurrency — fixed

Agreed, and the failure mode is as you describe. Worth stating the constraint plainly first: Dataset.checkpoint reads SparkContext.checkpointDir, and Spark exposes no per-call checkpoint path. Writing to a caller-chosen directory requires mutating session-wide state. The only alternatives are localCheckpoint (loses reliability, which is the point of the checkpoint here) or replacing checkpoint() with an explicit write.parquet / read.parquet round trip (a much larger semantic change, and it drops Spark's internal serializer for arbitrary schemas).

So the sequence is serialized rather than removed:

with _CHECKPOINT_LOCK:
    previous = spark_context.getCheckpointDir()
    spark_context.setCheckpointDir(checkpoint_dir)
    try:
        return df.checkpoint(eager=True)
    finally:
        _restore_checkpoint_dir(spark_context, previous)

A SparkContext is a per-process singleton, so a process-wide lock covers every caller that can reach the state. Separate driver processes have separate contexts and no shared state to race on.

Also added: when the caller has already configured a directory and passes no checkpoint_dir, _checkpoint now returns early and mutates nothing at all. That path never takes the lock.

test_concurrent_checkpoints_do_not_interleave covers it, on fakes so it is fast and deterministic. Removing the lock fails it with exactly the interleaving you predicted:

['worker-a', 'worker-a', 'worker-b', 'worker-b', 'worker-a', 'worker-b']

— worker-a checkpointing while worker-b's directory was set.

2. Swallowed exceptions — fixed

It now warns, naming the directory and telling the caller what state they are left in:

except Exception as error:  # noqa: BLE001 - Py4J raises outside one hierarchy
    warnings.warn(
        "replicas could not restore the Spark checkpoint directory to "
        f"{target!r}: {error!r}. ...",
        RuntimeWarning,
        stacklevel=3,
    )

It warns rather than raises on purpose: the bootstrap has already succeeded by that point, and discarding a completed job to report a state-cleanup failure is the worse trade. The catch stays broad because Py4J errors do not share one hierarchy — the defect you identified was the silence, not the breadth.

3. Path parsing — removed

The regex is gone. The backend now remembers the root it set alongside the value Spark reported for it, so restoration never reads Spark's path format.

One consequence I want to be explicit about, because it is a real trade rather than a clean win: a directory this package never set has no known root, so the first restore of a caller-configured directory still costs one level (<C>/<u1> → <C>/<u1>/<u2>). Every call after that reuses the remembered root, so the depth stops there. The old regex restored <C> exactly but only by assuming Spark's generated format.

Bounded at +1 and format-independent beats exact and format-dependent, but say the word if you would rather have it exact. test_repeated_checkpoints_keep_the_caller_directory_at_one_depth pins the behavior over 50 calls.

4. Per-user fallback — kept, and separable

This one I would push back on gently. It is not hardening for a hypothetical: /tmp/replicas is created by whoever runs bootstrap first, with that user's ownership, and every other user on the box then fails to create or write into it. Shared build agents and shared analysis hosts are ordinary for this workload. The cost is one 12-line function with no callers other than the fallback branch.

It is fully separable from the concurrency work, so if you would rather ship it on its own, it reverts in one commit and I will open it separately.

5. Notebook checker — images

Correct, and deliberate. Figure PNG bytes shift with matplotlib and font versions, so byte comparison would be red on unrelated dependency bumps rather than on defects.

Worth noting what is covered: display_data text/plain includes the figure reprs, so geometry and axes count are compared —

<Figure size 778x600 with 4 Axes>

A facet that stops rendering, or a grid that changes shape, fails the job. Pixel-level regressions do not. If you want real image diffing, pytest-mpl with a tolerance is the usual route; happy to add it as a separate change.

Verification

• ruff check . — All checks passed!; ruff format --check . — 30 files already formatted.
• pytest -q — 108 passed (95 on main, 13 new).
• scripts/check_example_notebook.py examples/quickstart.ipynb — clean, matches committed outputs.

[hamed]

Thanks for the re-review. 45d90d8 covers the one remaining item: the checkpoint-root edge case is now documented as a known limitation rather than living only in a test.

It is stated in four places, so a reader meets it before a surprising path does:

• _restore_target — the mechanism, why the reported <root>/<uuid> is the best available value, and why stripping the id back off would depend on Spark's internal path format.
• The bootstrap docstring — one sentence, pointing at the README.
• README, under "Core API" — the user-facing statement: a directory configured outside replicas ends up one generated level deeper the first time a call with an explicit checkpoint_dir restores it, and later calls hold the depth.
• CHANGELOG, on the checkpoint entry.

No behavior change in this commit. ruff check . and ruff format --check . clean; pytest -q still 108 passed.

The PR is ready to merge whenever you are.