[SPARK-45204][CONNECT] Add optional ExecuteHolder to SparkConnectPlanner #43311
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -86,14 +86,29 @@ final case class InvalidCommandInput( | |
| private val cause: Throwable = null) | ||
| extends Exception(message, cause) | ||
|
|
||
| class SparkConnectPlanner( | ||
| val sessionHolder: SessionHolder, | ||
| val executeHolderOpt: Option[ExecuteHolder] = None) | ||
| extends Logging { | ||
|
|
||
| def this(executeHolder: ExecuteHolder) = { | ||
| this(executeHolder.sessionHolder, Some(executeHolder)) | ||
| } | ||
|
|
||
| if (!executeHolderOpt.forall { e => e.sessionHolder == sessionHolder }) { | ||
| throw new IllegalArgumentException("executeHolder does not belong to sessionHolder") | ||
| } | ||
|
|
||
| private[connect] def session: SparkSession = sessionHolder.session | ||
|
|
||
| private[connect] def userId: String = sessionHolder.userId | ||
|
|
||
| private[connect] def sessionId: String = sessionHolder.sessionId | ||
|
|
||
| lazy val executeHolder = executeHolderOpt.getOrElse { | ||
| throw new IllegalArgumentException("executeHolder is not set") | ||
| } | ||
|
|
||
| private lazy val pythonExec = | ||
| sys.env.getOrElse("PYSPARK_PYTHON", sys.env.getOrElse("PYSPARK_DRIVER_PYTHON", "python3")) | ||
|
|
||
|
|
@@ -2461,48 +2476,39 @@ class SparkConnectPlanner(val sessionHolder: SessionHolder) extends Logging { | |
|
|
||
| def process( | ||
| command: proto.Command, | ||
| responseObserver: StreamObserver[ExecutePlanResponse]): Unit = { | ||
|
Comment on lines
2477
to
+2479
Contributor
There was a problem hiding this comment. Conceptually, this shouldn't be an issue because it's not really a public API. but it's not backwards compatible. This is particularly interesting because you went through some length to fix this problem in the constructor though. @hvanhovell Do we consider the SparkConenct planner to be public? Probably, not correct?
Contributor
There was a problem hiding this comment. FWIW, between Spark 3.4 and Spark 3.5, these interfaces also changed.
Contributor
There was a problem hiding this comment. No, I don't think we should. In the end that will hamper our ability to evolve the interface. |
||
| command.getCommandTypeCase match { | ||
| case proto.Command.CommandTypeCase.REGISTER_FUNCTION => | ||
| handleRegisterUserDefinedFunction(command.getRegisterFunction) | ||
| case proto.Command.CommandTypeCase.REGISTER_TABLE_FUNCTION => | ||
| handleRegisterUserDefinedTableFunction(command.getRegisterTableFunction) | ||
| case proto.Command.CommandTypeCase.WRITE_OPERATION => | ||
| handleWriteOperation(command.getWriteOperation) | ||
| case proto.Command.CommandTypeCase.CREATE_DATAFRAME_VIEW => | ||
| handleCreateViewCommand(command.getCreateDataframeView) | ||
| case proto.Command.CommandTypeCase.WRITE_OPERATION_V2 => | ||
| handleWriteOperationV2(command.getWriteOperationV2) | ||
| case proto.Command.CommandTypeCase.EXTENSION => | ||
| handleCommandPlugin(command.getExtension) | ||
| case proto.Command.CommandTypeCase.SQL_COMMAND => | ||
| handleSqlCommand(command.getSqlCommand, responseObserver) | ||
| case proto.Command.CommandTypeCase.WRITE_STREAM_OPERATION_START => | ||
| handleWriteStreamOperationStart(command.getWriteStreamOperationStart, responseObserver) | ||
| case proto.Command.CommandTypeCase.STREAMING_QUERY_COMMAND => | ||
| handleStreamingQueryCommand(command.getStreamingQueryCommand, responseObserver) | ||
| case proto.Command.CommandTypeCase.STREAMING_QUERY_MANAGER_COMMAND => | ||
| handleStreamingQueryManagerCommand( | ||
| command.getStreamingQueryManagerCommand, | ||
| responseObserver) | ||
| case proto.Command.CommandTypeCase.GET_RESOURCES_COMMAND => | ||
| handleGetResourcesCommand(responseObserver) | ||
| case _ => throw new UnsupportedOperationException(s"$command not supported.") | ||
| } | ||
| } | ||
|
|
||
| def handleSqlCommand( | ||
| getSqlCommand: SqlCommand, | ||
| responseObserver: StreamObserver[ExecutePlanResponse]): Unit = { | ||
| // Eagerly execute commands of the provided SQL string. | ||
| val args = getSqlCommand.getArgsMap | ||
| val namedArguments = getSqlCommand.getNamedArgumentsMap | ||
|
|
@@ -2600,8 +2606,7 @@ class SparkConnectPlanner(val sessionHolder: SessionHolder) extends Logging { | |
| } | ||
|
|
||
| private def handleRegisterUserDefinedFunction( | ||
| fun: proto.CommonInlineUserDefinedFunction): Unit = { | ||
| fun.getFunctionCase match { | ||
| case proto.CommonInlineUserDefinedFunction.FunctionCase.PYTHON_UDF => | ||
| handleRegisterPythonUDF(fun) | ||
|
|
@@ -2617,8 +2622,7 @@ class SparkConnectPlanner(val sessionHolder: SessionHolder) extends Logging { | |
| } | ||
|
|
||
| private def handleRegisterUserDefinedTableFunction( | ||
| fun: proto.CommonInlineUserDefinedTableFunction): Unit = { | ||
| fun.getFunctionCase match { | ||
| case proto.CommonInlineUserDefinedTableFunction.FunctionCase.PYTHON_UDTF => | ||
| val function = createPythonUserDefinedTableFunction(fun) | ||
|
|
@@ -2685,7 +2689,7 @@ class SparkConnectPlanner(val sessionHolder: SessionHolder) extends Logging { | |
| session.udf.register(fun.getFunctionName, udf) | ||
| } | ||
|
|
||
| private def handleCommandPlugin(extension: ProtoAny): Unit = { | ||
| SparkConnectPluginRegistry.commandRegistry | ||
| // Lazily traverse the collection. | ||
| .view | ||
|
|
@@ -2698,9 +2702,7 @@ class SparkConnectPlanner(val sessionHolder: SessionHolder) extends Logging { | |
| executeHolder.eventsManager.postFinished() | ||
| } | ||
|
|
||
| private def handleCreateViewCommand(createView: proto.CreateDataFrameViewCommand): Unit = { | ||
| val viewType = if (createView.getIsGlobal) GlobalTempView else LocalTempView | ||
|
|
||
| val tableIdentifier = | ||
|
|
@@ -2736,9 +2738,7 @@ class SparkConnectPlanner(val sessionHolder: SessionHolder) extends Logging { | |
| * | ||
| * @param writeOperation | ||
| */ | ||
| private def handleWriteOperation(writeOperation: proto.WriteOperation): Unit = { | ||
| // Transform the input plan into the logical plan. | ||
| val plan = transformRelation(writeOperation.getInput) | ||
| // And create a Dataset from the plan. | ||
|
|
@@ -2810,9 +2810,7 @@ class SparkConnectPlanner(val sessionHolder: SessionHolder) extends Logging { | |
| * | ||
| * @param writeOperation | ||
| */ | ||
| def handleWriteOperationV2(writeOperation: proto.WriteOperationV2): Unit = { | ||
| // Transform the input plan into the logical plan. | ||
| val plan = transformRelation(writeOperation.getInput) | ||
| // And create a Dataset from the plan. | ||
|
|
@@ -2873,8 +2871,7 @@ class SparkConnectPlanner(val sessionHolder: SessionHolder) extends Logging { | |
|
|
||
| def handleWriteStreamOperationStart( | ||
| writeOp: WriteStreamOperationStart, | ||
| responseObserver: StreamObserver[ExecutePlanResponse]): Unit = { | ||
| val plan = transformRelation(writeOp.getInput) | ||
| val tracker = executeHolder.eventsManager.createQueryPlanningTracker | ||
| val dataset = Dataset.ofRows(session, plan, tracker) | ||
|
|
@@ -2999,8 +2996,7 @@ class SparkConnectPlanner(val sessionHolder: SessionHolder) extends Logging { | |
|
|
||
| def handleStreamingQueryCommand( | ||
| command: StreamingQueryCommand, | ||
| responseObserver: StreamObserver[ExecutePlanResponse]): Unit = { | ||
|
|
||
| val id = command.getQueryId.getId | ||
| val runId = command.getQueryId.getRunId | ||
|
|
@@ -3177,8 +3173,7 @@ class SparkConnectPlanner(val sessionHolder: SessionHolder) extends Logging { | |
|
|
||
| def handleStreamingQueryManagerCommand( | ||
| command: StreamingQueryManagerCommand, | ||
| responseObserver: StreamObserver[ExecutePlanResponse]): Unit = { | ||
| val respBuilder = StreamingQueryManagerCommandResult.newBuilder() | ||
|
|
||
| command.getCommandCase match { | ||
|
|
@@ -3257,8 +3252,7 @@ class SparkConnectPlanner(val sessionHolder: SessionHolder) extends Logging { | |
| } | ||
|
|
||
| def handleGetResourcesCommand( | ||
| responseObserver: StreamObserver[proto.ExecutePlanResponse]): Unit = { | ||
| executeHolder.eventsManager.postFinished() | ||
| responseObserver.onNext( | ||
| proto.ExecutePlanResponse | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why not just pass in the ExecuteHolder, and create a getter for SessionHolder? IMO it is a bit weird to create this structure for just a single unit test.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
it's used in AnalyzePlan that doesn't have an ExecuteHolder
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ok, makes sense.