zhengruifeng commented on code in PR #59540:
URL: https://github.com/apache/spark/pull/59540#discussion_r4343608850
##########
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:
Addressed in 796ca7a45ff. Client-supplied duplicate IDs are now rejected
before relation planning or checkpoint work with
CONNECT_INVALID_PLAN.DATAFRAME_ALREADY_EXISTS. The final atomic insertion
reports the same condition if it loses a race, while server-generated and
foreachBatch collisions retain INTERNAL_ERROR. Added admission, race, and
cache-preservation coverage in the existing planner and session-holder suites.
##########
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:
Addressed in 796ca7a45ff. Strict ID enforcement is intentional. The bare
assert is replaced with PySparkAssertionError, and a mismatched server-returned
relation is explicitly removed before raising. The server requirement is
documented, and the mismatch cleanup test also passes under python -O.
##########
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:
Addressed in 796ca7a45ff. Added committed coverage in existing suites for
supplied, omitted, empty, duplicate, and racing IDs, including cache
preservation. Added Python client coverage that verifies relation_id is absent
when unset and present when supplied. The previously removed standalone Scala
suite and test_connect_plan changes remain absent.
##########
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:
Addressed in 796ca7a45ff. Added Connect integration coverage for reliable
checkpoint cleanup and for both checkpoint variants using their originating
session while another session is active. The duplicated Python implementation
is also consolidated into a shared _checkpoint helper.
--
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]
