This is an automated email from the ASF dual-hosted git repository. ptoth pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push: new 3b405948ee4 [SPARK-44934][SQL] Use outputSet instead of output to check if column pruning occurred in PushdownPredicateAndPruneColumnsForCTEDef 3b405948ee4 is described below commit 3b405948ee47702e5a7250dc27430836145b0e19 Author: Wen Yuen Pang <wenyuen.p...@databricks.com> AuthorDate: Thu Aug 24 17:00:04 2023 +0200 [SPARK-44934][SQL] Use outputSet instead of output to check if column pruning occurred in PushdownPredicateAndPruneColumnsForCTEDef ### What changes were proposed in this pull request? Originally, when a CTE has duplicate expression IDs in its output, the rule PushdownPredicatesAndPruneColumnsForCTEDef wrongly assesses that the columns in the CTE were pruned, as it compares the size of the attribute set containing the union of columns (which is unique) and the original output of the CTE (which contains duplicate columns) and notices that the former is less than the latter. This causes incorrect pruning of the CTE output, resulting in a missing reference and causing [...] This PR changes the logic to use the needsPruning function to assess whether a CTE has been pruned, which uses the outputSet to check if any columns has been pruned instead of the output. ### Why are the changes needed? The incorrect behaviour of PushdownPredicatesAndPruneColumnsForCTEDef in CTEs with duplicate expression IDs in its output causes a crash when such a query is run. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Unit test for the crashing case was added. ### Was this patch authored or co-authored using generative AI tooling? No Closes #42635 from wenyuen-db/SPARK-44934. Authored-by: Wen Yuen Pang <wenyuen.p...@databricks.com> Signed-off-by: Peter Toth <peter.t...@gmail.com> --- ...ushdownPredicatesAndPruneColumnsForCTEDef.scala | 2 +- .../org/apache/spark/sql/CTEInlineSuite.scala | 33 ++++++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala index f351ba0b39a..e643a1af363 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala @@ -143,7 +143,7 @@ object PushdownPredicatesAndPruneColumnsForCTEDef extends Rule[LogicalPlan] { case cteRef @ CTERelationRef(cteId, _, output, _) => val (cteDef, _, _, newAttrSet) = cteMap(cteId) - if (newAttrSet.size < output.size) { + if (needsPruning(cteDef.child, newAttrSet)) { val indices = newAttrSet.toSeq.map(cteDef.output.indexOf) val newOutput = indices.map(output) cteRef.copy(output = newOutput) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CTEInlineSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CTEInlineSuite.scala index eb141eaf3a9..5f6c4479265 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CTEInlineSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CTEInlineSuite.scala @@ -645,6 +645,39 @@ abstract class CTEInlineSuiteBase } } } + + test("SPARK-44934: CTE column pruning handles duplicate exprIds in CTE") { + withTempView("t") { + Seq((0, 1, 2), (1, 2, 3)).toDF("c1", "c2", "c3").createOrReplaceTempView("t") + val query = + """ + |with cte as ( + | select c1, c1, c2, c3 from t where random() > 0 + |) + |select cte.c1, cte2.c1, cte.c2, cte2.c3 from + | (select c1, c2 from cte) cte + | inner join + | (select c1, c3 from cte) cte2 + | on cte.c1 = cte2.c1 + """.stripMargin + + val df = sql(query) + checkAnswer(df, Row(0, 0, 1, 2) :: Row(1, 1, 2, 3) :: Nil) + assert( + df.queryExecution.analyzed.collect { + case WithCTE(_, cteDefs) => cteDefs + }.head.length == 1, + "With-CTE should contain 1 CTE def after analysis.") + val cteRepartitions = df.queryExecution.optimizedPlan.collect { + case r: RepartitionOperation => r + } + assert(cteRepartitions.length == 2, + "CTE should not be inlined after optimization.") + assert(cteRepartitions.head.collectFirst { + case p: Project if p.projectList.length == 4 => p + }.isDefined, "CTE columns should not be pruned.") + } + } } class CTEInlineSuiteAEOff extends CTEInlineSuiteBase with DisableAdaptiveExecutionSuite --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@spark.apache.org For additional commands, e-mail: commits-h...@spark.apache.org