diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala index c0f6e75ae4bfa..2415646acbbd9 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala @@ -1257,9 +1257,9 @@ object CollapseWindow extends Rule[LogicalPlan] { */ object TransposeWindow extends Rule[LogicalPlan] { private def compatiblePartitions(ps1 : Seq[Expression], ps2: Seq[Expression]): Boolean = { - ps1.length < ps2.length && ps2.take(ps1.length).permutations.exists(ps1.zip(_).forall { - case (l, r) => l.semanticEquals(r) - }) + ps1.length < ps2.length && ps1.forall { expr1 => + ps2.exists(expr1.semanticEquals) + } } private def windowsCompatible(w1: Window, w2: Window): Boolean = { diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/TransposeWindowSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/TransposeWindowSuite.scala index 60869d7f94842..a9796141c0c7e 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/TransposeWindowSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/TransposeWindowSuite.scala @@ -143,4 +143,21 @@ class TransposeWindowSuite extends PlanTest { comparePlans(optimized, analyzed) } + test("SPARK-38034: transpose two adjacent windows with compatible partitions " + + "which is not a prefix") { + val query = testRelation + .window(Seq(sum(c).as('sum_a_2)), partitionSpec4, orderSpec2) + .window(Seq(sum(c).as('sum_a_1)), partitionSpec3, orderSpec1) + + val analyzed = query.analyze + val optimized = Optimize.execute(analyzed) + + val correctAnswer = testRelation + .window(Seq(sum(c).as('sum_a_1)), partitionSpec3, orderSpec1) + .window(Seq(sum(c).as('sum_a_2)), partitionSpec4, orderSpec2) + .select('a, 'b, 'c, 'd, 'sum_a_2, 'sum_a_1) + + comparePlans(optimized, correctAnswer.analyze) + } + }