HyukjinKwon commented on code in PR #59540:
URL: https://github.com/apache/spark/pull/59540#discussion_r4324072429
##########
python/pyspark/sql/connect/dataframe.py:
##########
@@ -2417,26 +2418,47 @@ def offset(self, num: int) -> ParentDataFrame:
return DataFrame(plan.Offset(child=self._plan, offset=num),
session=self._session)
def checkpoint(self, eager: bool = True) -> ParentDataFrame:
- cmd = plan.Checkpoint(child=self._plan, local=False, eager=eager)
+ relation_id = str(uuid.uuid4())
+ cmd = plan.Checkpoint(
+ child=self._plan,
+ local=False,
+ eager=eager,
+ relation_id=relation_id,
+ )
_, properties, self._execution_info =
self._session.client.execute_command(
cmd.command(self._session.client)
)
assert "checkpoint_command_result" in properties
- checkpointed = properties["checkpoint_command_result"]
- assert isinstance(checkpointed._plan, plan.CachedRemoteRelation)
- return checkpointed
+ assert properties["checkpoint_command_result"] == relation_id, (
Review Comment:
**Nit (P3):** If the response carries a different ID than the one requested
(for example from a released server that ignores the new field), the server has
already run the checkpoint and cached the result under the returned ID. This
`assert` fails before anything wraps that ID in a `CachedRemoteRelation`, and
that relation's finalizer is the client's only release path, so the server-side
entry is never removed and stays pinned until the session closes. For an eager
`localCheckpoint()` those are persisted blocks on the executors, and every
retry adds another.
Because this is a bare `assert`, `python -O` removes the check: the method
then returns a DataFrame over the requested ID, which the server never cached,
so its first action fails with `CONNECT_INVALID_PLAN.DATAFRAME_NOT_FOUND` while
the real entry is still orphaned. The same applies to `localCheckpoint()` at
line 2455. This applies as long as the methods keep rejecting a mismatched ID.
**Verification:**
- **Regression:** When the response carries a relation ID other than the
requested one (for example, because the server did not honor the request's
relation_id), checkpoint() and localCheckpoint() each raise
PySparkAssertionError, and afterwards the server-side session cache no longer
holds the returned ID.
- **Inspection:** Neither method relies on an assert statement to validate
the checkpoint response, so Python's optimization mode cannot remove the
rejection.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala:
##########
@@ -4083,6 +4083,14 @@ class SparkConnectPlanner(
private def handleCheckpointCommand(
checkpointCommand: CheckpointCommand,
responseObserver: StreamObserver[proto.ExecutePlanResponse]): Unit = {
+ val dfId = if (checkpointCommand.hasRelationId) {
+ if (checkpointCommand.getRelationId.isEmpty) {
+ throw
InvalidPlanInput("CONNECT_INVALID_PLAN.CHECKPOINT_RELATION_ID_EMPTY", Map.empty)
+ }
+ checkpointCommand.getRelationId
Review Comment:
**Non-blocking (P2):** This accepts any non-empty supplied ID, including one
that is already cached in this session. The handler then plans the relation and
runs the checkpoint (eager by default), and only
`sessionHolder.cacheDataFrameById` notices the existing entry, via
`putIfAbsent`, and throws `SparkException.internalError`. The client gets
`INTERNAL_ERROR` (SQLSTATE XX000), which Spark's error guide reserves for bugs
in Spark itself, even though `commands.proto` now documents uniqueness as a
client precondition and the empty-ID case just above gets a
`CONNECT_INVALID_PLAN` condition. Clients and libraries can't tell this apart
from a genuine internal failure without matching the message text, and a client
that resends a checkpoint after losing the first response recomputes the whole
checkpoint (reliable-checkpoint files or persisted local-checkpoint blocks)
only to have it discarded. The existing entry is not overwritten, so nothing is
corrupted.
The description says the existing duplicate-ID error is kept on purpose, but
clients could not reach it before because the server always chose the ID. Once
this ships, the error becomes part of the new field's observable contract.
Requests without `relation_id` are unaffected.
**Recommended change:** In handleCheckpointCommand, when relation_id is
supplied and already identifies a DataFrame in the session cache, reject it
before the relation is transformed or checkpointed. Use a new user-facing
CONNECT_INVALID_PLAN subcondition that names the ID. If the final atomic
insertion of a supplied ID loses a race, raise that same condition instead of
INTERNAL_ERROR. Server-generated IDs (absent field and foreachBatch) keep
cacheDataFrameById's INTERNAL_ERROR. Add the subcondition to the error catalog
next to CHECKPOINT_RELATION_ID_EMPTY, and add a regression test for the Spark
Connect server checkpoint command.
**Why this works:** The handler is the only layer that knows an ID was
client-supplied. Before any work, it performs a side-effect-free lookup in the
session cache, which rejects serial reuse without paying for a checkpoint. The
session holder then reports whether the atomic insertion happened, as an
explicit outcome rather than through the internal-error message. The handler
turns a failed insertion of a supplied ID into the same structured
InvalidPlanInput. Because putIfAbsent remains the sole uniqueness decision, no
entry is ever overwritten, and collisions of server-generated IDs keep their
internal-bug meaning.
**Scope:** Move duplicate-ID admission for client-supplied checkpoint
relation IDs into the Connect checkpoint handler. The handler raises a
structured user-facing condition before any work, and also when it loses an
insertion race. The condition is added to the error catalog, with a server
regression test.
**Compatibility:** Requests that omit relation_id, and foreachBatch caching,
keep server-generated UUIDs and raise INTERNAL_ERROR on an unexpected
collision. Rejecting a duplicate never replaces or removes the entry already
cached under that ID.
**Risks:** The pre-check is not atomic, so on its own it would let a
concurrent duplicate reach INTERNAL_ERROR. Mapping the insertion outcome is
therefore required, not optional. The new subcondition becomes part of the
released relation_id protocol contract, so its name and parameters should be
settled before release. The PR description says the existing duplicate-ID error
is kept intentionally; this repair still rejects duplicates and changes only
the error's classification and when it is raised.
**Constraints:** Never overwrite or evict an existing dataFrameCache entry
when rejecting a duplicate. Keep INTERNAL_ERROR for collisions of
server-generated IDs (absent field and foreachBatch).
**Success:** A CheckpointCommand whose supplied relation_id already
identifies a cached DataFrame in the session fails with a CONNECT_INVALID_PLAN
subcondition that names the ID, rather than INTERNAL_ERROR. It fails before its
relation is planned or checkpointed, and the existing entry is unchanged. When
two requests race with the same supplied ID, exactly one is cached, and the
other fails with the same CONNECT_INVALID_PLAN subcondition. Requests that omit
relation_id, and foreachBatch caching, keep server-generated IDs and their
existing INTERNAL_ERROR on collision.
**Verification:**
- **Regression:** First, a checkpoint command caches a DataFrame under a
supplied ID. A second checkpoint command in the same session with that ID then
fails with the new CONNECT_INVALID_PLAN subcondition and its parameters, before
its relation is planned or checkpointed. The first cached DataFrame remains
resolvable by that ID.
- **Inspection:** A failed atomic insertion of a client-supplied ID reports
the same structured subcondition, and the insertion path for server-generated
IDs keeps INTERNAL_ERROR.
- **Compatibility:** A checkpoint command without relation_id still returns
and caches a server-generated ID, and foreachBatch DataFrame caching behaves as
before.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala:
##########
@@ -4083,6 +4083,14 @@ class SparkConnectPlanner(
private def handleCheckpointCommand(
checkpointCommand: CheckpointCommand,
responseObserver: StreamObserver[proto.ExecutePlanResponse]): Unit = {
+ val dfId = if (checkpointCommand.hasRelationId) {
+ if (checkpointCommand.getRelationId.isEmpty) {
Review Comment:
**Nit (P3):** The cited `test_connect_plan` suite has no checkpoint test,
and the existing Connect checkpoint tests only cover the happy path: the Python
client always sends a fresh UUID and the Scala client never sends one. Nothing
sends an empty or already-cached ID, checks that a supplied ID is honored or
that an omitted one is still server-generated, or builds `plan.Checkpoint` with
and without `relation_id`. So deleting this guard (an empty ID would then be
cached under `""`), misspelling or re-parameterizing
`CHECKPOINT_RELATION_ID_EMPTY` (its message is resolved only when the error is
thrown), changing what a duplicate ID produces, or making the builder set the
field when it is `None` would all leave CI green. This is the contract the PR
offers other protocol clients, and today it is protected only indirectly by the
Python happy path.
**Verification:**
- **Behavior:** A checkpoint command with a fresh supplied ID caches the
DataFrame under exactly that ID and returns it in the result. A command without
the field returns a server-generated ID that is cached.
- **Regression:** A checkpoint command with an empty supplied ID fails with
CONNECT_INVALID_PLAN.CHECKPOINT_RELATION_ID_EMPTY and leaves the session cache
unchanged. A command with an already-cached supplied ID fails with the
duplicate-ID condition, and the existing entry remains resolvable.
- **Compatibility:** When no ID is given, the Python checkpoint builder
leaves relation_id unset (field presence false), so callers that omit it keep
server-generated IDs. When an ID is given, the builder sets the field to it.
##########
python/pyspark/sql/connect/dataframe.py:
##########
@@ -2417,26 +2418,47 @@ def offset(self, num: int) -> ParentDataFrame:
return DataFrame(plan.Offset(child=self._plan, offset=num),
session=self._session)
def checkpoint(self, eager: bool = True) -> ParentDataFrame:
- cmd = plan.Checkpoint(child=self._plan, local=False, eager=eager)
+ relation_id = str(uuid.uuid4())
+ cmd = plan.Checkpoint(
+ child=self._plan,
+ local=False,
+ eager=eager,
+ relation_id=relation_id,
+ )
_, properties, self._execution_info =
self._session.client.execute_command(
cmd.command(self._session.client)
)
assert "checkpoint_command_result" in properties
- checkpointed = properties["checkpoint_command_result"]
- assert isinstance(checkpointed._plan, plan.CachedRemoteRelation)
- return checkpointed
+ assert properties["checkpoint_command_result"] == relation_id, (
+ "Checkpoint relation ID does not match the requested ID"
+ )
+ return DataFrame(
+ plan=plan.CachedRemoteRelation(relation_id, self._session),
Review Comment:
**Nit (P3):** `test_parity_frame_spark.py` skips `test_checkpoint` and the
`df.checkpoint(False)` doctest is `+SKIP`, so in Connect mode only
`localCheckpoint()` is ever exercised. Each method now has its own copy of ID
generation, propagation, echo check, and result construction, so dropping
`relation_id=relation_id` from the builder call above would make every Connect
`df.checkpoint()` and `psdf.spark.checkpoint()` fail against the in-tree server
without any test noticing.
The binding to `self._session` on this line and at line 2459 is not checked
either, because every checkpoint test uses a single session. Reverting it to
`SparkSession.active()` would hand a second session's checkpoint back bound to
the first session, so its first action fails with `DATAFRAME_NOT_FOUND` and the
release goes to the wrong session. The description mentions smoke checks for
UUID and session retention in both methods, but they were not committed. The
parity skip reason (no SparkContext in Connect) does not apply to suites like
`test_connect_basic.py`, which already reach the local server's SparkContext
through the classic session. This is independent of whether the methods keep
enforcing the echoed ID.
**Verification:**
- **Behavior:** In Connect mode, reliable checkpoint() returns a usable
DataFrame whose relation ID is present in the server session's DataFrame cache,
and the ID is removed after the DataFrame is garbage-collected.
- **Regression:** With another session active, a newly created second
Connect session uses both checkpoint variants. Each variant caches its result
in the second session and returns a DataFrame bound to it that executes
successfully, so binding to SparkSession.active() instead fails the test.
--
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]
