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
