Changes based on feedback.

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

@@ -627,13 +627,19 @@ object CollapseWindow extends Rule[LogicalPlan] {
 
 /**
  * Transpose Adjacent Window Expressions.
- * - If the partition spec of the parent Window expression is a subset of the partition spec
+ * - If the partition spec of the parent Window expression is compatible with the partition spec
  *   of the child window expression, transpose them.
  */
 object TransposeWindow extends Rule[LogicalPlan] {
+  private def compatibleParititions(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)
+    })
+  }
+
   def apply(plan: LogicalPlan): LogicalPlan = plan transformUp {
     case w1 @ Window(we1, ps1, os1, w2 @ Window(we2, ps2, os2, grandChild))
-        if ps1.length < ps2.length && ps2.containsSlice(ps1) =>
+        if w1.references.intersect(w2.windowOutputSet).isEmpty && compatibleParititions(ps1, ps2) =>
       Project(w1.output, Window(we2, ps2, os2, Window(we1, ps1, os1, grandChild)))
   }
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/TransposeWindowSuite.scala

@@ -42,32 +42,12 @@ class TransposeWindowSuite extends PlanTest {
   val partitionSpec1 = Seq(a)
   val partitionSpec2 = Seq(a, b)
   val partitionSpec3 = Seq(d)
+  val partitionSpec4 = Seq(b, a, d)
 
   val orderSpec1 = Seq(d.asc)
   val orderSpec2 = Seq(d.desc)
 
-  test("flip two adjacent windows with compatible partitions in multiple selects") {
-    val wexpr1 = windowExpr(sum('c), windowSpec(partitionSpec2, Seq.empty, UnspecifiedFrame))
-    val wexpr2 = windowExpr(sum('c), windowSpec(partitionSpec1, Seq.empty, UnspecifiedFrame))
-
-    val query = testRelation
-      .select('a, 'b, 'c, wexpr1.as('sum_a_2))
-      .select('a, 'b, 'c, 'sum_a_2, wexpr2.as('sum_a_1))
-
-    val optimized = Optimize.execute(query.analyze)
-
-    val query2 = testRelation
-      .select('a, 'b, 'c)
-      .select('a, 'b, 'c, wexpr2.as('sum_a_1))
-      .select('a, 'b, 'c, wexpr1.as('sum_a_2), 'sum_a_1)
-      .select('a, 'b, 'c, 'sum_a_2, 'sum_a_1)
-
-    val correctAnswer = Optimize.execute(query2.analyze)
-
-    comparePlans(optimized, correctAnswer)
-  }
-
-  test("flip two adjacent windows with compatible partitions") {
+  test("transpose two adjacent windows with compatible partitions") {
     val query = testRelation
       .window(Seq(sum(c).as('sum_a_2)), partitionSpec2, orderSpec2)
       .window(Seq(sum(c).as('sum_a_1)), partitionSpec1, orderSpec1)
@@ -83,19 +63,31 @@ class TransposeWindowSuite extends PlanTest {
     comparePlans(optimized, correctAnswer.analyze)
   }
 
-  test("don't flip two adjacent windows with incompatible partitions") {
+  test("transpose two adjacent windows with differently ordered compatible partitions") {
     val query = testRelation
-      .window(Seq(sum(c).as('sum_a_2)), partitionSpec3, Seq.empty)
-      .window(Seq(sum(c).as('sum_a_1)), partitionSpec1, Seq.empty)
+      .window(Seq(sum(c).as('sum_a_2)), partitionSpec4, Seq.empty)
+      .window(Seq(sum(c).as('sum_a_1)), partitionSpec2, Seq.empty)
 
     val analyzed = query.analyze
     val optimized = Optimize.execute(analyzed)
 
     val correctAnswer = testRelation
+      .window(Seq(sum(c).as('sum_a_1)), partitionSpec2, Seq.empty)
+      .window(Seq(sum(c).as('sum_a_2)), partitionSpec4, Seq.empty)
+      .select('a, 'b, 'c, 'd, 'sum_a_2, 'sum_a_1)
+
+    comparePlans(optimized, correctAnswer.analyze)
+  }
+
+  test("don't transpose two adjacent windows with incompatible partitions") {
+    val query = testRelation
       .window(Seq(sum(c).as('sum_a_2)), partitionSpec3, Seq.empty)
       .window(Seq(sum(c).as('sum_a_1)), partitionSpec1, Seq.empty)
 
-    comparePlans(optimized, correctAnswer.analyze)
+    val analyzed = query.analyze
+    val optimized = Optimize.execute(analyzed)
+
+    comparePlans(optimized, analyzed)
   }
 
 }