Skip to content

Commit 714884c

Browse files
committed
[SPARK-51625][SQL] Command in CTE relations should trigger inline
### What changes were proposed in this pull request? This PR fixes a small issue in the rule `CTESubstitution`. The intention is to fall back to CTE inline if there are commands in the plan. However, the code simply does `plan.collect` which misses the fact that `UnresolvedWith` does not include the CTE relations in its children. To fix it, we should do a manually recursion to handle the CTE relations. ### Why are the changes needed? A potential bug as the code does not work as expected. ### Does this PR introduce _any_ user-facing change? No, there is no user-facing API that can construct such a CTE plan. ### How was this patch tested? a new test ### Was this patch authored or co-authored using generative AI tooling? no Closes #50425 from cloud-fan/cte. Authored-by: Wenchen Fan <wenchen@databricks.com> Signed-off-by: Wenchen Fan <wenchen@databricks.com>
1 parent 0f36edc commit 714884c

File tree

2 files changed

+20
-2
lines changed

2 files changed

+20
-2
lines changed

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CTESubstitution.scala

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,9 +56,15 @@ object CTESubstitution extends Rule[LogicalPlan] {
5656
return plan
5757
}
5858

59-
val commands = plan.collect {
60-
case c @ (_: Command | _: ParsedStatement | _: InsertIntoDir) => c
59+
def collectCommands(p: LogicalPlan): Seq[LogicalPlan] = p match {
60+
case c @ (_: Command | _: ParsedStatement | _: InsertIntoDir) => Seq(c)
61+
case u: UnresolvedWith =>
62+
collectCommands(u.child) ++ u.cteRelations.flatMap {
63+
case (_, relation) => collectCommands(relation)
64+
}
65+
case p => p.children.flatMap(collectCommands)
6166
}
67+
val commands = collectCommands(plan)
6268
val forceInline = if (commands.length == 1) {
6369
if (conf.getConf(SQLConf.LEGACY_INLINE_CTE_IN_COMMANDS)) {
6470
// The legacy behavior always inlines the CTE relations for queries in commands.

sql/core/src/test/scala/org/apache/spark/sql/CTEInlineSuite.scala

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
package org.apache.spark.sql
1919

20+
import org.apache.spark.sql.catalyst.analysis.{CurrentNamespace, UnresolvedRelation}
2021
import org.apache.spark.sql.catalyst.expressions.{Alias, And, GreaterThan, LessThan, Literal, Or, Rand}
2122
import org.apache.spark.sql.catalyst.optimizer.InlineCTE
2223
import org.apache.spark.sql.catalyst.plans.logical._
@@ -834,6 +835,17 @@ abstract class CTEInlineSuiteBase
834835
}
835836
}
836837
}
838+
839+
test("SPARK-51625: command in CTE relations should trigger inline") {
840+
val plan = UnresolvedWith(
841+
child = UnresolvedRelation(Seq("t")),
842+
cteRelations = Seq("t" -> SubqueryAlias("t", ShowTables(CurrentNamespace, pattern = None)))
843+
)
844+
assert(!spark.sessionState.analyzer.execute(plan).exists {
845+
case _: WithCTE => true
846+
case _ => false
847+
})
848+
}
837849
}
838850

839851
class CTEInlineSuiteAEOff extends CTEInlineSuiteBase with DisableAdaptiveExecutionSuite

0 commit comments

Comments
 (0)