Changes based on feedback.

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

@@ -88,7 +88,7 @@ abstract class Optimizer(sessionCatalog: SessionCatalog)
       CollapseRepartition,
       CollapseProject,
       CollapseWindow,
-      FlipWindow,
+      TransposeWindow,
       CombineFilters,
       CombineLimits,
       CombineUnions,
@@ -626,15 +626,15 @@ object CollapseWindow extends Rule[LogicalPlan] {
 }
 
 /**
- * Flip Adjacent Window Expressions.
+ * Transpose Adjacent Window Expressions.
  * - If the partition spec of the parent Window expression is a subset of the partition spec
- *   of the child window expression, flip them.
+ *   of the child window expression, transpose them.
  */
-object FlipWindow extends Rule[LogicalPlan] {
+object TransposeWindow extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transformUp {
     case w1 @ Window(we1, ps1, os1, w2 @ Window(we2, ps2, os2, grandChild))
         if ps2.containsSlice(ps1) =>
-      Window(we2, ps2, os2, Window(we1, ps1, os1, grandChild))
+      Project(w1.output, Window(we2, ps2, os2, Window(we1, ps1, os1, grandChild)))
   }
 }
 

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

@@ -25,11 +25,11 @@ import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.RuleExecutor
 
 
-class FlipWindowSuite extends PlanTest {
+class TransposeWindowSuite extends PlanTest {
   object Optimize extends RuleExecutor[LogicalPlan] {
     val batches =
       Batch("CollapseProject", FixedPoint(100), CollapseProject, RemoveRedundantProject) ::
-      Batch("FlipWindow", Once, CollapseWindow, FlipWindow) :: Nil
+      Batch("FlipWindow", Once, CollapseWindow, TransposeWindow) :: Nil
   }
 
   val testRelation = LocalRelation('a.string, 'b.string, 'c.int, 'd.string)
@@ -51,13 +51,13 @@ class FlipWindowSuite extends PlanTest {
       .select('a, 'b, 'c,  windowExpr(sum('c), windowSpec(partitionSpec2, Seq.empty, UnspecifiedFrame)).as('sum_a_2))
       .select('a, 'b, 'c, 'sum_a_2, windowExpr(sum('c), windowSpec(partitionSpec1, Seq.empty, UnspecifiedFrame)).as('sum_a_1))
 
-    val analyzed = query.analyze
-    val optimized = Optimize.execute(analyzed)
+    val optimized = Optimize.execute(query.analyze)
 
     val query2 = testRelation
-      .select('a, 'b, 'c, windowExpr(sum('c), windowSpec(partitionSpec2, Seq.empty, UnspecifiedFrame)).as('sum_a_2),
-        windowExpr(sum('c), windowSpec(partitionSpec1, Seq.empty, UnspecifiedFrame)).as('sum_a_1)
-      )
+      .select('a, 'b, 'c)
+      .select('a, 'b, 'c,  windowExpr(sum('c), windowSpec(partitionSpec1, Seq.empty, UnspecifiedFrame)).as('sum_a_1))
+      .select('a, 'b, 'c,  windowExpr(sum('c), windowSpec(partitionSpec2, Seq.empty, UnspecifiedFrame)).as('sum_a_2), 'sum_a_1)
+      .select('a, 'b, 'c, 'sum_a_2, 'sum_a_1)
 
     val correctAnswer = Optimize.execute(query2.analyze)
 
@@ -75,6 +75,7 @@ class FlipWindowSuite extends PlanTest {
     val correctAnswer = testRelation
       .window(Seq(sum(c).as('sum_a_1)), partitionSpec1, orderSpec1)
       .window(Seq(sum(c).as('sum_a_2)), partitionSpec2, orderSpec2)
+      .select('a, 'b, 'c, 'd, 'sum_a_2, 'sum_a_1)
 
     comparePlans(optimized, correctAnswer.analyze)
   }

sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala

@@ -468,4 +468,25 @@ class DataFrameWindowFunctionsSuite extends QueryTest with SharedSQLContext {
       spark.read.schema(sampleSchema).json(input.toDS()).select(c0, c1).foreach { _ => () }
     }
   }
+
+  test("window functions in multiple selects") {
+    val df = Seq(
+      ("S1", "P1", 100),
+      ("S1", "P1", 700),
+      ("S2", "P1", 200),
+      ("S2", "P2", 300)
+    ).toDF("sno", "pno", "qty")
+
+    val w1 = Window.partitionBy("sno")
+    val w2 = Window.partitionBy("sno", "pno")
+
+    checkAnswer(
+      df.select($"sno", $"pno", $"qty", sum($"qty").over(w2).alias("sum_qty_2"))
+        .select($"sno", $"pno", $"qty", col("sum_qty_2"), sum("qty").over(w1).alias("sum_qty_1")),
+      Seq(
+        Row("S1", "P1", 100, 800, 800),
+        Row("S1", "P1", 700, 800, 800),
+        Row("S2", "P1", 200, 200, 500),
+        Row("S2", "P2", 300, 300, 500)))
+  }
 }