[SPARK-55145][SS] Support Avro for timestamp based RocksDB state key encoders

### What changes were proposed in this pull request?

This PR proposes to support Avro for timestamp key encoders, TimestampAsPrefixKeyStateEncoder(Spec) and TimestampAsPostfixKeyStateEncoder(Spec).

This PR does not cover the scope of enabling Avro for stream-stream join operators; this would trigger more scope than expected e.g. schema evolution and warrant separate effort.

### Why are the changes needed?

To resolve TODO we left in #53911

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

Existing UTs expanded from UnsafeRow only to UnsafeRow and Avro.

### Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude 4.6 opus

Closes #54844 from HeartSaVioR/SPARK-55145.

Authored-by: Jungtaek Lim <[email protected]>
Signed-off-by: Jungtaek Lim <[email protected]>

Author: HeartSaVioR

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

@@ -624,6 +624,8 @@ class UnsafeRowDataEncoder(
         decodeToUnsafeRow(bytes, reusedKeyRow)
       case PrefixKeyScanStateEncoderSpec(_, numColsPrefixKey) =>
         decodeToUnsafeRow(bytes, numFields = numColsPrefixKey)
+      case _: TimestampAsPrefixKeyStateEncoderSpec | _: TimestampAsPostfixKeyStateEncoderSpec =>
+        decodeToUnsafeRow(bytes, numFields = keySchema.length - 1)
       case _ => throw unsupportedOperationForKeyStateEncoder("decodeKey")
     }
   }
@@ -748,8 +750,15 @@ class AvroStateEncoder(
   )
 
   // Avro schema used by the avro encoders
-  private lazy val keyAvroType: Schema = SchemaConverters.toAvroTypeWithDefaults(keySchema)
-  private lazy val keyProj = UnsafeProjection.create(keySchema)
+  // For timestamp specs, the key part excludes the timestamp column (always the last field).
+  private lazy val effectiveKeySchema: StructType = keyStateEncoderSpec match {
+    case TimestampAsPrefixKeyStateEncoderSpec(s) => StructType(s.dropRight(1))
+    case TimestampAsPostfixKeyStateEncoderSpec(s) => StructType(s.dropRight(1))
+    case _ => keySchema
+  }
+  private lazy val keyAvroType: Schema =
+    SchemaConverters.toAvroTypeWithDefaults(effectiveKeySchema)
+  private lazy val keyProj = UnsafeProjection.create(effectiveKeySchema)
 
   private lazy val valueAvroType: Schema = SchemaConverters.toAvroTypeWithDefaults(valueSchema)
   private lazy val valueProj = UnsafeProjection.create(valueSchema)
@@ -847,8 +856,10 @@ class AvroStateEncoder(
           }
         }
         StructType(remainingSchema)
-      case _ =>
-        throw unsupportedOperationForKeyStateEncoder("createAvroEnc")
+      case TimestampAsPrefixKeyStateEncoderSpec(schema) =>
+        StructType(schema.dropRight(1))
+      case TimestampAsPostfixKeyStateEncoderSpec(schema) =>
+        StructType(schema.dropRight(1))
     }
 
     // Handle suffix key schema for prefix scan case
@@ -1005,6 +1016,11 @@ class AvroStateEncoder(
           StateSchemaIdRow(currentKeySchemaId, avroRow))
       case PrefixKeyScanStateEncoderSpec(_, _) =>
         encodeUnsafeRowToAvro(row, avroEncoder.keySerializer, prefixKeyAvroType, out)
+      case _: TimestampAsPrefixKeyStateEncoderSpec | _: TimestampAsPostfixKeyStateEncoderSpec =>
+        val avroRow =
+          encodeUnsafeRowToAvro(row, avroEncoder.keySerializer, keyAvroType, out)
+        encodeWithStateSchemaId(
+          StateSchemaIdRow(currentKeySchemaId, avroRow))
       case _ => throw unsupportedOperationForKeyStateEncoder("encodeKey")
     }
     prependVersionByte(keyBytes)
@@ -1179,6 +1195,10 @@ class AvroStateEncoder(
       case PrefixKeyScanStateEncoderSpec(_, _) =>
         decodeFromAvroToUnsafeRow(
           bytes, avroEncoder.keyDeserializer, prefixKeyAvroType, prefixKeyProj)
+      case _: TimestampAsPrefixKeyStateEncoderSpec | _: TimestampAsPostfixKeyStateEncoderSpec =>
+        val schemaIdRow = decodeStateSchemaIdRow(bytes)
+        decodeFromAvroToUnsafeRow(
+          schemaIdRow.bytes, avroEncoder.keyDeserializer, keyAvroType, keyProj)
       case _ => throw unsupportedOperationForKeyStateEncoder("decodeKey")
     }
   }
@@ -1782,9 +1802,7 @@ abstract class TimestampKeyStateEncoder(
       rowBytes, Platform.BYTE_ARRAY_OFFSET,
       rowBytesLength
     )
-    // The encoded row does not include the timestamp (it's stored separately),
-    // so decode with keySchema.length - 1 fields.
-    dataEncoder.decodeToUnsafeRow(rowBytes, keySchema.length - 1)
+    dataEncoder.decodeKey(rowBytes)
   }
 
   // NOTE: We reuse the ByteBuffer to avoid allocating a new one for every encoding/decoding,

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

@@ -673,7 +673,10 @@ case class RangeKeyScanStateEncoderSpec(
   }
 }
 
-/** The encoder specification for [[TimestampAsPrefixKeyStateEncoder]]. */
+/**
+ * The encoder specification for [[TimestampAsPrefixKeyStateEncoder]].
+ * The encoder expects the provided key schema to have [original key fields..., timestamp field].
+ */
 case class TimestampAsPrefixKeyStateEncoderSpec(keySchema: StructType)
   extends KeyStateEncoderSpec {
 
@@ -688,7 +691,10 @@ case class TimestampAsPrefixKeyStateEncoderSpec(keySchema: StructType)
   }
 }
 
-/** The encoder specification for [[TimestampAsPostfixKeyStateEncoder]]. */
+/**
+ * The encoder specification for [[TimestampAsPostfixKeyStateEncoder]].
+ * The encoder expects the provided key schema to have [original key fields..., timestamp field].
+ */
 case class TimestampAsPostfixKeyStateEncoderSpec(keySchema: StructType)
   extends KeyStateEncoderSpec {
 

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

@@ -68,9 +68,7 @@ class RocksDBTimestampEncoderOperationsSuite extends SharedSparkSession
 
   private def newDir(): String = Utils.createTempDir().getCanonicalPath
 
-  // TODO: [SPARK-55145] Address the new state format with Avro and enable the test with Avro
-  //  encoding
-  Seq("unsaferow").foreach { encoding =>
+  Seq("unsaferow", "avro").foreach { encoding =>
     Seq("prefix", "postfix").foreach { encoderType =>
       test(s"Event time as $encoderType: basic put and get operations (encoding = $encoding)") {
         tryWithProviderResource(
@@ -223,9 +221,7 @@ class RocksDBTimestampEncoderOperationsSuite extends SharedSparkSession
     }
   }
 
-  // TODO: [SPARK-55145] Address the new state format with Avro and enable the test with Avro
-  //  encoding
-  Seq("unsaferow").foreach { encoding =>
+  Seq("unsaferow", "avro").foreach { encoding =>
     test(s"Event time as prefix: iterator operations (encoding = $encoding)") {
       tryWithProviderResource(
         newStoreProviderWithTimestampEncoder(
@@ -558,9 +554,7 @@ class RocksDBTimestampEncoderOperationsSuite extends SharedSparkSession
     }
   }
 
-  // TODO: [SPARK-55145] Address the new state format with Avro and enable the test with Avro
-  //  encoding
-  Seq("unsaferow").foreach { encoding =>
+  Seq("unsaferow", "avro").foreach { encoding =>
     Seq("prefix", "postfix").foreach { encoderType =>
       Seq(false, true).foreach { useMultipleValuesPerKey =>
         val multiValueSuffix = if (useMultipleValuesPerKey) " and multiple values" else ""