initial commit

Author: xuanyuanking

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -1545,6 +1545,15 @@ object SQLConf {
       .booleanConf
       .createWithDefault(true)
 
+  val STREAMING_STATE_FORMAT_CHECK_ENABLED =
+    buildConf("spark.sql.streaming.stateFormatCheck.enabled")
+      .doc("Whether to detect a streaming query may try to use an invalid UnsafeRow in the " +
+        "state store.")
+      .version("3.1.0")
+      .internal()
+      .booleanConf
+      .createWithDefault(true)
+
   val PARALLEL_FILE_LISTING_IN_STATS_COMPUTATION =
     buildConf("spark.sql.statistics.parallelFileListingInStatsComputation.enabled")
       .internal()

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

@@ -17,9 +17,11 @@
 
 package org.apache.spark.sql.execution.streaming.state
 
+import org.apache.spark.SparkException
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.catalyst.expressions.{Attribute, UnsafeRow}
 import org.apache.spark.sql.catalyst.expressions.codegen.{GenerateUnsafeProjection, GenerateUnsafeRowJoiner}
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.StructType
 
 /**
@@ -59,6 +61,9 @@ sealed trait StreamingAggregationStateManager extends Serializable {
 
   /** Return an iterator containing all the values in target state store. */
   def values(store: StateStore): Iterator[UnsafeRow]
+
+  /** Check the UnsafeRow format with the expected schema */
+  def unsafeRowFormatValidation(row: UnsafeRow, schema: StructType): Unit
 }
 
 object StreamingAggregationStateManager extends Logging {
@@ -77,13 +82,24 @@ object StreamingAggregationStateManager extends Logging {
   }
 }
 
+/**
+ * An exception thrown when an invalid UnsafeRow is detected.
+ */
+class InvalidUnsafeRowException
+  extends SparkException("The UnsafeRow format is invalid. This may happen when using the old " +
+    "version or broken checkpoint file. To resolve this problem, you can try to restart the " +
+    "application or use the legacy way to process streaming state.", null)
+
 abstract class StreamingAggregationStateManagerBaseImpl(
     protected val keyExpressions: Seq[Attribute],
     protected val inputRowAttributes: Seq[Attribute]) extends StreamingAggregationStateManager {
 
   @transient protected lazy val keyProjector =
     GenerateUnsafeProjection.generate(keyExpressions, inputRowAttributes)
 
+  // Consider about the cost, only check the UnsafeRow format for the first row
+  private var checkFormat = true
+
   override def getKey(row: UnsafeRow): UnsafeRow = keyProjector(row)
 
   override def commit(store: StateStore): Long = store.commit()
@@ -94,6 +110,28 @@ abstract class StreamingAggregationStateManagerBaseImpl(
     // discard and don't convert values to avoid computation
     store.getRange(None, None).map(_.key)
   }
+
+  override def unsafeRowFormatValidation(row: UnsafeRow, schema: StructType): Unit = {
+    if (checkFormat && SQLConf.get.getConf(
+        SQLConf.STREAMING_STATE_FORMAT_CHECK_ENABLED) && row != null) {
+      if (schema.fields.length != row.numFields) {
+        throw new InvalidUnsafeRowException
+      }
+      schema.fields.zipWithIndex
+        .filterNot(field => UnsafeRow.isFixedLength(field._1.dataType)).foreach {
+        case (_, index) =>
+          val offsetAndSize = row.getLong(index)
+          val offset = (offsetAndSize >> 32).toInt
+          val size = offsetAndSize.toInt
+          if (size < 0 ||
+              offset < UnsafeRow.calculateBitSetWidthInBytes(row.numFields) + 8 * row.numFields ||
+              offset + size > row.getSizeInBytes) {
+            throw new InvalidUnsafeRowException
+          }
+      }
+      checkFormat = false
+    }
+  }
 }
 
 /**
@@ -114,7 +152,9 @@ class StreamingAggregationStateManagerImplV1(
   override def getStateValueSchema: StructType = inputRowAttributes.toStructType
 
   override def get(store: StateStore, key: UnsafeRow): UnsafeRow = {
-    store.get(key)
+    val res = store.get(key)
+    unsafeRowFormatValidation(res, inputRowAttributes.toStructType)
+    res
   }
 
   override def put(store: StateStore, row: UnsafeRow): Unit = {
@@ -173,7 +213,9 @@ class StreamingAggregationStateManagerImplV2(
       return savedState
     }
 
-    restoreOriginalRow(key, savedState)
+    val res = restoreOriginalRow(key, savedState)
+    unsafeRowFormatValidation(res, inputRowAttributes.toStructType)
+    res
   }
 
   override def put(store: StateStore, row: UnsafeRow): Unit = {

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

@@ -20,7 +20,7 @@ package org.apache.spark.sql.execution.streaming.state
 import org.apache.spark.sql.catalyst.expressions.{Attribute, SpecificInternalRow, UnsafeProjection, UnsafeRow}
 import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeProjection
 import org.apache.spark.sql.streaming.StreamTest
-import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
 
 class StreamingAggregationStateManagerSuite extends StreamTest {
   // ============================ fields and method for test data ============================
@@ -123,4 +123,24 @@ class StreamingAggregationStateManagerSuite extends StreamTest {
     // state manager should return row which is same as input row regardless of format version
     assert(inputRow === stateManager.get(memoryStateStore, keyRow))
   }
+
+  test("UnsafeRow format invalidation") {
+    // Pass the checking
+    val stateManager0 = StreamingAggregationStateManager.createStateManager(testKeyAttributes,
+      testOutputAttributes, 2)
+    stateManager0.unsafeRowFormatValidation(testRow, testOutputSchema)
+    // Fail for fields number not match
+    val stateManager1 = StreamingAggregationStateManager.createStateManager(testKeyAttributes,
+      testOutputAttributes, 2)
+    assertThrows[InvalidUnsafeRowException](stateManager1.unsafeRowFormatValidation(
+      testRow, StructType(testKeys.map(createIntegerField))))
+    // Fail for invalid schema
+    val stateManager2 = StreamingAggregationStateManager.createStateManager(testKeyAttributes,
+      testOutputAttributes, 2)
+    val invalidSchema = StructType(testKeys.map(createIntegerField) ++
+      Seq(StructField("struct", StructType(Seq(StructField("value1", StringType, true))), true),
+      StructField("value2", IntegerType, false)))
+    assertThrows[InvalidUnsafeRowException](stateManager2.unsafeRowFormatValidation(
+      testRow, invalidSchema))
+  }
 }