sarutak commented on code in PR #59131:
URL: https://github.com/apache/spark/pull/59131#discussion_r4131158395
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala:
##########
@@ -677,7 +677,18 @@ class Analyzer(
ResolveTranspiledPythonUDFOptions) ++
Seq(ResolveUpdateEventTimeWatermarkColumn) ++
extendedResolutionRules ++
- Seq(NameStreamingSources) : _*),
+ Seq(NameStreamingSources) : _*)
+
+ /**
+ * Runs only the "Resolution" batch. [[ResolveReferences]] uses it to
resolve a lazily planned
+ * `withColumns` in isolation from its child.
+ */
+ private lazy val resolutionBatchExecutor = new RuleExecutor[LogicalPlan] {
Review Comment:
`resolutionBatch` is declared to return `Batch`, so the
`asInstanceOf[Batch]` in `resolutionBatchExecutor.batches` looks redundant. If
it is needed because of the path-dependent `Batch` type across the two
`RuleExecutor` instances, a short comment would explain it; otherwise it can be
dropped.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala:
##########
@@ -1032,6 +1032,75 @@ class AnalysisSuite extends AnalysisTest with Matchers {
}
}
+ private def withColumns(child: LogicalPlan, name: String, expr: Expression):
LogicalPlan =
+ Project(Seq(UnresolvedStarWithColumns(Seq(name), Seq(expr))), child)
+
+ private def fn(name: String, args: Expression*): UnresolvedFunction =
+ UnresolvedFunction(name, args, isDistinct = false)
+
+ // Analyzes the plan with the lazily planned `withColumns` resolved in
isolation, and checks
+ // that the result is the same as resolving them in the normal iterations.
+ private def checkWithColumnsResolution(plan: LogicalPlan): Unit = {
+ val expected = withSQLConf(
+ SQLConf.ANALYZER_RESOLVE_WITH_COLUMNS_IN_ISOLATION.key -> "false",
+ SQLConf.ANALYZER_MAX_ITERATIONS.key -> "1000") {
+ getAnalyzer.executeAndCheck(plan, new QueryPlanningTracker)
+ }
+ comparePlans(getAnalyzer.executeAndCheck(plan, new QueryPlanningTracker),
expected)
+ }
+
+ test("SPARK-59855: resolve a long chain of lazily planned withColumns") {
+ // Like a chain of Spark Connect `withColumn` calls. Each call adds a
column that references
+ // the column added by the previous call and takes several rules to be
resolved.
+ val chain = (1 to 120).foldLeft[LogicalPlan](testRelation) { (child, i) =>
+ val prev = UnresolvedAttribute(if (i == 1) "a" else s"c${i - 1}")
+ withColumns(child, s"c$i", fn("abs", prev + 1L))
+ }
+ withSQLConf(SQLConf.ANALYZER_RESOLVE_WITH_COLUMNS_IN_ISOLATION.key ->
"false") {
+ // The normal iterations resolve at most one `withColumns` of the chain
per iteration.
+ val message =
intercept[RuntimeException](getAnalyzer.execute(chain)).getMessage
+ assert(message.startsWith("Max iterations (100) reached for batch
Resolution"))
+ }
+ checkWithColumnsResolution(chain)
+ }
+
+ test("SPARK-59855: lazily planned withColumns that can't be resolved in
isolation") {
Review Comment:
The added tests cover the 120/200 chains, the six "cannot resolve in
isolation" cases, and error parity. How about adding cases that exercise the
resolve-in-child vs. fall-back-to-normal split explicitly?
- a `withColumns` referencing a **metadata column** (e.g.
`_metadata.file_path`);
- a Connect **plan-id reference** (`df[...]`) inside a chain (the
description says this is resolved by `resolveReferencesInProject`; a test would
pin it);
- a **SQL variable** reference;
- a **streaming** chain (the `isStreaming = p.child.isStreaming` path on the
leaf);
- optionally, a `TempResolvedColumn`-like case, to confirm the isolated run
never marks a not-fully-resolved expression as resolved. (In practice
`TempResolvedColumn` is only produced under Aggregate/Filter/Sort, not under a
`Project` over a `LocalRelation`, so this is likely already impossible but a
test would document that.)
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala:
##########
@@ -1997,6 +2010,60 @@ class Analyzer(
q.mapExpressions(resolveExpressionByPlanChildren(_, q,
includeLastResort = true))
}
+ private def resolveReferencesInProject(p: Project): Project = {
+ val resolvedBasic = p.projectList.map(resolveExpressionByPlanChildren(_,
p))
+ // Lateral column alias has higher priority than outer reference.
+ val resolvedWithLCA = resolveLateralColumnAlias(resolvedBasic)
+ val resolvedFinal = resolvedWithLCA.map(resolveColsLastResort)
+ p.copy(projectList =
+ resolvedFinal.map(e =>
aliasIfOuterReference(e.asInstanceOf[NamedExpression])))
+ }
+
+ /**
+ * Resolves a `withColumns` whose star has just been expanded.
+ *
+ * A lazily planned chain of `withColumns`, like the one Spark Connect
creates for a chain of
+ * `withColumn` calls, is a stack of Projects. A Project waits for its
child to be resolved,
+ * and then takes several rules to be resolved itself, so the normal
iterations resolve at most
+ * one Project of the chain per iteration, and a long chain hits the max
iterations limit.
+ * Instead, we resolve the Project right away: resolve its column
references against the child
+ * as usual, then run the "Resolution" batch on the expressions that are
still unresolved,
+ * over a leaf node with the same output as the child. This costs time
proportional to the
+ * width of the Project instead of the size of the child, and lets the
parent Project be
+ * resolved in the same pass. Expressions that can't be resolved this way,
e.g. a window
+ * function that needs the plan to be rewritten, are left to the normal
iterations.
+ */
+ private def resolveWithColumns(p: Project): Project = {
+ if (containsStar(p.projectList) || containsDeserializer(p.projectList) ||
+ p.projectList.exists(SubqueryExpression.hasSubquery)) {
+ p
+ } else {
+ val withReferences = resolveReferencesInProject(p)
+ val unresolved = withReferences.projectList.filterNot(_.resolved)
+ // Lateral column alias references need the whole project list to be
resolved.
+ if (unresolved.isEmpty ||
+
unresolved.exists(_.containsPattern(LATERAL_COLUMN_ALIAS_REFERENCE))) {
+ withReferences
+ } else {
+ val leaf = LocalRelation(p.child.output, isStreaming =
p.child.isStreaming)
+ try {
+ resolutionBatchExecutor.execute(Project(unresolved, leaf)) match {
+ case r @ Project(resolvedList, l) if (l eq leaf) && r.resolved &&
+ r.missingInput.isEmpty && resolvedList.length ==
unresolved.length =>
+ val resolvedIter = resolvedList.iterator
+ withReferences.copy(projectList =
withReferences.projectList.map { e =>
+ if (e.resolved) e else resolvedIter.next()
+ })
+ case _ => withReferences
+ }
+ } catch {
+ // Leave the error to the normal iterations, so that it's reported
as usual.
+ case _: AnalysisException => withReferences
Review Comment:
`resolveWithColumns` wraps the isolated run in `catch { case _:
AnalysisException => withReferences }`. But when a batch exceeds
`maxIterations`, `RuleExecutor.execute` throws a plain `RuntimeException`, not
an `AnalysisException`.
https://github.com/apache/spark/blob/b03b53dcacb42c9c31c0cd2b622d495d60c3f59f/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleExecutor.scala#L299
The `Resolution` batch uses `fixedPoint` with `errorOnExceed = true`, and
under `Utils.isTesting` the throw happens unconditionally. So if
`resolutionBatchExecutor.execute(Project(unresolved, leaf))` does not converge
within `maxIterations`, the `RuntimeException` escapes the `catch` and
propagates to the user, instead of falling back to the normal iterations as the
design intends. The same applies to non-`AnalysisException` errors that
coercion rules can raise (e.g. `SparkException`).
The isolated run is a best-effort optimization whose failures should always
fall back to the normal path. Suggest widening the catch:
```scala
catch { case NonFatal(_) => withReferences }
```
(or at least adding `RuntimeException`). Worth a test that forces the
isolated run to not converge and asserts it falls back rather than surfacing a
raw `RuntimeException`. Note the fallback here still runs the normal
iterations, which for a genuinely divergent case may itself hit `maxIterations`
but that is the pre-existing behavior and is the correct thing to reproduce.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala:
##########
@@ -1997,6 +2010,60 @@ class Analyzer(
q.mapExpressions(resolveExpressionByPlanChildren(_, q,
includeLastResort = true))
}
+ private def resolveReferencesInProject(p: Project): Project = {
+ val resolvedBasic = p.projectList.map(resolveExpressionByPlanChildren(_,
p))
+ // Lateral column alias has higher priority than outer reference.
+ val resolvedWithLCA = resolveLateralColumnAlias(resolvedBasic)
+ val resolvedFinal = resolvedWithLCA.map(resolveColsLastResort)
+ p.copy(projectList =
+ resolvedFinal.map(e =>
aliasIfOuterReference(e.asInstanceOf[NamedExpression])))
+ }
+
+ /**
+ * Resolves a `withColumns` whose star has just been expanded.
+ *
+ * A lazily planned chain of `withColumns`, like the one Spark Connect
creates for a chain of
+ * `withColumn` calls, is a stack of Projects. A Project waits for its
child to be resolved,
+ * and then takes several rules to be resolved itself, so the normal
iterations resolve at most
+ * one Project of the chain per iteration, and a long chain hits the max
iterations limit.
+ * Instead, we resolve the Project right away: resolve its column
references against the child
+ * as usual, then run the "Resolution" batch on the expressions that are
still unresolved,
+ * over a leaf node with the same output as the child. This costs time
proportional to the
+ * width of the Project instead of the size of the child, and lets the
parent Project be
+ * resolved in the same pass. Expressions that can't be resolved this way,
e.g. a window
+ * function that needs the plan to be rewritten, are left to the normal
iterations.
+ */
+ private def resolveWithColumns(p: Project): Project = {
+ if (containsStar(p.projectList) || containsDeserializer(p.projectList) ||
+ p.projectList.exists(SubqueryExpression.hasSubquery)) {
+ p
+ } else {
+ val withReferences = resolveReferencesInProject(p)
+ val unresolved = withReferences.projectList.filterNot(_.resolved)
+ // Lateral column alias references need the whole project list to be
resolved.
+ if (unresolved.isEmpty ||
+
unresolved.exists(_.containsPattern(LATERAL_COLUMN_ALIAS_REFERENCE))) {
+ withReferences
+ } else {
+ val leaf = LocalRelation(p.child.output, isStreaming =
p.child.isStreaming)
+ try {
+ resolutionBatchExecutor.execute(Project(unresolved, leaf)) match {
+ case r @ Project(resolvedList, l) if (l eq leaf) && r.resolved &&
+ r.missingInput.isEmpty && resolvedList.length ==
unresolved.length =>
+ val resolvedIter = resolvedList.iterator
+ withReferences.copy(projectList =
withReferences.projectList.map { e =>
Review Comment:
The re-splice is correct: `unresolved =
withReferences.projectList.filterNot(_.resolved)` and the `map { if
(e.resolved) e else resolvedIter.next() }` walk the same Seq in the same
direction, so the i-th unresolved element maps to the i-th resolved element,
and `resolvedList.length == unresolved.length` guards the count.
However, this relies on the `Resolution` batch preserving the order and
count of the top-level project list of `Project(unresolved, leaf)` (the count
is guarded, the order is not). A one-line comment stating that dependency would
help, since a future rule that reorders the top-level list would silently break
the mapping.
--
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]
