cloud-fan opened a new pull request, #59131:
URL: https://github.com/apache/spark/pull/59131
### What changes were proposed in this pull request?
Spark Connect plans each `withColumns` / `withColumnsRenamed` call as a lazy
`Project(Seq(UnresolvedStarWithColumns(...)), child)` (SPARK-50601), so a loop
of `df = df.withColumn(...)` becomes a stack of such Projects. The fixed-point
analyzer resolves a Project only after its child is resolved, and resolving one
Project takes several rules (column resolution, function resolution, type
coercion). So it resolves at most one Project of the chain per iteration of the
`Resolution` batch, and a chain of about 100 calls hits
`spark.sql.analyzer.maxIterations`. SPARK-53505 worked around this in the
Connect planner: it analyzes each `withColumns` eagerly, and falls back to the
lazy plan when that fails.
This PR fixes the analyzer and removes the workaround:
1. When `ResolveReferences` expands the star of a lazily planned
`withColumns` whose child is resolved, it resolves that Project right away:
- It resolves the column references against the child, the same way it
does for any Project. This covers names, Spark Connect `df["col"]` plan ids,
metadata columns, outer references, and SQL variables.
- It runs only the `Resolution` batch on the expressions that are still
unresolved, over a `LocalRelation` with the same output as the child. Then it
puts the resolved expressions back into the Project.
The cost depends on the width of the Project, not the depth of its child.
The parent Project is then resolved in the same traversal, so the whole chain
is resolved in one pass.
2. Some expressions can't be resolved this way, and those Projects are left
to the normal iterations, so results and errors stay the same. This covers:
- generators and window functions, which need the plan to be rewritten;
- subqueries;
- stars inside expressions;
- lateral column aliases;
- errors.
3. A new internal config, `spark.sql.analyzer.resolveWithColumnsInIsolation`
(default `true`), turns this off.
4. The Spark Connect planner plans `withColumns` and `withColumnsRenamed`
lazily again, without the eager analysis added by SPARK-53505.
### Why are the changes needed?
The eager analysis in the Connect planner is a workaround with its own
problems:
- **It is quadratic.** Each call analyzes the plan built so far again. In a
local benchmark of 1000 chained `withColumn` calls, analysis took about 50-65 s
with the eager analysis, and about 7-9 s with this PR:
| 1000 chained calls | eager analysis | this PR |
|---|---|---|
| `abs(id)` | 52 s | 7.3 s |
| `abs(prev + 1)` | 64 s | 8.8 s |
| `abs(df[prev] + 1)` (plan-id reference) | n/a | 7.7 s |
- **It binds the plan when it's built, not when it's analyzed.** For
example, in Spark Declarative Pipelines, take a flow
`spark.read.table("src").withColumn(...)` where `src` is a dataset that already
exists in the catalog from a previous run. The flow is bound to the existing
table and loses its dependency on `src` (see #59100).
### Does this PR introduce _any_ user-facing change?
Yes, in Spark Connect:
- Long chains of `withColumn` / `withColumns` / `withColumnRenamed` calls
are analyzed faster.
- `withColumns` and `withColumnsRenamed` no longer bind their input when the
plan is built. For example, a pipeline flow that calls them on a dataset of the
pipeline now keeps its dependency on that dataset.
Results and error messages are unchanged.
### How was this patch tested?
New tests:
- `AnalysisSuite`:
- A chain of 120 lazily planned `withColumns` fails with max iterations
without the change. With the change, it resolves to the same plan as the normal
iterations.
- `withColumns` that can't be resolved in isolation still resolve to the
same plans as the normal iterations. The cases are a generator, a window
function, a lateral column alias, a subquery, a star, and a mix with
`withColumnsRenamed`.
- Errors are the same as with the normal iterations.
- `SparkConnectPlannerSuite`: a chain of 200 `withColumns` is planned
lazily, analyzed, and returns the correct result. It references columns both by
name and through `df[...]`.
Existing suites run locally:
- catalyst: `org.apache.spark.sql.catalyst.analysis.*`
- Connect: `SparkConnectPlannerSuite`, `SparkConnectProtoSuite`
- sql/core: `DataFrameSuite`, `DataFrameSubquerySuite`,
`ColumnExpressionSuite`, `DataFrameWindowFunctionsSuite`,
`GeneratorFunctionSuite`, `UDFSuite`, `DataFrameSelfJoinSuite`,
`FileMetadataStructSuite`, `MetadataColumnSuite`, `DatasetSuite`
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5.5)
This pull request and its description were written by Isaac.
--
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]
