HeartSaVioR commented on code in PR #49277:
URL: https://github.com/apache/spark/pull/49277#discussion_r1917947186
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala:
##########
@@ -496,6 +498,486 @@ class RocksDBStateStoreSuite extends
StateStoreSuiteBase[RocksDBStateStoreProvid
}
}
+ test("AvroStateEncoder - add field") {
+ val keySchema = StructType(Seq(
+ StructField("k", StringType)
+ ))
+
+ val initialValueSchema = StructType(Seq(
+ StructField("value", IntegerType, false)
+ ))
+
+ val evolvedValueSchema = StructType(Seq(
+ StructField("value", IntegerType, false),
+ StructField("timestamp", LongType, true)
+ ))
+
+ // Create test state schema provider
+ val testProvider = new TestStateSchemaProvider()
+
+ // Add initial schema version
+ testProvider.captureSchema(
+ StateStore.DEFAULT_COL_FAMILY_NAME,
+ keySchema,
+ initialValueSchema,
+ valueSchemaId = 0
+ )
+
+ // Create encoder with initial schema
+ val encoder1 = new AvroStateEncoder(
+ NoPrefixKeyStateEncoderSpec(keySchema),
+ initialValueSchema,
+ Some(testProvider),
+ Some(ColumnFamilyInfo(StateStore.DEFAULT_COL_FAMILY_NAME, 1))
+ )
+
+ // Create test data
+ val proj = UnsafeProjection.create(initialValueSchema)
+ val row1 = proj.apply(InternalRow(1))
+
+ // Encode with schema v0
+ val encoded = encoder1.encodeValue(row1)
+
+ // Add evolved schema
+ testProvider.captureSchema(
+ StateStore.DEFAULT_COL_FAMILY_NAME,
+ keySchema,
+ evolvedValueSchema,
+ valueSchemaId = 1
+ )
+
+ // Create encoder with initial schema
+ val encoder2 = new AvroStateEncoder(
+ NoPrefixKeyStateEncoderSpec(keySchema),
+ evolvedValueSchema,
+ Some(testProvider),
+ Some(ColumnFamilyInfo(StateStore.DEFAULT_COL_FAMILY_NAME, 1))
+ )
+
+ // Decode with evolved schema
+ val decoded = encoder2.decodeValue(encoded)
+
+ // Should be able to read old format
+ assert(decoded.getInt(0) === 1)
+ assert(decoded.getLong(1) === 0L) // New field should be null
+
+ // Encode with new schema
+ val proj2 = UnsafeProjection.create(evolvedValueSchema)
+ val row2 = proj2.apply(InternalRow(2, 100L))
+ val encoded2 = encoder2.encodeValue(row2)
+
+ // Should write with new schema version
+ val decoded2 = encoder2.decodeValue(encoded2)
+ assert(decoded2.getInt(0) === 2)
+ assert(decoded2.getLong(1) === 100L)
+ }
+
+ test("AvroStateEncoder - add field with null") {
+ val keySchema = StructType(Seq(
+ StructField("k", StringType)
+ ))
+
+ val initialValueSchema = StructType(Seq(
+ StructField("value", IntegerType, true) // Made nullable
+ ))
+
+ val evolvedValueSchema = StructType(Seq(
Review Comment:
See python tests - we allow users to put the schema of the state. They can
of course change the existing column from nullable to non-nullable. You may
argue that who will ever want to do that, but since this is technically
feasible, we should consider about this scenario.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
