[SPARK-52813][CONNECT] Allow DAGs in Spark Connect #51516
+1,491
−328
Conversation
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
hvanhovell
commented
Jul 16, 2025
| getSubqueryExpression.getSubqueryType match { | ||
| case proto.SubqueryExpression.SubqueryType.SUBQUERY_TYPE_SCALAR => | ||
| UnresolvedScalarSubqueryPlanId(planId) |
There was a problem hiding this comment.
We can remove the Unresolved*PlanId expressions.
hvanhovell
commented
Jul 16, 2025
| case other => | ||
| throw InvalidInputErrors.sqlCommandExpectsSqlOrWithRelations(other) | ||
| // Only allow a SQL relation or a SQL relation nested in a WithRelations relation. | ||
| if (!relation.hasSql && |
There was a problem hiding this comment.
This check is a bit arbitrary. Do we really care?
hvanhovell
commented
Jul 16, 2025
| } | ||
| val df = Dataset.ofRows(session, transformRelation(relation), tracker) |
There was a problem hiding this comment.
Capture the parse time here, use the transformRelation time as a proxy. Long term we may want to look at the use of QueryPlanTracker in general.
hvanhovell
commented
Jul 16, 2025
| val plan = if (relation.getCommon.hasPlanId) { | ||
| getCachedRelation(relation.getCommon.getPlanId) | ||
| } else { | ||
| transformRelation(relation) |
There was a problem hiding this comment.
I am being paranoid here. I hope this does not happen...
hvanhovell
commented
Jul 16, 2025
| @@ -2881,79 +2890,6 @@ class SparkConnectPlanner( | |||
| } | |||
| } | |||
|
|
|||
| private def isValidSQLWithRefs(query: proto.WithRelations): Boolean = { | |||
There was a problem hiding this comment.
All of this code is replaced by CTEs...
hvanhovell
commented
Jul 16, 2025
| transformSqlWithRefs(getWithRelations) | ||
| // Register the plans in the relation cache, so they can be resolved while | ||
| // transforming the root relation into a LogicalPlan. | ||
| val namedReferences = mutable.Buffer.empty[(String, proto.Relation)] |
There was a problem hiding this comment.
Should we throw if there are duplicate names?
hvanhovell
commented
Jul 16, 2025
| throw InvalidInputErrors.invalidWithRelationReference() | ||
| } | ||
| if (common.hasPlanId) { | ||
| relationCache.put(common.getPlanId, ref) |
There was a problem hiding this comment.
Should we throw if the id is already registered to a different relation?
hvanhovell
changed the title
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
What changes were proposed in this pull request?
This PR adds support to turn connect plans into DAGs. Currently we inline the structure of a reused query fragment during serialization. This can cause massive bloat, increase plan size. This hurts over the wire performance, and it hurts planning performance.
In order to allow DAGs we add reference relation type. This reference relation contains an plan_id that points to a plan defined in an parent/root WithRelations node.
I have created an optimization rule for scala that leverages this functionality by DAGifing a plan before sending it over the wire.
Why are the changes needed?
This changes in this PR can decrease plan size significantly. This benefits over the wire performance, and it allows for faster planning.
Does this PR introduce any user-facing change?
No.
How was this patch tested?
I have added unit tests for the PlanOptimization in PlanOptimizationSuite.
I have added an end to end test to ClientDatasetSuite.
I still need to add a self contained server side test
Was this patch authored or co-authored using generative AI tooling?
No.