[SPARK-19873][SS] Record num shuffle partitions in offset log and enforce in next batch.

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/OffsetSeq.scala

@@ -71,7 +71,10 @@ object OffsetSeq {
  * @param batchTimestampMs: The current batch processing timestamp.
  * Time unit: milliseconds
  */
-case class OffsetSeqMetadata(var batchWatermarkMs: Long = 0, var batchTimestampMs: Long = 0) {
+case class OffsetSeqMetadata(
+    var batchWatermarkMs: Long = 0,
+    var batchTimestampMs: Long = 0,
+    var numShufflePartitions: Int = 0) {
   def json: String = Serialization.write(this)(OffsetSeqMetadata.format)
 }
 

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamExecution.scala

@@ -35,6 +35,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, Curre
 import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
 import org.apache.spark.sql.execution.QueryExecution
 import org.apache.spark.sql.execution.command.StreamingExplainCommand
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.streaming._
 import org.apache.spark.util.{Clock, UninterruptibleThread, Utils}
 
@@ -117,7 +118,8 @@ class StreamExecution(
   }
 
   /** Metadata associated with the offset seq of a batch in the query. */
-  protected var offsetSeqMetadata = OffsetSeqMetadata()
+  protected var offsetSeqMetadata =
+    OffsetSeqMetadata(numShufflePartitions = sparkSession.conf.get(SQLConf.SHUFFLE_PARTITIONS))
 
   override val id: UUID = UUID.fromString(streamMetadata.id)
 
@@ -380,7 +382,20 @@ class StreamExecution(
         logInfo(s"Resuming streaming query, starting with batch $batchId")
         currentBatchId = batchId
         availableOffsets = nextOffsets.toStreamProgress(sources)
-        offsetSeqMetadata = nextOffsets.metadata.getOrElse(OffsetSeqMetadata())
+        val numShufflePartitionsFromConf = sparkSession.conf.get(SQLConf.SHUFFLE_PARTITIONS)
+        offsetSeqMetadata = nextOffsets
+          .metadata
+          .getOrElse(OffsetSeqMetadata(0, 0, numShufflePartitionsFromConf))
+
+        /*
+         * For backwards compatibility, if # partitions was not recorded in the offset log, then
+         * ensure it is non-zero. The new value is picked up from the conf.
+         */
+        if (offsetSeqMetadata.numShufflePartitions == 0) {
+          offsetSeqMetadata.numShufflePartitions = numShufflePartitionsFromConf
+          logDebug("Number of shuffle partitions from previous run not found in checkpoint data. " +
+            s"Using the value from the conf, $numShufflePartitionsFromConf partitions.")
+        }
         logDebug(s"Found possibly unprocessed offsets $availableOffsets " +
           s"at batch timestamp ${offsetSeqMetadata.batchTimestampMs}")
 
@@ -542,9 +557,16 @@ class StreamExecution(
           cd.dataType, cd.timeZoneId)
     }
 
+    // Fork a cloned session and set confs to disallow change in number of partitions
+    val sparkSessionForCurrentBatch = sparkSession.cloneSession()
+    sparkSessionForCurrentBatch.conf.set(
+      SQLConf.SHUFFLE_PARTITIONS.key,
+      offsetSeqMetadata.numShufflePartitions.toString)
+    sparkSessionForCurrentBatch.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "false")
+
     reportTimeTaken("queryPlanning") {
       lastExecution = new IncrementalExecution(
-        sparkSession,
+        sparkSessionForCurrentBatch,
         triggerLogicalPlan,
         outputMode,
         checkpointFile("state"),
@@ -554,7 +576,8 @@ class StreamExecution(
     }
 
     val nextBatch =
-      new Dataset(sparkSession, lastExecution, RowEncoder(lastExecution.analyzed.schema))
+      new Dataset(sparkSessionForCurrentBatch, lastExecution,
+        RowEncoder(lastExecution.analyzed.schema))
 
     reportTimeTaken("addBatch") {
       sink.addBatch(currentBatchId, nextBatch)

sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/OffsetSeqLogSuite.scala

@@ -28,12 +28,22 @@ class OffsetSeqLogSuite extends SparkFunSuite with SharedSQLContext {
   case class StringOffset(override val json: String) extends Offset
 
   test("OffsetSeqMetadata - deserialization") {
-    assert(OffsetSeqMetadata(0, 0) === OffsetSeqMetadata("""{}"""))
-    assert(OffsetSeqMetadata(1, 0) === OffsetSeqMetadata("""{"batchWatermarkMs":1}"""))
-    assert(OffsetSeqMetadata(0, 2) === OffsetSeqMetadata("""{"batchTimestampMs":2}"""))
+    assert(OffsetSeqMetadata(0, 0, 0) === OffsetSeqMetadata("""{}"""))
+    assert(OffsetSeqMetadata(1, 0, 0) === OffsetSeqMetadata("""{"batchWatermarkMs":1}"""))
+    assert(OffsetSeqMetadata(0, 2, 0) === OffsetSeqMetadata("""{"batchTimestampMs":2}"""))
+    assert(OffsetSeqMetadata(0, 0, 2) === OffsetSeqMetadata("""{"numShufflePartitions":2}"""))
     assert(
-      OffsetSeqMetadata(1, 2) ===
+      OffsetSeqMetadata(1, 2, 0) ===
         OffsetSeqMetadata("""{"batchWatermarkMs":1,"batchTimestampMs":2}"""))
+    assert(
+      OffsetSeqMetadata(1, 0, 3) ===
+        OffsetSeqMetadata("""{"batchWatermarkMs":1,"numShufflePartitions":3}"""))
+    assert(
+      OffsetSeqMetadata(0, 2, 3) ===
+        OffsetSeqMetadata("""{"batchTimestampMs":2,"numShufflePartitions":3}"""))
+    assert(
+      OffsetSeqMetadata(1, 2, 3) === OffsetSeqMetadata(
+          """{"batchWatermarkMs":1,"batchTimestampMs":2,"numShufflePartitions":3}"""))
   }
 
   test("OffsetSeqLog - serialization - deserialization") {

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

@@ -389,6 +389,20 @@ class StreamSuite extends StreamTest {
     query.stop()
     assert(query.exception.isEmpty)
   }
+
+  test("SPARK-19873: streaming agg with change in number of partitions") {
+    val inputData = MemoryStream[(Int, Int)]
+    val agg = inputData.toDS().groupBy("_1").count()
+
+    testStream(agg, OutputMode.Complete())(
+      AddData(inputData, (1, 1), (2, 1), (1, 2)),
+      StartStream(additionalConfs = Map("spark.sql.shuffle.partitions" -> "2")),
+      CheckAnswer((1, 2), (2, 1)),
+      StopStream,
+      AddData(inputData, (3, 1), (2, 2), (1, 1)),
+      StartStream(additionalConfs = Map("spark.sql.shuffle.partitions" -> "5")),
+      CheckAnswer((1, 3), (2, 2), (3, 1)))
+  }
 }
 
 abstract class FakeSource extends StreamSourceProvider {

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

@@ -151,9 +151,13 @@ class StreamingQuerySuite extends StreamTest with BeforeAndAfter with Logging wi
       TestAwaitTermination(ExpectException[SparkException], timeoutMs = 10),
       AssertOnQuery(q => {
         q.exception.get.startOffset ===
-          q.committedOffsets.toOffsetSeq(Seq(inputData), OffsetSeqMetadata()).toString &&
+          q.committedOffsets.toOffsetSeq(
+            Seq(inputData),
+            OffsetSeqMetadata(0, 0, spark.conf.get(SQLConf.SHUFFLE_PARTITIONS))).toString &&
           q.exception.get.endOffset ===
-            q.availableOffsets.toOffsetSeq(Seq(inputData), OffsetSeqMetadata()).toString
+            q.availableOffsets.toOffsetSeq(
+              Seq(inputData),
+              OffsetSeqMetadata(0, 0, spark.conf.get(SQLConf.SHUFFLE_PARTITIONS))).toString
       }, "incorrect start offset or end offset on exception")
     )
   }