[SPARK-51065][SQL] Disallowing non-nullable schema when Avro encoding is used for TransformWithState

common/utils/src/main/resources/error/error-conditions.json

@@ -4669,6 +4669,13 @@
     ],
     "sqlState" : "42K06"
   },
+  "STATE_STORE_SCHEMA_MUST_BE_NULLABLE" : {
+    "message" : [
+      "If schema evolution is enabled, all the fields in the schema for column family <columnFamilyName> must be nullable.",
+      "Please make the schema nullable. Current schema: <schema>"
+    ],
+    "sqlState" : "XXKST"
+  },
   "STATE_STORE_STATE_SCHEMA_FILES_THRESHOLD_EXCEEDED" : {
     "message" : [
       "The number of state schema files <numStateSchemaFiles> exceeds the maximum number of state schema files for this query: <maxStateSchemaFiles>.",

python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py

@@ -1470,6 +1470,39 @@ def check_exception(error):
                     check_exception=check_exception,
                 )
 
+    def test_not_nullable_fails(self):
+        with self.sql_conf({"spark.sql.streaming.stateStore.encodingFormat": "avro"}):
+            with tempfile.TemporaryDirectory() as checkpoint_dir:
+                input_path = tempfile.mkdtemp()
+                self._prepare_test_resource1(input_path)
+
+                df = self._build_test_df(input_path)
+
+                def check_basic_state(batch_df, batch_id):
+                    result = batch_df.collect()[0]
+                    assert result.value["id"] == 0  # First ID from test data
+                    assert result.value["name"] == "name-0"
+
+                def check_exception(error):
+                    from pyspark.errors.exceptions.captured import StreamingQueryException
+
+                    if not isinstance(error, StreamingQueryException):
+                        return False
+
+                    error_msg = str(error)
+                    return (
+                        "[STATE_STORE_SCHEMA_MUST_BE_NULLABLE]" in error_msg
+                        and "column family state must be nullable" in error_msg
+                    )
+
+                self._run_evolution_test(
+                    BasicProcessorNotNullable(),
+                    checkpoint_dir,
+                    check_basic_state,
+                    df,
+                    check_exception=check_exception,
+                )
+
 
 class SimpleStatefulProcessorWithInitialState(StatefulProcessor):
     # this dict is the same as input initial state dataframe
@@ -1893,6 +1926,27 @@ def close(self) -> None:
         pass
 
 
+class BasicProcessorNotNullable(StatefulProcessor):
+    # Schema definitions
+    state_schema = StructType(
+        [StructField("id", IntegerType(), False), StructField("name", StringType(), False)]
+    )
+
+    def init(self, handle):
+        self.state = handle.getValueState("state", self.state_schema)
+
+    def handleInputRows(self, key, rows, timerValues) -> Iterator[pd.DataFrame]:
+        for pdf in rows:
+            pass
+        id_val = int(key[0])
+        name = f"name-{id_val}"
+        self.state.update((id_val, name))
+        yield pd.DataFrame({"id": [key[0]], "value": [{"id": id_val, "name": name}]})
+
+    def close(self) -> None:
+        pass
+
+
 class AddFieldsProcessor(StatefulProcessor):
     state_schema = StructType(
         [

sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasExec.scala

@@ -117,7 +117,7 @@ case class TransformWithStateInPandasExec(
   override def getColFamilySchemas(
       setNullableFields: Boolean
   ): Map[String, StateStoreColFamilySchema] = {
-    driverProcessorHandle.getColumnFamilySchemas(setNullableFields)
+    driverProcessorHandle.getColumnFamilySchemas(true)
   }
 
   override def getStateVariableInfos(): Map[String, TransformWithStateVariableInfo] = {

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

@@ -38,7 +38,7 @@ trait ListStateMetricsImpl {
   // We keep track of the count of entries in the list in a separate column family
   // to avoid scanning the entire list to get the count.
   private val counterCFValueSchema: StructType =
-    StructType(Seq(StructField("count", LongType, nullable = false)))
+    StructType(Seq(StructField("count", LongType)))
 
   private val counterCFProjection = UnsafeProjection.create(counterCFValueSchema)
 

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

@@ -47,7 +47,7 @@ object StateStoreColumnFamilySchemaUtils {
         // Byte type is converted to Int in Avro, which doesn't work for us as Avro
         // uses zig-zag encoding as opposed to big-endian for Ints
         Seq(
-          StructField(s"${field.name}_marker", BinaryType, nullable = false),
+          StructField(s"${field.name}_marker", BinaryType),
           field.copy(name = s"${field.name}_value", BinaryType)
         )
       } else {
@@ -117,7 +117,7 @@ object StateStoreColumnFamilySchemaUtils {
       getRowCounterCFName(stateName), keySchemaId = 0,
       keyEncoder.schema,
       valueSchemaId = 0,
-      StructType(Seq(StructField("count", LongType, nullable = false))),
+      StructType(Seq(StructField("count", LongType))),
       Some(NoPrefixKeyStateEncoderSpec(keyEncoder.schema)))
     schemas.put(counterSchema.colFamilyName, counterSchema)
 
@@ -149,7 +149,7 @@ object StateStoreColumnFamilySchemaUtils {
         keySchemaId = 0,
         keyEncoder.schema,
         valueSchemaId = 0,
-        StructType(Seq(StructField("count", LongType, nullable = false))),
+        StructType(Seq(StructField("count", LongType))),
         Some(NoPrefixKeyStateEncoderSpec(keyEncoder.schema)))
       schemas.put(countSchema.colFamilyName, countSchema)
     }

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

@@ -364,17 +364,20 @@ class DriverStatefulProcessorHandleImpl(timeMode: TimeMode, keyExprEnc: Expressi
   }
 
   def getColumnFamilySchemas(
-      setNullableFields: Boolean
+      shouldCheckNullable: Boolean
   ): Map[String, StateStoreColFamilySchema] = {
     val schemas = columnFamilySchemas.toMap
-    if (setNullableFields) {
-      schemas.map { case (colFamilyName, stateStoreColFamilySchema) =>
-        colFamilyName -> stateStoreColFamilySchema.copy(
-          valueSchema = stateStoreColFamilySchema.valueSchema.toNullable
-        )
+    schemas.map { case (colFamilyName, schema) =>
+      // assert that each field is nullable if schema evolution is enabled
+      schema.valueSchema.fields.foreach { field =>
+        if (!field.nullable && shouldCheckNullable) {
+          throw StateStoreErrors.stateStoreSchemaMustBeNullable(
+            schema.colFamilyName, schema.valueSchema.toString())
+        }
       }
-    } else {
-      schemas
+      colFamilyName -> schema.copy(
+        valueSchema = schema.valueSchema.toNullable
+      )
     }
   }
 
@@ -549,7 +552,7 @@ class DriverStatefulProcessorHandleImpl(timeMode: TimeMode, keyExprEnc: Expressi
       elementKeySchema: StructType): StateStoreColFamilySchema = {
     val countIndexName = s"$$count_$stateName"
     val countValueSchema = StructType(Seq(
-      StructField("count", LongType, nullable = false)
+      StructField("count", LongType)
     ))
 
     StateStoreColFamilySchema(

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

@@ -357,7 +357,7 @@ abstract class OneToManyTTLState(
   // Schema of the entry count index: elementKey -> count
   private val COUNT_INDEX = "$count_" + stateName
   private val COUNT_INDEX_VALUE_SCHEMA: StructType =
-    StructType(Seq(StructField("count", LongType, nullable = false)))
+    StructType(Seq(StructField("count", LongType)))
   private val countIndexValueProjector = UnsafeProjection.create(COUNT_INDEX_VALUE_SCHEMA)
 
   // Reused internal row that we use to create an UnsafeRow with the schema of

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

@@ -150,7 +150,7 @@ case class TransformWithStateExec(
       Some(NoPrefixKeyStateEncoderSpec(keySchema)))
 
     val columnFamilySchemas = getDriverProcessorHandle()
-      .getColumnFamilySchemas(setNullableFields) ++
+      .getColumnFamilySchemas(false) ++
         Map(StateStore.DEFAULT_COL_FAMILY_NAME -> defaultSchema)
     closeProcessorHandle()
     columnFamilySchemas

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

@@ -145,6 +145,12 @@ object StateStoreErrors {
     new StateStoreValueSchemaNotCompatible(storedValueSchema, newValueSchema)
   }
 
+  def stateStoreSchemaMustBeNullable(
+      columnFamilyName: String,
+      schema: String): StateStoreSchemaMustBeNullable = {
+    new StateStoreSchemaMustBeNullable(columnFamilyName, schema)
+  }
+
   def stateStoreInvalidValueSchemaEvolution(
       oldValueSchema: String,
       newValueSchema: String): StateStoreInvalidValueSchemaEvolution = {
@@ -346,6 +352,15 @@ class StateStoreValueSchemaNotCompatible(
       "storedValueSchema" -> storedValueSchema,
       "newValueSchema" -> newValueSchema))
 
+class StateStoreSchemaMustBeNullable(
+    columnFamilyName: String,
+    schema: String)
+  extends SparkUnsupportedOperationException(
+    errorClass = "STATE_STORE_SCHEMA_MUST_BE_NULLABLE",
+    messageParameters = Map(
+      "columnFamilyName" -> columnFamilyName,
+      "schema" -> schema))
+
 class StateStoreInvalidValueSchemaEvolution(
     oldValueSchema: String,
     newValueSchema: String)

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

@@ -1815,7 +1815,6 @@ class TransformWithStateSuite extends StateStoreMetricsTest
           TransformWithStateSuiteUtils.NUM_SHUFFLE_PARTITIONS.toString) {
         withTempDir { checkpointDir =>
           // When Avro is used, we want to set the StructFields to nullable
-          val shouldBeNullable = usingAvroEncoding()
           val metadataPathPostfix = "state/0/_stateSchema/default"
           val stateSchemaPath = new Path(checkpointDir.toString,
             s"$metadataPathPostfix")
@@ -1826,15 +1825,15 @@ class TransformWithStateSuite extends StateStoreMetricsTest
           val schema0 = StateStoreColFamilySchema(
             "countState", 0,
             keySchema, 0,
-            new StructType().add("value", LongType, nullable = shouldBeNullable),
+            new StructType().add("value", LongType, nullable = true),
             Some(NoPrefixKeyStateEncoderSpec(keySchema)),
             None
           )
           val schema1 = StateStoreColFamilySchema(
             "listState", 0,
             keySchema, 0,
             new StructType()
-              .add("id", LongType, nullable = shouldBeNullable)
+              .add("id", LongType, nullable = true)
               .add("name", StringType),
             Some(NoPrefixKeyStateEncoderSpec(keySchema)),
             None
@@ -1857,7 +1856,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
           val schema3 = StateStoreColFamilySchema(
             "$rowCounter_listState", 0,
             keySchema, 0,
-            new StructType().add("count", LongType, nullable = shouldBeNullable),
+            new StructType().add("count", LongType, nullable = true),
             Some(NoPrefixKeyStateEncoderSpec(keySchema)),
             None
           )
@@ -2209,9 +2208,9 @@ class TransformWithStateSuite extends StateStoreMetricsTest
               t.asInstanceOf[SparkUnsupportedOperationException],
               condition = "STATE_STORE_VALUE_SCHEMA_NOT_COMPATIBLE",
               parameters = Map(
-                "storedValueSchema" -> "StructType(StructField(value,LongType,false))",
+                "storedValueSchema" -> "StructType(StructField(value,LongType,true))",
                 "newValueSchema" ->
-                  ("StructType(StructField(value,StructType(StructField(value,LongType,false))," +
+                  ("StructType(StructField(value,StructType(StructField(value,LongType,true))," +
                     "true),StructField(ttlExpirationMs,LongType,true))")
               )
             )

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

@@ -273,7 +273,6 @@ class TransformWithValueStateTTLSuite extends TransformWithStateTTLTest {
     ) {
       withTempDir { checkpointDir =>
         // When Avro is used, we want to set the StructFields to nullable
-        val shouldBeNullable = usingAvroEncoding()
         val metadataPathPostfix = "state/0/_stateSchema/default"
         val stateSchemaPath = new Path(checkpointDir.toString, s"$metadataPathPostfix")
         val hadoopConf = spark.sessionState.newHadoopConf()
@@ -317,15 +316,15 @@ class TransformWithValueStateTTLSuite extends TransformWithStateTTLTest {
         val schema2 = StateStoreColFamilySchema(
           "$count_listState", 0,
           keySchema, 0,
-          new StructType().add("count", LongType, nullable = shouldBeNullable),
+          new StructType().add("count", LongType, nullable = true),
           Some(NoPrefixKeyStateEncoderSpec(keySchema)),
           None
         )
 
         val schema3 = StateStoreColFamilySchema(
           "$rowCounter_listState", 0,
           keySchema, 0,
-          new StructType().add("count", LongType, nullable = shouldBeNullable),
+          new StructType().add("count", LongType, nullable = true),
           Some(NoPrefixKeyStateEncoderSpec(keySchema)),
           None
         )
@@ -409,7 +408,7 @@ class TransformWithValueStateTTLSuite extends TransformWithStateTTLTest {
           "valueStateTTL", 0,
           keySchema, 0,
           new StructType()
-            .add("value", new StructType().add("value", IntegerType, nullable = shouldBeNullable))
+            .add("value", new StructType().add("value", IntegerType, nullable = true))
             .add("ttlExpirationMs", LongType),
           Some(NoPrefixKeyStateEncoderSpec(keySchema)),
           None
@@ -418,7 +417,7 @@ class TransformWithValueStateTTLSuite extends TransformWithStateTTLTest {
         val schema10 = StateStoreColFamilySchema(
           "valueState", 0,
           keySchema, 0,
-          new StructType().add("value", IntegerType, nullable = shouldBeNullable),
+          new StructType().add("value", IntegerType, nullable = true),
           Some(NoPrefixKeyStateEncoderSpec(keySchema)),
           None
         )
@@ -428,7 +427,7 @@ class TransformWithValueStateTTLSuite extends TransformWithStateTTLTest {
           keySchema, 0,
           new StructType()
             .add("value", new StructType()
-              .add("id", LongType, nullable = shouldBeNullable)
+              .add("id", LongType, nullable = true)
               .add("name", StringType))
             .add("ttlExpirationMs", LongType),
           Some(NoPrefixKeyStateEncoderSpec(keySchema)),