Added more checks

Author: brkyvz

sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateSuite.scala

@@ -30,8 +30,9 @@ import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeProj
 import org.apache.spark.sql.catalyst.plans.logical.FlatMapGroupsWithState
 import org.apache.spark.sql.catalyst.plans.physical.UnknownPartitioning
 import org.apache.spark.sql.catalyst.streaming.InternalOutputModes._
-import org.apache.spark.sql.execution.RDDScanExec
-import org.apache.spark.sql.execution.streaming.{FlatMapGroupsWithStateExec, GroupStateImpl, MemoryStream}
+import org.apache.spark.sql.execution.{RDDScanExec, WholeStageCodegenExec}
+import org.apache.spark.sql.execution.exchange.ShuffleExchange
+import org.apache.spark.sql.execution.streaming.{FlatMapGroupsWithStateExec, GroupStateImpl, MemoryStream, StreamingQueryWrapper}
 import org.apache.spark.sql.execution.streaming.state.{StateStore, StateStoreId, StateStoreMetrics, UnsafeRowPair}
 import org.apache.spark.sql.streaming.FlatMapGroupsWithStateSuite.MemoryStateStore
 import org.apache.spark.sql.streaming.util.{MockSourceProvider, StreamManualClock}
@@ -917,6 +918,13 @@ class FlatMapGroupsWithStateSuite extends StateStoreMetricsTest with BeforeAndAf
           inputSource.releaseLock()
           sq.processAllAvailable()
 
+          val restore1 = sq.asInstanceOf[StreamingQueryWrapper].streamingQuery
+            .lastExecution.executedPlan
+            .collect { case ss: FlatMapGroupsWithStateExec => ss }
+            .head
+          assert(restore1.child.outputPartitioning.numPartitions ===
+            spark.sessionState.conf.numShufflePartitions)
+
           checkDataset(
             spark.read.parquet(data).as[Int],
             1)
@@ -939,6 +947,13 @@ class FlatMapGroupsWithStateSuite extends StateStoreMetricsTest with BeforeAndAf
           inputSource.releaseLock()
           sq2.processAllAvailable()
 
+          val restore2 = sq.asInstanceOf[StreamingQueryWrapper].streamingQuery
+            .lastExecution.executedPlan
+            .collect { case ss: FlatMapGroupsWithStateExec => ss }
+            .head
+          assert(restore2.child.outputPartitioning.numPartitions ===
+            spark.sessionState.conf.numShufflePartitions)
+
           checkDataset(
             spark.read.parquet(data).as[Int],
             4, 3, 1)

sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala

@@ -30,7 +30,9 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.plans.logical.Aggregate
 import org.apache.spark.sql.catalyst.plans.physical.SinglePartition
 import org.apache.spark.sql.catalyst.util.DateTimeUtils
-import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.{SparkPlan, WholeStageCodegenExec}
+import org.apache.spark.sql.execution.aggregate.HashAggregateExec
+import org.apache.spark.sql.execution.exchange.ShuffleExchange
 import org.apache.spark.sql.execution.streaming._
 import org.apache.spark.sql.execution.streaming.state.StateStore
 import org.apache.spark.sql.expressions.scalalang.typed
@@ -417,10 +419,30 @@ class StreamingAggregationSuite extends StateStoreMetricsTest
             spark.table("agg_test").as[Long],
             1L)
 
+          val restore1 = sq.asInstanceOf[StreamingQueryWrapper].streamingQuery
+            .lastExecution.executedPlan
+            .collect { case ss: StateStoreRestoreExec => ss }
+            .head
+          restore1.child match {
+            case wscg: WholeStageCodegenExec =>
+              assert(wscg.outputPartitioning.numPartitions === 1)
+              assert(wscg.child.isInstanceOf[HashAggregateExec], "Shouldn't require shuffling")
+            case _ => fail("Expected no shuffling")
+          }
+
           inputSource.addData()
           inputSource.releaseLock()
           sq.processAllAvailable()
 
+          val restore2 = sq.asInstanceOf[StreamingQueryWrapper].streamingQuery
+            .lastExecution.executedPlan
+            .collect { case ss: StateStoreRestoreExec => ss }
+            .head
+          restore2.child match {
+            case shuffle: ShuffleExchange => assert(shuffle.newPartitioning.numPartitions === 1)
+            case _ => fail("Expected shuffling when there was no data")
+          }
+
           checkDataset(
             spark.table("agg_test").as[Long],
             1L)
@@ -472,6 +494,17 @@ class StreamingAggregationSuite extends StateStoreMetricsTest
           inputSource.releaseLock()
           sq.processAllAvailable()
 
+          val restore1 = sq.asInstanceOf[StreamingQueryWrapper].streamingQuery
+            .lastExecution.executedPlan
+            .collect { case ss: StateStoreRestoreExec => ss }
+            .head
+          restore1.child match {
+            case shuffle: ShuffleExchange =>
+              assert(shuffle.newPartitioning.numPartitions ===
+                spark.sessionState.conf.numShufflePartitions)
+            case _ => fail(s"Expected shuffling but got: ${restore1.child}")
+          }
+
           checkDataset(
             spark.table("agg_test").as[(Long, Long)],
             (0L, 1L))
@@ -502,6 +535,19 @@ class StreamingAggregationSuite extends StateStoreMetricsTest
           inputSource.releaseLock()
           sq2.processAllAvailable()
 
+          val restore2 = sq2.asInstanceOf[StreamingQueryWrapper].streamingQuery
+            .lastExecution.executedPlan
+            .collect { case ss: StateStoreRestoreExec => ss }
+            .head
+          restore2.child match {
+            case wscg: WholeStageCodegenExec =>
+              assert(wscg.outputPartitioning.numPartitions ===
+                spark.sessionState.conf.numShufflePartitions)
+            case _ =>
+              fail("Shouldn't require shuffling as HashAggregateExec should have asked for a " +
+                s"shuffle. But got: ${restore2.child}")
+          }
+
           checkDataset(
             spark.table("agg_test").as[(Long, Long)],
             (0L, 4L))