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 5d670ecdf1d1a..db2e64c558f84 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 @@ -1176,9 +1176,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 a53e04da19d41..8efdcf0d6c6af 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 @@ -142,4 +142,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) + } + }