From 9f038aaff64514bd2e97d989fdbb1046866f28e2 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Thu, 24 Oct 2024 20:47:14 -0700
Subject: [PATCH 01/30] [SPARK-50112] Moving Avro files to sql/core so they can
 be used by TransformWithState operator

---
 .../org/apache/spark/sql/avro/AvroCompressionCodec.java     | 0
 .../java}/org/apache/spark/sql/avro/AvroDeserializer.scala  | 0
 .../java}/org/apache/spark/sql/avro/AvroFileFormat.scala    | 3 +--
 .../main/java}/org/apache/spark/sql/avro/AvroOptions.scala  | 0
 .../java}/org/apache/spark/sql/avro/AvroOutputWriter.scala  | 0
 .../org/apache/spark/sql/avro/AvroOutputWriterFactory.scala | 0
 .../java}/org/apache/spark/sql/avro/AvroSerializer.scala    | 6 ++----
 .../main/java}/org/apache/spark/sql/avro/AvroUtils.scala    | 0
 .../main/java/org/apache/spark/sql/avro/CustomDecimal.scala | 3 +--
 .../java}/org/apache/spark/sql/avro/SchemaConverters.scala  | 5 ++---
 .../org/apache/spark/sql/avro/SparkAvroKeyOutputFormat.java | 0
 11 files changed, 6 insertions(+), 11 deletions(-)
 rename {connector/avro => sql/core}/src/main/java/org/apache/spark/sql/avro/AvroCompressionCodec.java (100%)
 rename {connector/avro/src/main/scala => sql/core/src/main/java}/org/apache/spark/sql/avro/AvroDeserializer.scala (100%)
 rename {connector/avro/src/main/scala => sql/core/src/main/java}/org/apache/spark/sql/avro/AvroFileFormat.scala (98%)
 rename {connector/avro/src/main/scala => sql/core/src/main/java}/org/apache/spark/sql/avro/AvroOptions.scala (100%)
 rename {connector/avro/src/main/scala => sql/core/src/main/java}/org/apache/spark/sql/avro/AvroOutputWriter.scala (100%)
 rename {connector/avro/src/main/scala => sql/core/src/main/java}/org/apache/spark/sql/avro/AvroOutputWriterFactory.scala (100%)
 rename {connector/avro/src/main/scala => sql/core/src/main/java}/org/apache/spark/sql/avro/AvroSerializer.scala (98%)
 rename {connector/avro/src/main/scala => sql/core/src/main/java}/org/apache/spark/sql/avro/AvroUtils.scala (100%)
 rename {connector/avro => sql/core}/src/main/java/org/apache/spark/sql/avro/CustomDecimal.scala (97%)
 rename {connector/avro/src/main/scala => sql/core/src/main/java}/org/apache/spark/sql/avro/SchemaConverters.scala (98%)
 rename {connector/avro => sql/core}/src/main/java/org/apache/spark/sql/avro/SparkAvroKeyOutputFormat.java (100%)

diff --git a/connector/avro/src/main/java/org/apache/spark/sql/avro/AvroCompressionCodec.java b/sql/core/src/main/java/org/apache/spark/sql/avro/AvroCompressionCodec.java
similarity index 100%
rename from connector/avro/src/main/java/org/apache/spark/sql/avro/AvroCompressionCodec.java
rename to sql/core/src/main/java/org/apache/spark/sql/avro/AvroCompressionCodec.java
diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala b/sql/core/src/main/java/org/apache/spark/sql/avro/AvroDeserializer.scala
similarity index 100%
rename from connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala
rename to sql/core/src/main/java/org/apache/spark/sql/avro/AvroDeserializer.scala
diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroFileFormat.scala b/sql/core/src/main/java/org/apache/spark/sql/avro/AvroFileFormat.scala
similarity index 98%
rename from connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroFileFormat.scala
rename to sql/core/src/main/java/org/apache/spark/sql/avro/AvroFileFormat.scala
index 264c3a1f48abe..3e1aa11b52b3a 100755
--- a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroFileFormat.scala
+++ b/sql/core/src/main/java/org/apache/spark/sql/avro/AvroFileFormat.scala
@@ -21,8 +21,7 @@ import java.io._
 
 import scala.util.control.NonFatal
 
-import org.apache.avro.{LogicalTypes, Schema}
-import org.apache.avro.LogicalType
+import org.apache.avro.{LogicalType, LogicalTypes, Schema}
 import org.apache.avro.file.DataFileReader
 import org.apache.avro.generic.{GenericDatumReader, GenericRecord}
 import org.apache.avro.mapred.FsInput
diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroOptions.scala b/sql/core/src/main/java/org/apache/spark/sql/avro/AvroOptions.scala
similarity index 100%
rename from connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroOptions.scala
rename to sql/core/src/main/java/org/apache/spark/sql/avro/AvroOptions.scala
diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroOutputWriter.scala b/sql/core/src/main/java/org/apache/spark/sql/avro/AvroOutputWriter.scala
similarity index 100%
rename from connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroOutputWriter.scala
rename to sql/core/src/main/java/org/apache/spark/sql/avro/AvroOutputWriter.scala
diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroOutputWriterFactory.scala b/sql/core/src/main/java/org/apache/spark/sql/avro/AvroOutputWriterFactory.scala
similarity index 100%
rename from connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroOutputWriterFactory.scala
rename to sql/core/src/main/java/org/apache/spark/sql/avro/AvroOutputWriterFactory.scala
diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroSerializer.scala b/sql/core/src/main/java/org/apache/spark/sql/avro/AvroSerializer.scala
similarity index 98%
rename from connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroSerializer.scala
rename to sql/core/src/main/java/org/apache/spark/sql/avro/AvroSerializer.scala
index 1d9eada94658e..814a28e24f522 100644
--- a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroSerializer.scala
+++ b/sql/core/src/main/java/org/apache/spark/sql/avro/AvroSerializer.scala
@@ -21,14 +21,12 @@ import java.nio.ByteBuffer
 
 import scala.jdk.CollectionConverters._
 
+import org.apache.avro.{LogicalTypes, Schema}
 import org.apache.avro.Conversions.DecimalConversion
-import org.apache.avro.LogicalTypes
 import org.apache.avro.LogicalTypes.{LocalTimestampMicros, LocalTimestampMillis, TimestampMicros, TimestampMillis}
-import org.apache.avro.Schema
 import org.apache.avro.Schema.Type
 import org.apache.avro.Schema.Type._
-import org.apache.avro.generic.GenericData.{EnumSymbol, Fixed}
-import org.apache.avro.generic.GenericData.Record
+import org.apache.avro.generic.GenericData.{EnumSymbol, Fixed, Record}
 import org.apache.avro.util.Utf8
 
 import org.apache.spark.internal.Logging
diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroUtils.scala b/sql/core/src/main/java/org/apache/spark/sql/avro/AvroUtils.scala
similarity index 100%
rename from connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroUtils.scala
rename to sql/core/src/main/java/org/apache/spark/sql/avro/AvroUtils.scala
diff --git a/connector/avro/src/main/java/org/apache/spark/sql/avro/CustomDecimal.scala b/sql/core/src/main/java/org/apache/spark/sql/avro/CustomDecimal.scala
similarity index 97%
rename from connector/avro/src/main/java/org/apache/spark/sql/avro/CustomDecimal.scala
rename to sql/core/src/main/java/org/apache/spark/sql/avro/CustomDecimal.scala
index fab3d4493e344..a5700a0481531 100644
--- a/connector/avro/src/main/java/org/apache/spark/sql/avro/CustomDecimal.scala
+++ b/sql/core/src/main/java/org/apache/spark/sql/avro/CustomDecimal.scala
@@ -17,8 +17,7 @@
 
 package org.apache.spark.sql.avro
 
-import org.apache.avro.LogicalType
-import org.apache.avro.Schema
+import org.apache.avro.{LogicalType, Schema}
 
 import org.apache.spark.sql.types.DecimalType
 
diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala b/sql/core/src/main/java/org/apache/spark/sql/avro/SchemaConverters.scala
similarity index 98%
rename from connector/avro/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala
rename to sql/core/src/main/java/org/apache/spark/sql/avro/SchemaConverters.scala
index 1168a887abd8e..495fc011df462 100644
--- a/connector/avro/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala
+++ b/sql/core/src/main/java/org/apache/spark/sql/avro/SchemaConverters.scala
@@ -23,13 +23,12 @@ import scala.collection.mutable
 import scala.jdk.CollectionConverters._
 
 import org.apache.avro.{LogicalTypes, Schema, SchemaBuilder}
-import org.apache.avro.LogicalTypes.{Date, Decimal, LocalTimestampMicros, LocalTimestampMillis, TimestampMicros, TimestampMillis}
+import org.apache.avro.LogicalTypes.{Decimal, _}
 import org.apache.avro.Schema.Type._
 
 import org.apache.spark.annotation.DeveloperApi
-import org.apache.spark.internal.Logging
+import org.apache.spark.internal.{Logging, MDC}
 import org.apache.spark.internal.LogKeys.{FIELD_NAME, FIELD_TYPE, RECURSIVE_DEPTH}
-import org.apache.spark.internal.MDC
 import org.apache.spark.sql.avro.AvroOptions.RECURSIVE_FIELD_MAX_DEPTH_LIMIT
 import org.apache.spark.sql.catalyst.parser.CatalystSqlParser
 import org.apache.spark.sql.types._
diff --git a/connector/avro/src/main/java/org/apache/spark/sql/avro/SparkAvroKeyOutputFormat.java b/sql/core/src/main/java/org/apache/spark/sql/avro/SparkAvroKeyOutputFormat.java
similarity index 100%
rename from connector/avro/src/main/java/org/apache/spark/sql/avro/SparkAvroKeyOutputFormat.java
rename to sql/core/src/main/java/org/apache/spark/sql/avro/SparkAvroKeyOutputFormat.java

From 28c3dbd213a98fe50ea2d97f257e46f4aff62fd6 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Thu, 24 Oct 2024 21:00:07 -0700
Subject: [PATCH 02/30] moving scala to scala dir

---
 .../org/apache/spark/sql/avro/AvroDeserializer.scala            | 2 +-
 .../org/apache/spark/sql/avro/AvroFileFormat.scala              | 0
 .../{java => scala}/org/apache/spark/sql/avro/AvroOptions.scala | 0
 .../org/apache/spark/sql/avro/AvroOutputWriter.scala            | 0
 .../org/apache/spark/sql/avro/AvroOutputWriterFactory.scala     | 0
 .../org/apache/spark/sql/avro/AvroSerializer.scala              | 0
 .../{java => scala}/org/apache/spark/sql/avro/AvroUtils.scala   | 0
 .../org/apache/spark/sql/avro/CustomDecimal.scala               | 0
 .../org/apache/spark/sql/avro/SchemaConverters.scala            | 0
 9 files changed, 1 insertion(+), 1 deletion(-)
 rename sql/core/src/main/{java => scala}/org/apache/spark/sql/avro/AvroDeserializer.scala (99%)
 rename sql/core/src/main/{java => scala}/org/apache/spark/sql/avro/AvroFileFormat.scala (100%)
 rename sql/core/src/main/{java => scala}/org/apache/spark/sql/avro/AvroOptions.scala (100%)
 rename sql/core/src/main/{java => scala}/org/apache/spark/sql/avro/AvroOutputWriter.scala (100%)
 rename sql/core/src/main/{java => scala}/org/apache/spark/sql/avro/AvroOutputWriterFactory.scala (100%)
 rename sql/core/src/main/{java => scala}/org/apache/spark/sql/avro/AvroSerializer.scala (100%)
 rename sql/core/src/main/{java => scala}/org/apache/spark/sql/avro/AvroUtils.scala (100%)
 rename sql/core/src/main/{java => scala}/org/apache/spark/sql/avro/CustomDecimal.scala (100%)
 rename sql/core/src/main/{java => scala}/org/apache/spark/sql/avro/SchemaConverters.scala (100%)

diff --git a/sql/core/src/main/java/org/apache/spark/sql/avro/AvroDeserializer.scala b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala
similarity index 99%
rename from sql/core/src/main/java/org/apache/spark/sql/avro/AvroDeserializer.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala
index ac20614553ca2..4e559f6eee887 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/avro/AvroDeserializer.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala
@@ -29,7 +29,7 @@ import org.apache.avro.Schema.Type._
 import org.apache.avro.generic._
 import org.apache.avro.util.Utf8
 
-import org.apache.spark.sql.avro.AvroUtils.{nonNullUnionBranches, toFieldStr, AvroMatchedField}
+import org.apache.spark.sql.avro.AvroUtils.{ nonNullUnionBranches, toFieldStr, AvroMatchedField}
 import org.apache.spark.sql.catalyst.{InternalRow, NoopFilters, StructFilters}
 import org.apache.spark.sql.catalyst.expressions.{SpecificInternalRow, UnsafeArrayData}
 import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, ArrayData, DateTimeUtils, GenericArrayData}
diff --git a/sql/core/src/main/java/org/apache/spark/sql/avro/AvroFileFormat.scala b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroFileFormat.scala
similarity index 100%
rename from sql/core/src/main/java/org/apache/spark/sql/avro/AvroFileFormat.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/avro/AvroFileFormat.scala
diff --git a/sql/core/src/main/java/org/apache/spark/sql/avro/AvroOptions.scala b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroOptions.scala
similarity index 100%
rename from sql/core/src/main/java/org/apache/spark/sql/avro/AvroOptions.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/avro/AvroOptions.scala
diff --git a/sql/core/src/main/java/org/apache/spark/sql/avro/AvroOutputWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroOutputWriter.scala
similarity index 100%
rename from sql/core/src/main/java/org/apache/spark/sql/avro/AvroOutputWriter.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/avro/AvroOutputWriter.scala
diff --git a/sql/core/src/main/java/org/apache/spark/sql/avro/AvroOutputWriterFactory.scala b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroOutputWriterFactory.scala
similarity index 100%
rename from sql/core/src/main/java/org/apache/spark/sql/avro/AvroOutputWriterFactory.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/avro/AvroOutputWriterFactory.scala
diff --git a/sql/core/src/main/java/org/apache/spark/sql/avro/AvroSerializer.scala b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroSerializer.scala
similarity index 100%
rename from sql/core/src/main/java/org/apache/spark/sql/avro/AvroSerializer.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/avro/AvroSerializer.scala
diff --git a/sql/core/src/main/java/org/apache/spark/sql/avro/AvroUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroUtils.scala
similarity index 100%
rename from sql/core/src/main/java/org/apache/spark/sql/avro/AvroUtils.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/avro/AvroUtils.scala
diff --git a/sql/core/src/main/java/org/apache/spark/sql/avro/CustomDecimal.scala b/sql/core/src/main/scala/org/apache/spark/sql/avro/CustomDecimal.scala
similarity index 100%
rename from sql/core/src/main/java/org/apache/spark/sql/avro/CustomDecimal.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/avro/CustomDecimal.scala
diff --git a/sql/core/src/main/java/org/apache/spark/sql/avro/SchemaConverters.scala b/sql/core/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala
similarity index 100%
rename from sql/core/src/main/java/org/apache/spark/sql/avro/SchemaConverters.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala

From 2e33fd170971f07a768f612aaeaca76031abe4e7 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Thu, 24 Oct 2024 21:06:36 -0700
Subject: [PATCH 03/30] adding deprecated one

---
 .../spark/sql/avro/SchemaConverters.scala     | 57 +++++++++++++++++++
 1 file changed, 57 insertions(+)
 create mode 100644 connector/avro/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala

diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala b/connector/avro/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala
new file mode 100644
index 0000000000000..cd00f19571385
--- /dev/null
+++ b/connector/avro/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala
@@ -0,0 +1,57 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.avro
+
+import org.apache.avro.Schema
+
+import org.apache.spark.annotation.Evolving
+import org.apache.spark.sql.avro.{SchemaConverters => CoreSchemaConverters}
+
+@deprecated("Use SchemaConverters in sql/core instead", "4.0.0")
+@Evolving
+object SchemaConverters {
+  @deprecated("Use SchemaConverters in sql/core instead", "4.0.0")
+  type SchemaType = CoreSchemaConverters.SchemaType
+
+  @deprecated("Use SchemaConverters in sql/core instead", "4.0.0")
+  def toSqlType(avroSchema: Schema): SchemaType = {
+    CoreSchemaConverters.toSqlType(avroSchema)
+  }
+
+  @deprecated("Use SchemaConverters in sql/core instead", "4.0.0")
+  def toSqlType(
+      avroSchema: Schema,
+      useStableIdForUnionType: Boolean,
+      stableIdPrefixForUnionType: String,
+      recursiveFieldMaxDepth: Int = -1): SchemaType = {
+    CoreSchemaConverters.toSqlType(
+      avroSchema,
+      useStableIdForUnionType,
+      stableIdPrefixForUnionType,
+      recursiveFieldMaxDepth)
+  }
+
+  @deprecated("Use SchemaConverters in sql/core instead", "4.0.0")
+  def toAvroType(
+      catalystType: org.apache.spark.sql.types.DataType,
+      nullable: Boolean = false,
+      recordName: String = "topLevelRecord",
+      nameSpace: String = ""): Schema = {
+    CoreSchemaConverters.toAvroType(catalystType, nullable, recordName, nameSpace)
+  }
+}
\ No newline at end of file

From b037859a5fac02d4540aa34e31d3f4854c51990a Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Thu, 24 Oct 2024 23:30:57 -0700
Subject: [PATCH 04/30] init

---
 .../spark/sql/avro/AvroDataToCatalyst.scala   |   1 +
 .../spark/sql/avro/SchemaConverters.scala     |  57 ----
 .../spark/sql/catalyst/StructFilters.scala    |   3 +-
 .../apache/spark/sql/internal/SQLConf.scala   |  10 +
 .../spark/sql/avro/AvroDeserializer.scala     |   4 +-
 .../spark/sql/avro/AvroSerializer.scala       |   2 +-
 .../execution/streaming/ListStateImpl.scala   | 274 ++++++++++++------
 .../streaming/ListStateImplWithTTL.scala      |   5 +-
 .../execution/streaming/MapStateImpl.scala    |   7 +-
 .../streaming/MapStateImplWithTTL.scala       |   5 +-
 .../StateStoreColumnFamilySchemaUtils.scala   |  41 ++-
 .../streaming/StateTypesEncoderUtils.scala    | 126 +++++++-
 .../StatefulProcessorHandleImpl.scala         |  58 +++-
 .../StreamingSymmetricHashJoinHelper.scala    |  65 ++++-
 .../streaming/TransformWithStateExec.scala    |  56 ++--
 .../execution/streaming/ValueStateImpl.scala  |  51 +++-
 .../streaming/ValueStateImplWithTTL.scala     |   5 +-
 .../state/HDFSBackedStateStoreProvider.scala  |  47 +++
 .../streaming/state/RocksDBStateEncoder.scala | 224 ++++++++++++++
 .../state/RocksDBStateStoreProvider.scala     |  92 ++++++
 .../StateSchemaCompatibilityChecker.scala     |  12 +-
 .../streaming/state/StateStore.scala          |  64 ++++
 .../execution/streaming/state/package.scala   |  43 +++
 .../streaming/statefulOperators.scala         |   3 +
 .../streaming/state/MemoryStateStore.scala    |  30 ++
 ...sDBStateStoreCheckpointFormatV2Suite.scala |  29 ++
 .../streaming/state/RocksDBSuite.scala        |  17 ++
 .../streaming/state/ValueStateSuite.scala     |  76 +++--
 .../TransformWithListStateSuite.scala         |  29 +-
 .../streaming/TransformWithStateSuite.scala   | 130 +++++----
 30 files changed, 1242 insertions(+), 324 deletions(-)
 delete mode 100644 connector/avro/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala

diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroDataToCatalyst.scala b/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroDataToCatalyst.scala
index 9c8b2d0375588..cfe98c13f8c07 100644
--- a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroDataToCatalyst.scala
+++ b/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroDataToCatalyst.scala
@@ -24,6 +24,7 @@ import org.apache.avro.generic.GenericDatumReader
 import org.apache.avro.io.{BinaryDecoder, DecoderFactory}
 
 import org.apache.spark.SparkException
+import org.apache.spark.sql.avro.SchemaConverters
 import org.apache.spark.sql.catalyst.expressions.{ExpectsInputTypes, Expression, SpecificInternalRow, UnaryExpression}
 import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodeGenerator, ExprCode}
 import org.apache.spark.sql.catalyst.util.{FailFastMode, ParseMode, PermissiveMode}
diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala b/connector/avro/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala
deleted file mode 100644
index cd00f19571385..0000000000000
--- a/connector/avro/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala
+++ /dev/null
@@ -1,57 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.sql.avro
-
-import org.apache.avro.Schema
-
-import org.apache.spark.annotation.Evolving
-import org.apache.spark.sql.avro.{SchemaConverters => CoreSchemaConverters}
-
-@deprecated("Use SchemaConverters in sql/core instead", "4.0.0")
-@Evolving
-object SchemaConverters {
-  @deprecated("Use SchemaConverters in sql/core instead", "4.0.0")
-  type SchemaType = CoreSchemaConverters.SchemaType
-
-  @deprecated("Use SchemaConverters in sql/core instead", "4.0.0")
-  def toSqlType(avroSchema: Schema): SchemaType = {
-    CoreSchemaConverters.toSqlType(avroSchema)
-  }
-
-  @deprecated("Use SchemaConverters in sql/core instead", "4.0.0")
-  def toSqlType(
-      avroSchema: Schema,
-      useStableIdForUnionType: Boolean,
-      stableIdPrefixForUnionType: String,
-      recursiveFieldMaxDepth: Int = -1): SchemaType = {
-    CoreSchemaConverters.toSqlType(
-      avroSchema,
-      useStableIdForUnionType,
-      stableIdPrefixForUnionType,
-      recursiveFieldMaxDepth)
-  }
-
-  @deprecated("Use SchemaConverters in sql/core instead", "4.0.0")
-  def toAvroType(
-      catalystType: org.apache.spark.sql.types.DataType,
-      nullable: Boolean = false,
-      recordName: String = "topLevelRecord",
-      nameSpace: String = ""): Schema = {
-    CoreSchemaConverters.toAvroType(catalystType, nullable, recordName, nameSpace)
-  }
-}
\ No newline at end of file
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/StructFilters.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/StructFilters.scala
index 1b2013d87eedf..34a23cabf6f7c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/StructFilters.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/StructFilters.scala
@@ -36,7 +36,8 @@ import org.apache.spark.util.ArrayImplicits._
  *                      the fields of the provided schema.
  * @param schema The required schema of records from datasource files.
  */
-abstract class StructFilters(pushedFilters: Seq[sources.Filter], schema: StructType) {
+abstract class StructFilters(
+    pushedFilters: Seq[sources.Filter], schema: StructType) extends Serializable {
 
   protected val filters = StructFilters.pushedFilters(pushedFilters.toArray, schema)
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index 1c9f5e85d1a06..a8d2af606e994 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -2204,6 +2204,16 @@ object SQLConf {
       .intConf
       .createWithDefault(3)
 
+  val STREAMING_STATE_STORE_ENCODING_FORMAT =
+    buildConf("spark.sql.streaming.stateStore.encodingFormat")
+      .doc("The encoding format used for stateful operators to store information" +
+        "in the state store")
+      .version("4.0.0")
+      .stringConf
+      .checkValue(v => Set("UnsafeRow", "Avro").contains(v),
+        "Valid versions are 'UnsafeRow' and 'Avro'")
+      .createWithDefault("UnsafeRow")
+
   // The feature is still in development, so it is still internal.
   val STATE_STORE_CHECKPOINT_FORMAT_VERSION =
     buildConf("spark.sql.streaming.stateStore.checkpointFormatVersion")
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala
index 4e559f6eee887..18eea38930b6d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala
@@ -52,7 +52,7 @@ private[sql] class AvroDeserializer(
     filters: StructFilters,
     useStableIdForUnionType: Boolean,
     stableIdPrefixForUnionType: String,
-    recursiveFieldMaxDepth: Int) {
+    recursiveFieldMaxDepth: Int) extends Serializable {
 
   def this(
       rootAvroType: Schema,
@@ -463,7 +463,7 @@ private[sql] class AvroDeserializer(
    * A base interface for updating values inside catalyst data structure like `InternalRow` and
    * `ArrayData`.
    */
-  sealed trait CatalystDataUpdater {
+  sealed trait CatalystDataUpdater extends Serializable {
     def set(ordinal: Int, value: Any): Unit
 
     def setNullAt(ordinal: Int): Unit = set(ordinal, null)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroSerializer.scala b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroSerializer.scala
index 814a28e24f522..3aefb47d20825 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroSerializer.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroSerializer.scala
@@ -47,7 +47,7 @@ private[sql] class AvroSerializer(
     rootAvroType: Schema,
     nullable: Boolean,
     positionalFieldMatch: Boolean,
-    datetimeRebaseMode: LegacyBehaviorPolicy.Value) extends Logging {
+    datetimeRebaseMode: LegacyBehaviorPolicy.Value) extends Logging with Serializable {
 
   def this(rootCatalystType: DataType, rootAvroType: Schema, nullable: Boolean) = {
     this(rootCatalystType, rootAvroType, nullable, positionalFieldMatch = false,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
index 77c481a8ba0ba..50294fa3d0587 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
@@ -20,7 +20,7 @@ import org.apache.spark.internal.Logging
 import org.apache.spark.sql.Encoder
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.execution.metric.SQLMetric
-import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, StateStore, StateStoreErrors}
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, NoPrefixKeyStateEncoderSpec, StateStore, StateStoreErrors}
 import org.apache.spark.sql.streaming.ListState
 import org.apache.spark.sql.types.StructType
 
@@ -40,6 +40,7 @@ class ListStateImpl[S](
      stateName: String,
      keyExprEnc: ExpressionEncoder[Any],
      valEncoder: Encoder[S],
+     avroEnc: Option[AvroEncoderSpec],
      metrics: Map[String, SQLMetric] = Map.empty)
   extends ListStateMetricsImpl
   with ListState[S]
@@ -49,101 +50,186 @@ class ListStateImpl[S](
   override def baseStateName: String = stateName
   override def exprEncSchema: StructType = keyExprEnc.schema
 
-  private val stateTypesEncoder = StateTypesEncoder(keyExprEnc, valEncoder, stateName)
+  // If we are using Avro, the avroSerde parameter must be populated
+  // else, we will default to using UnsafeRow.
+  private val usingAvro: Boolean = avroEnc.isDefined
+  private val avroTypesEncoder = new AvroTypesEncoder[S](
+    keyExprEnc, valEncoder, stateName, hasTtl = false, avroEnc)
+  private val unsafeRowTypesEncoder = new UnsafeRowTypesEncoder[S](
+    keyExprEnc, valEncoder, stateName, hasTtl = false)
 
   store.createColFamilyIfAbsent(stateName, keyExprEnc.schema, valEncoder.schema,
     NoPrefixKeyStateEncoderSpec(keyExprEnc.schema), useMultipleValuesPerKey = true)
 
   /** Whether state exists or not. */
-   override def exists(): Boolean = {
-     val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
-     val stateValue = store.get(encodedGroupingKey, stateName)
-     stateValue != null
-   }
-
-   /**
-    * Get the state value if it exists. If the state does not exist in state store, an
-    * empty iterator is returned.
-    */
-   override def get(): Iterator[S] = {
-     val encodedKey = stateTypesEncoder.encodeGroupingKey()
-     val unsafeRowValuesIterator = store.valuesIterator(encodedKey, stateName)
-     new Iterator[S] {
-       override def hasNext: Boolean = {
-         unsafeRowValuesIterator.hasNext
-       }
-
-       override def next(): S = {
-         val valueUnsafeRow = unsafeRowValuesIterator.next()
-         stateTypesEncoder.decodeValue(valueUnsafeRow)
-       }
-     }
-   }
-
-   /** Update the value of the list. */
-   override def put(newState: Array[S]): Unit = {
-     validateNewState(newState)
-
-     val encodedKey = stateTypesEncoder.encodeGroupingKey()
-     var isFirst = true
-     var entryCount = 0L
-     TWSMetricsUtils.resetMetric(metrics, "numUpdatedStateRows")
-
-     newState.foreach { v =>
-       val encodedValue = stateTypesEncoder.encodeValue(v)
-       if (isFirst) {
-         store.put(encodedKey, encodedValue, stateName)
-         isFirst = false
-       } else {
-         store.merge(encodedKey, encodedValue, stateName)
-       }
-       entryCount += 1
-       TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
-     }
-     updateEntryCount(encodedKey, entryCount)
-   }
-
-   /** Append an entry to the list. */
-   override def appendValue(newState: S): Unit = {
-     StateStoreErrors.requireNonNullStateValue(newState, stateName)
-     val encodedKey = stateTypesEncoder.encodeGroupingKey()
-     val entryCount = getEntryCount(encodedKey)
-     store.merge(encodedKey,
-         stateTypesEncoder.encodeValue(newState), stateName)
-     TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
-     updateEntryCount(encodedKey, entryCount + 1)
-   }
-
-   /** Append an entire list to the existing value. */
-   override def appendList(newState: Array[S]): Unit = {
-     validateNewState(newState)
-
-     val encodedKey = stateTypesEncoder.encodeGroupingKey()
-     var entryCount = getEntryCount(encodedKey)
-     newState.foreach { v =>
-       val encodedValue = stateTypesEncoder.encodeValue(v)
-       store.merge(encodedKey, encodedValue, stateName)
-       entryCount += 1
-       TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
-     }
-     updateEntryCount(encodedKey, entryCount)
-   }
-
-   /** Remove this state. */
-   override def clear(): Unit = {
-     val encodedKey = stateTypesEncoder.encodeGroupingKey()
-     store.remove(encodedKey, stateName)
-     val entryCount = getEntryCount(encodedKey)
-     TWSMetricsUtils.incrementMetric(metrics, "numRemovedStateRows", entryCount)
-     removeEntryCount(encodedKey)
-   }
-
-   private def validateNewState(newState: Array[S]): Unit = {
-     StateStoreErrors.requireNonNullStateValue(newState, stateName)
-     StateStoreErrors.requireNonEmptyListStateValue(newState, stateName)
-
-     newState.foreach { v =>
-       StateStoreErrors.requireNonNullStateValue(v, stateName)
-     }
-   }
- }
+  override def exists(): Boolean = {
+    if (usingAvro) {
+      val encodedKey: Array[Byte] = avroTypesEncoder.encodeGroupingKey()
+      store.get(encodedKey, stateName) != null
+    } else {
+      val encodedKey = unsafeRowTypesEncoder.encodeGroupingKey()
+      store.get(encodedKey, stateName) != null
+    }
+  }
+
+  /**
+   * Get the state value if it exists. If the state does not exist in state store, an
+   * empty iterator is returned.
+   */
+  override def get(): Iterator[S] = {
+    if (usingAvro) {
+      getAvro()
+    } else {
+      getUnsafeRow()
+    }
+  }
+
+  private def getAvro(): Iterator[S] = {
+    val encodedKey: Array[Byte] = avroTypesEncoder.encodeGroupingKey()
+    val avroValuesIterator = store.valuesIterator(encodedKey, stateName)
+    new Iterator[S] {
+      override def hasNext: Boolean = {
+        avroValuesIterator.hasNext
+      }
+
+      override def next(): S = {
+        val valueRow = avroValuesIterator.next()
+        avroTypesEncoder.decodeValue(valueRow)
+      }
+    }
+  }
+
+  private def getUnsafeRow(): Iterator[S] = {
+    val encodedKey = unsafeRowTypesEncoder.encodeGroupingKey()
+    val unsafeRowValuesIterator = store.valuesIterator(encodedKey, stateName)
+    new Iterator[S] {
+      override def hasNext: Boolean = {
+        unsafeRowValuesIterator.hasNext
+      }
+
+      override def next(): S = {
+        val valueUnsafeRow = unsafeRowValuesIterator.next()
+        unsafeRowTypesEncoder.decodeValue(valueUnsafeRow)
+      }
+    }
+  }
+
+  /** Update the value of the list. */
+  override def put(newState: Array[S]): Unit = {
+    validateNewState(newState)
+
+    if (usingAvro) {
+      putAvro(newState)
+    } else {
+      putUnsafeRow(newState)
+    }
+  }
+
+  private def putAvro(newState: Array[S]): Unit = {
+    val encodedKey: Array[Byte] = avroTypesEncoder.encodeGroupingKey()
+    var isFirst = true
+    var entryCount = 0L
+    TWSMetricsUtils.resetMetric(metrics, "numUpdatedStateRows")
+
+    newState.foreach { v =>
+      val encodedValue = avroTypesEncoder.encodeValue(v)
+      if (isFirst) {
+        store.put(encodedKey, encodedValue, stateName)
+        isFirst = false
+      } else {
+        store.merge(encodedKey, encodedValue, stateName)
+      }
+      entryCount += 1
+      TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
+    }
+  }
+
+  private def putUnsafeRow(newState: Array[S]): Unit = {
+    val encodedKey = unsafeRowTypesEncoder.encodeGroupingKey()
+    var isFirst = true
+    var entryCount = 0L
+    TWSMetricsUtils.resetMetric(metrics, "numUpdatedStateRows")
+
+    newState.foreach { v =>
+      val encodedValue = unsafeRowTypesEncoder.encodeValue(v)
+      if (isFirst) {
+        store.put(encodedKey, encodedValue, stateName)
+        isFirst = false
+      } else {
+        store.merge(encodedKey, encodedValue, stateName)
+      }
+      entryCount += 1
+      TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
+    }
+    updateEntryCount(encodedKey, entryCount)
+  }
+
+  /** Append an entry to the list. */
+  override def appendValue(newState: S): Unit = {
+    StateStoreErrors.requireNonNullStateValue(newState, stateName)
+
+    if (usingAvro) {
+      val encodedKey: Array[Byte] = avroTypesEncoder.encodeGroupingKey()
+      val encodedValue = avroTypesEncoder.encodeValue(newState)
+      store.merge(encodedKey, encodedValue, stateName)
+      TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
+    } else {
+      val encodedKey = unsafeRowTypesEncoder.encodeGroupingKey()
+      val entryCount = getEntryCount(encodedKey)
+      val encodedValue = unsafeRowTypesEncoder.encodeValue(newState)
+      store.merge(encodedKey, encodedValue, stateName)
+      TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
+      updateEntryCount(encodedKey, entryCount + 1)
+    }
+  }
+
+  /** Append an entire list to the existing value. */
+  override def appendList(newState: Array[S]): Unit = {
+    validateNewState(newState)
+
+    if (usingAvro) {
+      val encodedKey: Array[Byte] = avroTypesEncoder.encodeGroupingKey()
+      newState.foreach { v =>
+        val encodedValue = avroTypesEncoder.encodeValue(v)
+        store.merge(encodedKey, encodedValue, stateName)
+        TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
+      }
+    } else {
+      val encodedKey = unsafeRowTypesEncoder.encodeGroupingKey()
+      var entryCount = getEntryCount(encodedKey)
+      newState.foreach { v =>
+        val encodedValue = unsafeRowTypesEncoder.encodeValue(v)
+        store.merge(encodedKey, encodedValue, stateName)
+        entryCount += 1
+        TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
+      }
+      updateEntryCount(encodedKey, entryCount)
+    }
+  }
+
+  /** Remove this state. */
+  override def clear(): Unit = {
+    if (usingAvro) {
+      val encodedKey: Array[Byte] = avroTypesEncoder.encodeGroupingKey()
+      store.remove(encodedKey, stateName)
+       val entryCount = getEntryCount(encodedKey)
+       TWSMetricsUtils.incrementMetric(metrics, "numRemovedStateRows", entryCount)
+       removeEntryCount(encodedKey)
+    } else {
+      val encodedKey = unsafeRowTypesEncoder.encodeGroupingKey()
+      store.remove(encodedKey, stateName)
+      val entryCount = getEntryCount(encodedKey)
+      TWSMetricsUtils.incrementMetric(metrics, "numRemovedStateRows", entryCount)
+      removeEntryCount(encodedKey)
+    }
+  }
+
+  private def validateNewState(newState: Array[S]): Unit = {
+    StateStoreErrors.requireNonNullStateValue(newState, stateName)
+    StateStoreErrors.requireNonEmptyListStateValue(newState, stateName)
+
+    newState.foreach { v =>
+      StateStoreErrors.requireNonNullStateValue(v, stateName)
+    }
+  }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala
index be47f566bc6a9..639683e5ff549 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala
@@ -21,7 +21,7 @@ import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, StateStore, StateStoreErrors}
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, NoPrefixKeyStateEncoderSpec, StateStore, StateStoreErrors}
 import org.apache.spark.sql.streaming.{ListState, TTLConfig}
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.util.NextIterator
@@ -46,6 +46,7 @@ class ListStateImplWithTTL[S](
     valEncoder: Encoder[S],
     ttlConfig: TTLConfig,
     batchTimestampMs: Long,
+    avroEnc: Option[AvroEncoderSpec], // TODO: Add Avro Encoding support for TTL
     metrics: Map[String, SQLMetric] = Map.empty)
   extends SingleKeyTTLStateImpl(stateName, store, keyExprEnc, batchTimestampMs)
   with ListStateMetricsImpl
@@ -55,7 +56,7 @@ class ListStateImplWithTTL[S](
   override def baseStateName: String = stateName
   override def exprEncSchema: StructType = keyExprEnc.schema
 
-  private lazy val stateTypesEncoder = StateTypesEncoder(keyExprEnc, valEncoder,
+  private lazy val stateTypesEncoder = UnsafeRowTypesEncoder(keyExprEnc, valEncoder,
     stateName, hasTtl = true)
 
   private lazy val ttlExpirationMs =
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala
index cb3db19496dd2..b723020d98e02 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala
@@ -21,7 +21,7 @@ import org.apache.spark.sql.Encoder
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{PrefixKeyScanStateEncoderSpec, StateStore, StateStoreErrors, UnsafeRowPair}
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, PrefixKeyScanStateEncoderSpec, StateStore, StateStoreErrors, UnsafeRowPair}
 import org.apache.spark.sql.streaming.MapState
 import org.apache.spark.sql.types.StructType
 
@@ -42,6 +42,7 @@ class MapStateImpl[K, V](
     keyExprEnc: ExpressionEncoder[Any],
     userKeyEnc: Encoder[K],
     valEncoder: Encoder[V],
+    avroEnc: Option[AvroEncoderSpec],
     metrics: Map[String, SQLMetric] = Map.empty) extends MapState[K, V] with Logging {
 
   // Pack grouping key and user key together as a prefixed composite key
@@ -49,8 +50,8 @@ class MapStateImpl[K, V](
     getCompositeKeySchema(keyExprEnc.schema, userKeyEnc.schema)
   }
   private val schemaForValueRow: StructType = valEncoder.schema
-  private val stateTypesEncoder = new CompositeKeyStateEncoder(
-    keyExprEnc, userKeyEnc, valEncoder, stateName)
+  private val stateTypesEncoder = new CompositeKeyUnsafeRowEncoder(
+    keyExprEnc, userKeyEnc, valEncoder, stateName, hasTtl = false)
 
   store.createColFamilyIfAbsent(stateName, schemaForCompositeKeyRow, schemaForValueRow,
     PrefixKeyScanStateEncoderSpec(schemaForCompositeKeyRow, 1))
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala
index 6a3685ad6c46c..4020b1b4fd904 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala
@@ -22,7 +22,7 @@ import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{PrefixKeyScanStateEncoderSpec, StateStore, StateStoreErrors}
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, PrefixKeyScanStateEncoderSpec, StateStore, StateStoreErrors}
 import org.apache.spark.sql.streaming.{MapState, TTLConfig}
 import org.apache.spark.util.NextIterator
 
@@ -49,12 +49,13 @@ class MapStateImplWithTTL[K, V](
     valEncoder: Encoder[V],
     ttlConfig: TTLConfig,
     batchTimestampMs: Long,
+    avroEnc: Option[AvroEncoderSpec], // TODO: Add Avro Encoding support for TTL
     metrics: Map[String, SQLMetric] = Map.empty)
   extends CompositeKeyTTLStateImpl[K](stateName, store,
     keyExprEnc, userKeyEnc, batchTimestampMs)
   with MapState[K, V] with Logging {
 
-  private val stateTypesEncoder = new CompositeKeyStateEncoder(
+  private val stateTypesEncoder = new CompositeKeyUnsafeRowEncoder(
     keyExprEnc, userKeyEnc, valEncoder, stateName, hasTtl = true)
 
   private val ttlExpirationMs =
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
index 7da8408f98b0f..2214af226b300 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
@@ -17,23 +17,50 @@
 package org.apache.spark.sql.execution.streaming
 
 import org.apache.spark.sql.Encoder
+import org.apache.spark.sql.avro.{AvroDeserializer, AvroOptions, AvroSerializer, SchemaConverters}
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
 import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, PrefixKeyScanStateEncoderSpec, StateStoreColFamilySchema}
+import org.apache.spark.sql.execution.streaming.state.AvroEncoderSpec
 import org.apache.spark.sql.types.StructType
 
 object StateStoreColumnFamilySchemaUtils {
 
+  def apply(initializeAvroSerde: Boolean): StateStoreColumnFamilySchemaUtils =
+    new StateStoreColumnFamilySchemaUtils(initializeAvroSerde)
+}
+
+class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) {
+
+  private def getAvroSerde(
+      keySchema: StructType, valSchema: StructType): Option[AvroEncoderSpec] = {
+    if (initializeAvroSerde) {
+      val avroType = SchemaConverters.toAvroType(valSchema)
+      val avroOptions = AvroOptions(Map.empty)
+      val keyAvroType = SchemaConverters.toAvroType(keySchema)
+      val keySer = new AvroSerializer(keySchema, keyAvroType, nullable = false)
+      val ser = new AvroSerializer(valSchema, avroType, nullable = false)
+      val de = new AvroDeserializer(avroType, valSchema,
+        avroOptions.datetimeRebaseModeInRead, avroOptions.useStableIdForUnionType,
+        avroOptions.stableIdPrefixForUnionType, avroOptions.recursiveFieldMaxDepth)
+      Some(AvroEncoderSpec(keySer, ser, de))
+    } else {
+      None
+    }
+  }
+
   def getValueStateSchema[T](
       stateName: String,
       keyEncoder: ExpressionEncoder[Any],
       valEncoder: Encoder[T],
       hasTtl: Boolean): StateStoreColFamilySchema = {
+   val valSchema = getValueSchemaWithTTL(valEncoder.schema, hasTtl)
    StateStoreColFamilySchema(
       stateName,
       keyEncoder.schema,
-      getValueSchemaWithTTL(valEncoder.schema, hasTtl),
-      Some(NoPrefixKeyStateEncoderSpec(keyEncoder.schema)))
+      valSchema,
+      Some(NoPrefixKeyStateEncoderSpec(keyEncoder.schema)),
+      avroEnc = getAvroSerde(keyEncoder.schema, valSchema))
   }
 
   def getListStateSchema[T](
@@ -41,11 +68,13 @@ object StateStoreColumnFamilySchemaUtils {
       keyEncoder: ExpressionEncoder[Any],
       valEncoder: Encoder[T],
       hasTtl: Boolean): StateStoreColFamilySchema = {
+  val valSchema = getValueSchemaWithTTL(valEncoder.schema, hasTtl)
   StateStoreColFamilySchema(
       stateName,
       keyEncoder.schema,
-      getValueSchemaWithTTL(valEncoder.schema, hasTtl),
-      Some(NoPrefixKeyStateEncoderSpec(keyEncoder.schema)))
+      valSchema,
+      Some(NoPrefixKeyStateEncoderSpec(keyEncoder.schema)),
+      avroEnc = getAvroSerde(keyEncoder.schema, valSchema))
   }
 
   def getMapStateSchema[K, V](
@@ -55,12 +84,14 @@ object StateStoreColumnFamilySchemaUtils {
       valEncoder: Encoder[V],
       hasTtl: Boolean): StateStoreColFamilySchema = {
     val compositeKeySchema = getCompositeKeySchema(keyEncoder.schema, userKeyEnc.schema)
+    val valSchema = getValueSchemaWithTTL(valEncoder.schema, hasTtl)
     StateStoreColFamilySchema(
       stateName,
       compositeKeySchema,
       getValueSchemaWithTTL(valEncoder.schema, hasTtl),
       Some(PrefixKeyScanStateEncoderSpec(compositeKeySchema, 1)),
-      Some(userKeyEnc.schema))
+      Some(userKeyEnc.schema),
+      avroEnc = getAvroSerde(compositeKeySchema, valSchema))
   }
 
   def getTimerStateSchema(
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala
index b70f9699195d4..dcfc797586aac 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala
@@ -17,12 +17,18 @@
 
 package org.apache.spark.sql.execution.streaming
 
+import java.io.ByteArrayOutputStream
+
+import org.apache.avro.generic.{GenericDatumReader, GenericDatumWriter}
+import org.apache.avro.io.{DecoderFactory, EncoderFactory}
+
 import org.apache.spark.sql.Encoder
+import org.apache.spark.sql.avro.SchemaConverters
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.encoders.{encoderFor, ExpressionEncoder}
 import org.apache.spark.sql.catalyst.expressions.{UnsafeProjection, UnsafeRow}
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.StateStoreErrors
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, StateStoreErrors}
 import org.apache.spark.sql.types._
 
 /**
@@ -59,6 +65,15 @@ object TransformWithStateKeyValueRowSchemaUtils {
   }
 }
 
+trait StateTypesEncoder[V, S] {
+  def encodeGroupingKey(): S
+  def encodeValue(value: V): S
+  def decodeValue(row: S): V
+  def encodeValue(value: V, expirationMs: Long): S
+  def decodeTtlExpirationMs(row: S): Option[Long]
+  def isExpired(row: S, batchTimestampMs: Long): Boolean
+}
+
 /**
  * Helper class providing APIs to encode the grouping key, and user provided values
  * to Spark [[UnsafeRow]].
@@ -73,11 +88,11 @@ object TransformWithStateKeyValueRowSchemaUtils {
  * @param stateName - name of logical state partition
  * @tparam V - value type
  */
-class StateTypesEncoder[V](
+class UnsafeRowTypesEncoder[V](
     keyEncoder: ExpressionEncoder[Any],
     valEncoder: Encoder[V],
     stateName: String,
-    hasTtl: Boolean) {
+    hasTtl: Boolean) extends StateTypesEncoder[V, UnsafeRow] {
 
   /** Variables reused for value conversions between spark sql and object */
   private val keySerializer = keyEncoder.createSerializer()
@@ -143,23 +158,118 @@ class StateTypesEncoder[V](
   }
 }
 
-object StateTypesEncoder {
+
+object UnsafeRowTypesEncoder {
   def apply[V](
       keyEncoder: ExpressionEncoder[Any],
       valEncoder: Encoder[V],
       stateName: String,
-      hasTtl: Boolean = false): StateTypesEncoder[V] = {
-    new StateTypesEncoder[V](keyEncoder, valEncoder, stateName, hasTtl)
+      hasTtl: Boolean = false): UnsafeRowTypesEncoder[V] = {
+    new UnsafeRowTypesEncoder[V](keyEncoder, valEncoder, stateName, hasTtl)
+  }
+}
+
+/**
+ * Helper class providing APIs to encode the grouping key, and user provided values
+ * to Spark [[UnsafeRow]].
+ *
+ * CAUTION: StateTypesEncoder class instance is *not* thread-safe.
+ * This class reuses the keyProjection and valueProjection for encoding grouping
+ * key and state value respectively. As UnsafeProjection is not thread safe, this
+ * class is also not thread safe.
+ *
+ * @param keyEncoder - SQL encoder for the grouping key, key type is implicit
+ * @param valEncoder - SQL encoder for value of type `S`
+ * @param stateName - name of logical state partition
+ * @tparam V - value type
+ */
+class AvroTypesEncoder[V](
+    keyEncoder: ExpressionEncoder[Any],
+    valEncoder: Encoder[V],
+    stateName: String,
+    hasTtl: Boolean,
+    avroSerde: Option[AvroEncoderSpec]) extends StateTypesEncoder[V, Array[Byte]] {
+
+  val out = new ByteArrayOutputStream
+
+  /** Variables reused for value conversions between spark sql and object */
+  private val keySerializer = keyEncoder.createSerializer()
+  private val valExpressionEnc = encoderFor(valEncoder)
+  private val objToRowSerializer = valExpressionEnc.createSerializer()
+  private val rowToObjDeserializer = valExpressionEnc.resolveAndBind().createDeserializer()
+
+  private val keySchema = keyEncoder.schema
+  private val keyAvroType = SchemaConverters.toAvroType(keySchema)
+
+  // case class -> dataType
+  private val valSchema: StructType = valEncoder.schema
+  // dataType -> avroType
+  private val valueAvroType = SchemaConverters.toAvroType(valSchema)
+
+  override def encodeGroupingKey(): Array[Byte] = {
+    val keyOption = ImplicitGroupingKeyTracker.getImplicitKeyOption
+    if (keyOption.isEmpty) {
+      throw StateStoreErrors.implicitKeyNotFound(stateName)
+    }
+
+    val keyRow = keySerializer.apply(keyOption.get).copy() // V -> InternalRow
+    val avroData = avroSerde.get.keySerializer.serialize(keyRow) // InternalRow -> GenericDataRecord
+
+    out.reset()
+    val encoder = EncoderFactory.get().directBinaryEncoder(out, null)
+    val writer = new GenericDatumWriter[Any](keyAvroType)
+
+    writer.write(avroData, encoder)
+    encoder.flush()
+    out.toByteArray
+  }
+
+  override def encodeValue(value: V): Array[Byte] = {
+    val objRow: InternalRow = objToRowSerializer.apply(value).copy() // V -> InternalRow
+    val avroData =
+      avroSerde.get.valueSerializer.serialize(objRow) // InternalRow -> GenericDataRecord
+    out.reset()
+
+    val encoder = EncoderFactory.get().directBinaryEncoder(out, null)
+    val writer = new GenericDatumWriter[Any](
+      valueAvroType) // Defining Avro writer for this struct type
+
+    writer.write(avroData, encoder) // GenericDataRecord -> bytes
+    encoder.flush()
+    out.toByteArray
+  }
+
+  override def decodeValue(row: Array[Byte]): V = {
+    val reader = new GenericDatumReader[Any](valueAvroType)
+    val decoder = DecoderFactory.get().binaryDecoder(row, 0, row.length, null)
+    val genericData = reader.read(null, decoder) // bytes -> GenericDataRecord
+    val internalRow = avroSerde.get.valueDeserializer.deserialize(
+      genericData).orNull.asInstanceOf[InternalRow] // GenericDataRecord -> InternalRow
+    if (hasTtl) {
+      rowToObjDeserializer.apply(internalRow.getStruct(0, valEncoder.schema.length))
+    } else rowToObjDeserializer.apply(internalRow)
+  }
+
+  override def encodeValue(value: V, expirationMs: Long): Array[Byte] = {
+    throw new UnsupportedOperationException
+  }
+
+  override def decodeTtlExpirationMs(row: Array[Byte]): Option[Long] = {
+    throw new UnsupportedOperationException
+  }
+
+  override def isExpired(row: Array[Byte], batchTimestampMs: Long): Boolean = {
+    throw new UnsupportedOperationException
   }
 }
 
-class CompositeKeyStateEncoder[K, V](
+class CompositeKeyUnsafeRowEncoder[K, V](
     keyEncoder: ExpressionEncoder[Any],
     userKeyEnc: Encoder[K],
     valEncoder: Encoder[V],
     stateName: String,
     hasTtl: Boolean = false)
-  extends StateTypesEncoder[V](keyEncoder, valEncoder, stateName, hasTtl) {
+  extends UnsafeRowTypesEncoder[V](keyEncoder, valEncoder, stateName, hasTtl) {
   import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
 
   /** Encoders */
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
index 762dfc7d08920..88f0be0b2269f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
@@ -104,7 +104,8 @@ class StatefulProcessorHandleImpl(
     timeMode: TimeMode,
     isStreaming: Boolean = true,
     batchTimestampMs: Option[Long] = None,
-    metrics: Map[String, SQLMetric] = Map.empty)
+    metrics: Map[String, SQLMetric] = Map.empty,
+    schemas: Map[String, StateStoreColFamilySchema] = Map.empty)
   extends StatefulProcessorHandleImplBase(timeMode, keyEncoder) with Logging {
   import StatefulProcessorHandleState._
 
@@ -139,11 +140,30 @@ class StatefulProcessorHandleImpl(
       stateName: String,
       valEncoder: Encoder[T]): ValueState[T] = {
     verifyStateVarOperations("get_value_state", CREATED)
-    val resultState = new ValueStateImpl[T](store, stateName, keyEncoder, valEncoder, metrics)
+    val resultState = new ValueStateImpl[T](
+      store, stateName, keyEncoder, valEncoder, schemas(stateName).avroEnc, metrics)
     TWSMetricsUtils.incrementMetric(metrics, "numValueStateVars")
     resultState
   }
 
+  // This method is for unit-testing ValueState, as the avroEnc will not be
+  // populated unless the handle is created through the TransformWithStateExec operator
+  private[sql] def getValueStateWithAvro[T](
+      stateName: String,
+      valEncoder: Encoder[T],
+      useAvro: Boolean): ValueState[T] = {
+    verifyStateVarOperations("get_value_state", CREATED)
+    val avroEnc = if (useAvro) {
+      new StateStoreColumnFamilySchemaUtils(true).getValueStateSchema[T](
+        stateName, keyEncoder, valEncoder, hasTtl = false).avroEnc
+    } else {
+      None
+    }
+    val resultState = new ValueStateImpl[T](
+      store, stateName, keyEncoder, valEncoder, avroEnc)
+    resultState
+  }
+
   override def getValueState[T](
       stateName: String,
       valEncoder: Encoder[T],
@@ -153,7 +173,7 @@ class StatefulProcessorHandleImpl(
 
     assert(batchTimestampMs.isDefined)
     val valueStateWithTTL = new ValueStateImplWithTTL[T](store, stateName,
-      keyEncoder, valEncoder, ttlConfig, batchTimestampMs.get, metrics)
+      keyEncoder, valEncoder, ttlConfig, batchTimestampMs.get, avroEnc = None, metrics)
     ttlStates.add(valueStateWithTTL)
     TWSMetricsUtils.incrementMetric(metrics, "numValueStateWithTTLVars")
 
@@ -232,7 +252,8 @@ class StatefulProcessorHandleImpl(
 
   override def getListState[T](stateName: String, valEncoder: Encoder[T]): ListState[T] = {
     verifyStateVarOperations("get_list_state", CREATED)
-    val resultState = new ListStateImpl[T](store, stateName, keyEncoder, valEncoder, metrics)
+    val resultState = new ListStateImpl[T](
+      store, stateName, keyEncoder, valEncoder, schemas(stateName).avroEnc, metrics)
     TWSMetricsUtils.incrementMetric(metrics, "numListStateVars")
     resultState
   }
@@ -262,7 +283,7 @@ class StatefulProcessorHandleImpl(
 
     assert(batchTimestampMs.isDefined)
     val listStateWithTTL = new ListStateImplWithTTL[T](store, stateName,
-      keyEncoder, valEncoder, ttlConfig, batchTimestampMs.get, metrics)
+      keyEncoder, valEncoder, ttlConfig, batchTimestampMs.get, avroEnc = None, metrics)
     TWSMetricsUtils.incrementMetric(metrics, "numListStateWithTTLVars")
     ttlStates.add(listStateWithTTL)
 
@@ -275,7 +296,7 @@ class StatefulProcessorHandleImpl(
       valEncoder: Encoder[V]): MapState[K, V] = {
     verifyStateVarOperations("get_map_state", CREATED)
     val resultState = new MapStateImpl[K, V](store, stateName, keyEncoder,
-      userKeyEnc, valEncoder, metrics)
+      userKeyEnc, valEncoder, avroEnc = None, metrics)
     TWSMetricsUtils.incrementMetric(metrics, "numMapStateVars")
     resultState
   }
@@ -290,7 +311,7 @@ class StatefulProcessorHandleImpl(
 
     assert(batchTimestampMs.isDefined)
     val mapStateWithTTL = new MapStateImplWithTTL[K, V](store, stateName, keyEncoder, userKeyEnc,
-      valEncoder, ttlConfig, batchTimestampMs.get, metrics)
+      valEncoder, ttlConfig, batchTimestampMs.get, avroEnc = None, metrics)
     TWSMetricsUtils.incrementMetric(metrics, "numMapStateWithTTLVars")
     ttlStates.add(mapStateWithTTL)
 
@@ -313,7 +334,8 @@ class StatefulProcessorHandleImpl(
  * actually done. We need this class because we can only collect the schemas after
  * the StatefulProcessor is initialized.
  */
-class DriverStatefulProcessorHandleImpl(timeMode: TimeMode, keyExprEnc: ExpressionEncoder[Any])
+class DriverStatefulProcessorHandleImpl(
+    timeMode: TimeMode, keyExprEnc: ExpressionEncoder[Any], initializeAvroEnc: Boolean)
   extends StatefulProcessorHandleImplBase(timeMode, keyExprEnc) {
 
   // Because this is only happening on the driver side, there is only
@@ -324,6 +346,12 @@ class DriverStatefulProcessorHandleImpl(timeMode: TimeMode, keyExprEnc: Expressi
   private val stateVariableInfos: mutable.Map[String, TransformWithStateVariableInfo] =
     new mutable.HashMap[String, TransformWithStateVariableInfo]()
 
+  // If we want use Avro serializers and deserializers, the schemaUtils will create and populate
+  // these objects as a part of the schema, and will add this to the map
+  // These serde objects will eventually be passed to the executors
+  private val schemaUtils: StateStoreColumnFamilySchemaUtils =
+    new StateStoreColumnFamilySchemaUtils(initializeAvroEnc)
+
   // If timeMode is not None, add a timer column family schema to the operator metadata so that
   // registered timers can be read using the state data source reader.
   if (timeMode != TimeMode.None()) {
@@ -343,7 +371,7 @@ class DriverStatefulProcessorHandleImpl(timeMode: TimeMode, keyExprEnc: Expressi
   private def addTimerColFamily(): Unit = {
     val stateName = TimerStateUtils.getTimerStateVarName(timeMode.toString)
     val timerEncoder = new TimerKeyEncoder(keyExprEnc)
-    val colFamilySchema = StateStoreColumnFamilySchemaUtils.
+    val colFamilySchema = schemaUtils.
       getTimerStateSchema(stateName, timerEncoder.schemaForKeyRow, timerEncoder.schemaForValueRow)
     columnFamilySchemas.put(stateName, colFamilySchema)
     val stateVariableInfo = TransformWithStateVariableUtils.getTimerState(stateName)
@@ -352,7 +380,7 @@ class DriverStatefulProcessorHandleImpl(timeMode: TimeMode, keyExprEnc: Expressi
 
   override def getValueState[T](stateName: String, valEncoder: Encoder[T]): ValueState[T] = {
     verifyStateVarOperations("get_value_state", PRE_INIT)
-    val colFamilySchema = StateStoreColumnFamilySchemaUtils.
+    val colFamilySchema = schemaUtils.
       getValueStateSchema(stateName, keyExprEnc, valEncoder, false)
     checkIfDuplicateVariableDefined(stateName)
     columnFamilySchemas.put(stateName, colFamilySchema)
@@ -367,7 +395,7 @@ class DriverStatefulProcessorHandleImpl(timeMode: TimeMode, keyExprEnc: Expressi
       valEncoder: Encoder[T],
       ttlConfig: TTLConfig): ValueState[T] = {
     verifyStateVarOperations("get_value_state", PRE_INIT)
-    val colFamilySchema = StateStoreColumnFamilySchemaUtils.
+    val colFamilySchema = schemaUtils.
       getValueStateSchema(stateName, keyExprEnc, valEncoder, true)
     checkIfDuplicateVariableDefined(stateName)
     columnFamilySchemas.put(stateName, colFamilySchema)
@@ -379,7 +407,7 @@ class DriverStatefulProcessorHandleImpl(timeMode: TimeMode, keyExprEnc: Expressi
 
   override def getListState[T](stateName: String, valEncoder: Encoder[T]): ListState[T] = {
     verifyStateVarOperations("get_list_state", PRE_INIT)
-    val colFamilySchema = StateStoreColumnFamilySchemaUtils.
+    val colFamilySchema = schemaUtils.
       getListStateSchema(stateName, keyExprEnc, valEncoder, false)
     checkIfDuplicateVariableDefined(stateName)
     columnFamilySchemas.put(stateName, colFamilySchema)
@@ -394,7 +422,7 @@ class DriverStatefulProcessorHandleImpl(timeMode: TimeMode, keyExprEnc: Expressi
       valEncoder: Encoder[T],
       ttlConfig: TTLConfig): ListState[T] = {
     verifyStateVarOperations("get_list_state", PRE_INIT)
-    val colFamilySchema = StateStoreColumnFamilySchemaUtils.
+    val colFamilySchema = schemaUtils.
       getListStateSchema(stateName, keyExprEnc, valEncoder, true)
     checkIfDuplicateVariableDefined(stateName)
     columnFamilySchemas.put(stateName, colFamilySchema)
@@ -409,7 +437,7 @@ class DriverStatefulProcessorHandleImpl(timeMode: TimeMode, keyExprEnc: Expressi
       userKeyEnc: Encoder[K],
       valEncoder: Encoder[V]): MapState[K, V] = {
     verifyStateVarOperations("get_map_state", PRE_INIT)
-    val colFamilySchema = StateStoreColumnFamilySchemaUtils.
+    val colFamilySchema = schemaUtils.
       getMapStateSchema(stateName, keyExprEnc, userKeyEnc, valEncoder, false)
     checkIfDuplicateVariableDefined(stateName)
     columnFamilySchemas.put(stateName, colFamilySchema)
@@ -425,7 +453,7 @@ class DriverStatefulProcessorHandleImpl(timeMode: TimeMode, keyExprEnc: Expressi
       valEncoder: Encoder[V],
       ttlConfig: TTLConfig): MapState[K, V] = {
     verifyStateVarOperations("get_map_state", PRE_INIT)
-    val colFamilySchema = StateStoreColumnFamilySchemaUtils.
+    val colFamilySchema = schemaUtils.
       getMapStateSchema(stateName, keyExprEnc, userKeyEnc, valEncoder, true)
     columnFamilySchemas.put(stateName, colFamilySchema)
     val stateVariableInfo = TransformWithStateVariableUtils.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamingSymmetricHashJoinHelper.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamingSymmetricHashJoinHelper.scala
index 497e71070a09a..468d0df75fee4 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamingSymmetricHashJoinHelper.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamingSymmetricHashJoinHelper.scala
@@ -27,7 +27,7 @@ import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeSet,
 import org.apache.spark.sql.catalyst.plans.logical.EventTimeWatermark._
 import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.execution.streaming.WatermarkSupport.watermarkExpression
-import org.apache.spark.sql.execution.streaming.state.{StateStoreCheckpointInfo, StateStoreCoordinatorRef, StateStoreProviderId}
+import org.apache.spark.sql.execution.streaming.state.{StateStoreCheckpointInfo, StateStoreColFamilySchema, StateStoreCoordinatorRef, StateStoreProviderId}
 
 
 /**
@@ -303,6 +303,56 @@ object StreamingSymmetricHashJoinHelper extends Logging {
     }
   }
 
+  /**
+   * A custom RDD that allows partitions to be "zipped" together, while ensuring the tasks'
+   * preferred location is based on which executors have the required join state stores already
+   * loaded. This class is a variant of [[org.apache.spark.rdd.ZippedPartitionsRDD2]] which only
+   * changes signature of `f` by taking in a map of column family schemas. This is used for
+   * passing the column family schemas when there is initial state for the TransformWithStateExec
+   * operator
+   */
+  class StateStoreAwareZipPartitionsRDDWithSchemas[A: ClassTag, B: ClassTag, V: ClassTag](
+      sc: SparkContext,
+      var f: (Int, Iterator[A], Iterator[B], Map[String, StateStoreColFamilySchema]) => Iterator[V],
+      var rdd1: RDD[A],
+      var rdd2: RDD[B],
+      stateInfo: StatefulOperatorStateInfo,
+      stateStoreNames: Seq[String],
+      @transient private val storeCoordinator: Option[StateStoreCoordinatorRef],
+      schemas: Map[String, StateStoreColFamilySchema])
+    extends ZippedPartitionsBaseRDD[V](sc, List(rdd1, rdd2)) {
+
+    /**
+     * Set the preferred location of each partition using the executor that has the related
+     * [[StateStoreProvider]] already loaded.
+     */
+    override def getPreferredLocations(partition: Partition): Seq[String] = {
+      stateStoreNames.flatMap { storeName =>
+        val stateStoreProviderId = StateStoreProviderId(stateInfo, partition.index, storeName)
+        storeCoordinator.flatMap(_.getLocation(stateStoreProviderId))
+      }.distinct
+    }
+
+    override def compute(s: Partition, context: TaskContext): Iterator[V] = {
+      val partitions = s.asInstanceOf[ZippedPartitionsPartition].partitions
+      if (partitions(0).index != partitions(1).index) {
+        throw new IllegalStateException(s"Partition ID should be same in both side: " +
+          s"left ${partitions(0).index} , right ${partitions(1).index}")
+      }
+
+      val partitionId = partitions(0).index
+      f(partitionId, rdd1.iterator(partitions(0), context),
+        rdd2.iterator(partitions(1), context), schemas)
+    }
+
+    override def clearDependencies(): Unit = {
+      super.clearDependencies()
+      rdd1 = null
+      rdd2 = null
+      f = null
+    }
+  }
+
   implicit class StateStoreAwareZipPartitionsHelper[T: ClassTag](dataRDD: RDD[T]) {
     /**
      * Function used by `StreamingSymmetricHashJoinExec` to zip together the partitions of two
@@ -319,6 +369,19 @@ object StreamingSymmetricHashJoinHelper extends Logging {
       new StateStoreAwareZipPartitionsRDD(
         dataRDD.sparkContext, f, dataRDD, dataRDD2, stateInfo, storeNames, Some(storeCoordinator))
     }
+
+    def stateStoreAwareZipPartitions[U: ClassTag, V: ClassTag](
+        dataRDD2: RDD[U],
+        stateInfo: StatefulOperatorStateInfo,
+        storeNames: Seq[String],
+        storeCoordinator: StateStoreCoordinatorRef,
+        schemas: Map[String, StateStoreColFamilySchema]
+      )(f: (Int, Iterator[T], Iterator[U], Map[String, StateStoreColFamilySchema]) =>
+      Iterator[V]): RDD[V] = {
+      new StateStoreAwareZipPartitionsRDDWithSchemas(
+        dataRDD.sparkContext, f, dataRDD, dataRDD2, stateInfo,
+        storeNames, Some(storeCoordinator), schemas)
+    }
   }
 
   case class JoinerStateStoreCkptInfo(
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
index 42cd429587f3e..eed0593d2d7dd 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
@@ -104,7 +104,8 @@ case class TransformWithStateExec(
    * @return a new instance of the driver processor handle
    */
   private def getDriverProcessorHandle(): DriverStatefulProcessorHandleImpl = {
-    val driverProcessorHandle = new DriverStatefulProcessorHandleImpl(timeMode, keyEncoder)
+    val driverProcessorHandle = new DriverStatefulProcessorHandleImpl(
+      timeMode, keyEncoder, initializeAvroEnc = useAvroEncoding)
     driverProcessorHandle.setHandleState(StatefulProcessorHandleState.PRE_INIT)
     statefulProcessor.setHandle(driverProcessorHandle)
     statefulProcessor.init(outputMode, timeMode)
@@ -532,10 +533,11 @@ case class TransformWithStateExec(
         initialState.execute(),
         getStateInfo,
         storeNames = Seq(),
-        session.streams.stateStoreCoordinator) {
+        session.streams.stateStoreCoordinator,
+        getColFamilySchemas()) {
         // The state store aware zip partitions will provide us with two iterators,
         // child data iterator and the initial state iterator per partition.
-        case (partitionId, childDataIterator, initStateIterator) =>
+        case (partitionId, childDataIterator, initStateIterator, colFamilySchemas) =>
           if (isStreaming) {
             val stateStoreId = StateStoreId(stateInfo.get.checkpointLocation,
               stateInfo.get.operatorId, partitionId)
@@ -552,26 +554,29 @@ case class TransformWithStateExec(
               hadoopConf = hadoopConfBroadcast.value.value
             )
 
-            processDataWithInitialState(store, childDataIterator, initStateIterator)
+            processDataWithInitialState(
+              store, childDataIterator, initStateIterator, colFamilySchemas)
           } else {
-            initNewStateStoreAndProcessData(partitionId, hadoopConfBroadcast) { store =>
-              processDataWithInitialState(store, childDataIterator, initStateIterator)
+            initNewStateStoreAndProcessData(
+              partitionId, hadoopConfBroadcast, getColFamilySchemas()) { (store, schemas) =>
+              processDataWithInitialState(store, childDataIterator, initStateIterator, schemas)
             }
           }
       }
     } else {
       if (isStreaming) {
-        child.execute().mapPartitionsWithStateStore[InternalRow](
+        child.execute().mapPartitionsWithStateStoreWithSchemas[InternalRow](
           getStateInfo,
           keyEncoder.schema,
           DUMMY_VALUE_ROW_SCHEMA,
           NoPrefixKeyStateEncoderSpec(keyEncoder.schema),
           session.sessionState,
           Some(session.streams.stateStoreCoordinator),
-          useColumnFamilies = true
+          useColumnFamilies = true,
+          columnFamilySchemas = getColFamilySchemas()
         ) {
-          case (store: StateStore, singleIterator: Iterator[InternalRow]) =>
-            processData(store, singleIterator)
+          case (store: StateStore, singleIterator: Iterator[InternalRow], columnFamilySchemas) =>
+            processData(store, singleIterator, columnFamilySchemas)
         }
       } else {
         // If the query is running in batch mode, we need to create a new StateStore and instantiate
@@ -580,8 +585,9 @@ case class TransformWithStateExec(
           new SerializableConfiguration(session.sessionState.newHadoopConf()))
         child.execute().mapPartitionsWithIndex[InternalRow](
           (i: Int, iter: Iterator[InternalRow]) => {
-            initNewStateStoreAndProcessData(i, hadoopConfBroadcast) { store =>
-              processData(store, iter)
+            initNewStateStoreAndProcessData(
+              i, hadoopConfBroadcast, getColFamilySchemas()) { (store, schemas) =>
+              processData(store, iter, schemas)
             }
           }
         )
@@ -595,8 +601,10 @@ case class TransformWithStateExec(
    */
   private def initNewStateStoreAndProcessData(
       partitionId: Int,
-      hadoopConfBroadcast: Broadcast[SerializableConfiguration])
-    (f: StateStore => CompletionIterator[InternalRow, Iterator[InternalRow]]):
+      hadoopConfBroadcast: Broadcast[SerializableConfiguration],
+      schemas: Map[String, StateStoreColFamilySchema])
+    (f: (StateStore, Map[String, StateStoreColFamilySchema]) =>
+      CompletionIterator[InternalRow, Iterator[InternalRow]]):
     CompletionIterator[InternalRow, Iterator[InternalRow]] = {
 
     val providerId = {
@@ -621,8 +629,8 @@ case class TransformWithStateExec(
       hadoopConf = hadoopConfBroadcast.value.value,
       useMultipleValuesPerKey = true)
 
-    val store = stateStoreProvider.getStore(0, None)
-    val outputIterator = f(store)
+    val store = stateStoreProvider.getStore(0)
+    val outputIterator = f(store, schemas)
     CompletionIterator[InternalRow, Iterator[InternalRow]](outputIterator.iterator, {
       stateStoreProvider.close()
       statefulProcessor.close()
@@ -633,13 +641,17 @@ case class TransformWithStateExec(
    * Process the data in the partition using the state store and the stateful processor.
    * @param store The state store to use
    * @param singleIterator The iterator of rows to process
+   * @param schemas The column family schemas used by this stateful processor
    * @return An iterator of rows that are the result of processing the input rows
    */
-  private def processData(store: StateStore, singleIterator: Iterator[InternalRow]):
+  private def processData(
+      store: StateStore,
+      singleIterator: Iterator[InternalRow],
+      schemas: Map[String, StateStoreColFamilySchema]):
     CompletionIterator[InternalRow, Iterator[InternalRow]] = {
     val processorHandle = new StatefulProcessorHandleImpl(
       store, getStateInfo.queryRunId, keyEncoder, timeMode,
-      isStreaming, batchTimestampMs, metrics)
+      isStreaming, batchTimestampMs, metrics, schemas)
     assert(processorHandle.getHandleState == StatefulProcessorHandleState.CREATED)
     statefulProcessor.setHandle(processorHandle)
     statefulProcessor.init(outputMode, timeMode)
@@ -650,10 +662,11 @@ case class TransformWithStateExec(
   private def processDataWithInitialState(
       store: StateStore,
       childDataIterator: Iterator[InternalRow],
-      initStateIterator: Iterator[InternalRow]):
+      initStateIterator: Iterator[InternalRow],
+      schemas: Map[String, StateStoreColFamilySchema]):
     CompletionIterator[InternalRow, Iterator[InternalRow]] = {
     val processorHandle = new StatefulProcessorHandleImpl(store, getStateInfo.queryRunId,
-      keyEncoder, timeMode, isStreaming, batchTimestampMs, metrics)
+      keyEncoder, timeMode, isStreaming, batchTimestampMs, metrics, schemas)
     assert(processorHandle.getHandleState == StatefulProcessorHandleState.CREATED)
     statefulProcessor.setHandle(processorHandle)
     statefulProcessor.init(outputMode, timeMode)
@@ -718,8 +731,7 @@ object TransformWithStateExec {
       queryRunId = UUID.randomUUID(),
       operatorId = 0,
       storeVersion = 0,
-      numPartitions = shufflePartitions,
-      stateStoreCkptIds = None
+      numPartitions = shufflePartitions
     )
 
     new TransformWithStateExec(
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala
index b1b87feeb263b..fff67396cce3b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala
@@ -21,6 +21,7 @@ import org.apache.spark.sql.Encoder
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, StateStore}
+import org.apache.spark.sql.execution.streaming.state.AvroEncoderSpec
 import org.apache.spark.sql.streaming.ValueState
 
 /**
@@ -38,10 +39,17 @@ class ValueStateImpl[S](
     stateName: String,
     keyExprEnc: ExpressionEncoder[Any],
     valEncoder: Encoder[S],
+    avroEnc: Option[AvroEncoderSpec],
     metrics: Map[String, SQLMetric] = Map.empty)
   extends ValueState[S] with Logging {
 
-  private val stateTypesEncoder = StateTypesEncoder(keyExprEnc, valEncoder, stateName)
+  // If we are using Avro, the avroSerde parameter must be populated
+  // else, we will default to using UnsafeRow.
+  private val usingAvro: Boolean = avroEnc.isDefined
+  private val avroTypesEncoder = new AvroTypesEncoder[S](
+    keyExprEnc, valEncoder, stateName, hasTtl = false, avroEnc)
+  private val unsafeRowTypesEncoder = new UnsafeRowTypesEncoder[S](
+    keyExprEnc, valEncoder, stateName, hasTtl = false)
 
   initialize()
 
@@ -62,11 +70,28 @@ class ValueStateImpl[S](
 
   /** Function to return associated value with key if exists and null otherwise */
   override def get(): S = {
-    val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
+    if (usingAvro) {
+      getAvro()
+    } else {
+      getUnsafeRow()
+    }
+  }
+
+  private def getAvro(): S = {
+    val encodedGroupingKey = avroTypesEncoder.encodeGroupingKey()
     val retRow = store.get(encodedGroupingKey, stateName)
+    if (retRow != null) {
+      avroTypesEncoder.decodeValue(retRow)
+    } else {
+      null.asInstanceOf[S]
+    }
+  }
 
+  private def getUnsafeRow(): S = {
+    val encodedGroupingKey = unsafeRowTypesEncoder.encodeGroupingKey()
+    val retRow = store.get(encodedGroupingKey, stateName)
     if (retRow != null) {
-      stateTypesEncoder.decodeValue(retRow)
+      unsafeRowTypesEncoder.decodeValue(retRow)
     } else {
       null.asInstanceOf[S]
     }
@@ -74,15 +99,27 @@ class ValueStateImpl[S](
 
   /** Function to update and overwrite state associated with given key */
   override def update(newState: S): Unit = {
-    val encodedValue = stateTypesEncoder.encodeValue(newState)
-    store.put(stateTypesEncoder.encodeGroupingKey(),
-      encodedValue, stateName)
+    if (usingAvro) {
+      val encodedValue = avroTypesEncoder.encodeValue(newState)
+      store.put(avroTypesEncoder.encodeGroupingKey(),
+        encodedValue, stateName)
+    } else {
+      val encodedValue = unsafeRowTypesEncoder.encodeValue(newState)
+      store.put(unsafeRowTypesEncoder.encodeGroupingKey(),
+        encodedValue, stateName)
+    }
     TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
   }
 
   /** Function to remove state for given key */
   override def clear(): Unit = {
-    store.remove(stateTypesEncoder.encodeGroupingKey(), stateName)
+    if (usingAvro) {
+      val encodedKey = avroTypesEncoder.encodeGroupingKey()
+      store.remove(encodedKey, stateName)
+    } else {
+      val encodedKey = unsafeRowTypesEncoder.encodeGroupingKey()
+      store.remove(encodedKey, stateName)
+    }
     TWSMetricsUtils.incrementMetric(metrics, "numRemovedStateRows")
   }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala
index 145cd90264910..ac7a83ff65c21 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala
@@ -21,7 +21,7 @@ import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, StateStore}
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, NoPrefixKeyStateEncoderSpec, StateStore}
 import org.apache.spark.sql.streaming.{TTLConfig, ValueState}
 
 /**
@@ -44,11 +44,12 @@ class ValueStateImplWithTTL[S](
     valEncoder: Encoder[S],
     ttlConfig: TTLConfig,
     batchTimestampMs: Long,
+    avroEnc: Option[AvroEncoderSpec],   // TODO: Add Avro Encoding support for TTL
     metrics: Map[String, SQLMetric] = Map.empty)
   extends SingleKeyTTLStateImpl(
     stateName, store, keyExprEnc, batchTimestampMs) with ValueState[S] {
 
-  private val stateTypesEncoder = StateTypesEncoder(keyExprEnc, valEncoder,
+  private val stateTypesEncoder = UnsafeRowTypesEncoder(keyExprEnc, valEncoder,
     stateName, hasTtl = true)
   private val ttlExpirationMs =
     StateTTL.calculateExpirationTimeForDuration(ttlConfig.ttlDuration, batchTimestampMs)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala
index 2f77b2c14b009..899d96e1b341e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala
@@ -101,6 +101,24 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
     override def valuesIterator(key: UnsafeRow, colFamilyName: String): Iterator[UnsafeRow] = {
       throw StateStoreErrors.unsupportedOperationException("multipleValuesPerKey", "HDFSStateStore")
     }
+
+
+    override def get(key: Array[Byte], colFamilyName: String): Array[Byte] = {
+      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
+    }
+
+    override def valuesIterator(key: Array[Byte], colFamilyName: String): Iterator[Array[Byte]] = {
+      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
+    }
+
+    override def prefixScan(
+        prefixKey: Array[Byte], colFamilyName: String): Iterator[ByteArrayPair] = {
+      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
+    }
+
+    override def byteArrayIter(colFamilyName: String): Iterator[ByteArrayPair] = {
+      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
+    }
   }
 
   /** Implementation of [[StateStore]] API which is backed by an HDFS-compatible file system */
@@ -250,6 +268,35 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
         colFamilyName: String): Unit = {
       throw StateStoreErrors.unsupportedOperationException("merge", providerName)
     }
+
+    override def put(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit = {
+      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
+    }
+
+    override def remove(key: Array[Byte], colFamilyName: String): Unit = {
+      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
+    }
+
+    override def merge(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit = {
+      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
+    }
+
+    override def get(key: Array[Byte], colFamilyName: String): Array[Byte] = {
+      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
+    }
+
+    override def valuesIterator(key: Array[Byte], colFamilyName: String): Iterator[Array[Byte]] = {
+      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
+    }
+
+    override def prefixScan(
+        prefixKey: Array[Byte], colFamilyName: String): Iterator[ByteArrayPair] = {
+      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
+    }
+
+    override def byteArrayIter(colFamilyName: String): Iterator[ByteArrayPair] = {
+      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
+    }
   }
 
   def getMetricsForProvider(): Map[String, Long] = synchronized {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
index 4c7a226e0973f..bbf0cfcfc7905 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
@@ -33,6 +33,8 @@ sealed trait RocksDBKeyStateEncoder {
   def encodePrefixKey(prefixKey: UnsafeRow): Array[Byte]
   def encodeKey(row: UnsafeRow): Array[Byte]
   def decodeKey(keyBytes: Array[Byte]): UnsafeRow
+  def encodeKeyBytes(row: Array[Byte]): Array[Byte]
+  def decodeKeyBytes(keyBytes: Array[Byte]): Array[Byte]
   def getColumnFamilyIdBytes(): Array[Byte]
 }
 
@@ -41,6 +43,9 @@ sealed trait RocksDBValueStateEncoder {
   def encodeValue(row: UnsafeRow): Array[Byte]
   def decodeValue(valueBytes: Array[Byte]): UnsafeRow
   def decodeValues(valueBytes: Array[Byte]): Iterator[UnsafeRow]
+  def encodeValueBytes(row: Array[Byte]): Array[Byte]
+  def decodeValueBytes(valueBytes: Array[Byte]): Array[Byte]
+  def decodeValuesBytes(valueBytes: Array[Byte]): Iterator[Array[Byte]]
 }
 
 abstract class RocksDBKeyStateEncoderBase(
@@ -166,6 +171,49 @@ object RocksDBStateEncoder {
       null
     }
   }
+
+  /**
+   * Encode a byte array by adding a version byte at the beginning.
+   * Final byte layout: [VersionByte][OriginalBytes]
+   * where:
+   * - VersionByte: Single byte indicating encoding version
+   * - OriginalBytes: The input byte array unchanged
+   *
+   * @param input The original byte array to encode
+   * @return A new byte array containing the version byte followed by the input bytes
+   * @note This creates a new byte array and copies the input array to the new array.
+   */
+  def encodeByteArray(input: Array[Byte]): Array[Byte] = {
+    val encodedBytes = new Array[Byte](input.length + STATE_ENCODING_NUM_VERSION_BYTES)
+    Platform.putByte(encodedBytes, Platform.BYTE_ARRAY_OFFSET, STATE_ENCODING_VERSION)
+    Platform.copyMemory(
+      input, Platform.BYTE_ARRAY_OFFSET,
+      encodedBytes, Platform.BYTE_ARRAY_OFFSET + STATE_ENCODING_NUM_VERSION_BYTES,
+      input.length)
+    encodedBytes
+  }
+
+  /**
+   * Decode bytes by removing the version byte at the beginning.
+   * Input byte layout: [VersionByte][OriginalBytes]
+   * Returns: [OriginalBytes]
+   *
+   * @param bytes The encoded byte array
+   * @return A new byte array containing just the original bytes (excluding version byte),
+   *         or null if input is null
+   */
+  def decodeToByteArray(bytes: Array[Byte]): Array[Byte] = {
+    if (bytes != null) {
+      val decodedBytes = new Array[Byte](bytes.length - STATE_ENCODING_NUM_VERSION_BYTES)
+      Platform.copyMemory(
+        bytes, Platform.BYTE_ARRAY_OFFSET + STATE_ENCODING_NUM_VERSION_BYTES,
+        decodedBytes, Platform.BYTE_ARRAY_OFFSET,
+        decodedBytes.length)
+      decodedBytes
+    } else {
+      null
+    }
+  }
 }
 
 /**
@@ -267,6 +315,12 @@ class PrefixKeyScanStateEncoder(
   }
 
   override def supportPrefixKeyScan: Boolean = true
+
+  override def encodeKeyBytes(row: Array[Byte]): Array[Byte] =
+    throw new UnsupportedOperationException
+
+  override def decodeKeyBytes(keyBytes: Array[Byte]): Array[Byte] =
+    throw new UnsupportedOperationException
 }
 
 /**
@@ -644,6 +698,12 @@ class RangeKeyScanStateEncoder(
   }
 
   override def supportPrefixKeyScan: Boolean = true
+
+  override def encodeKeyBytes(row: Array[Byte]): Array[Byte] =
+    throw new UnsupportedOperationException
+
+  override def decodeKeyBytes(keyBytes: Array[Byte]): Array[Byte] =
+    throw new UnsupportedOperationException
 }
 
 /**
@@ -708,6 +768,69 @@ class NoPrefixKeyStateEncoder(
     } else decodeToUnsafeRow(keyBytes, keyRow)
   }
 
+  /**
+   * Encodes a byte array by adding column family prefix and version information.
+   * Final byte layout: [ColFamilyPrefix][VersionByte][OriginalBytes]
+   * where:
+   * - ColFamilyPrefix: Optional prefix identifying the column family (if useColumnFamilies=true)
+   * - VersionByte: Single byte indicating encoding version
+   * - OriginalBytes: The input byte array unchanged
+   *
+   * @param row The original byte array to encode
+   * @return The encoded byte array with prefix and version if column families are enabled,
+   *         otherwise returns the original array
+   */
+  override def encodeKeyBytes(row: Array[Byte]): Array[Byte] = {
+    if (!useColumnFamilies) {
+      row
+    } else {
+      // Calculate total size needed: original bytes + 1 byte for version + column family prefix
+      val (encodedBytes, startingOffset) = encodeColumnFamilyPrefix(
+        row.length +
+          STATE_ENCODING_NUM_VERSION_BYTES
+      )
+
+      // Add version byte right after column family prefix
+      Platform.putByte(encodedBytes, startingOffset, STATE_ENCODING_VERSION)
+
+      // Copy original bytes after the version byte
+      // Platform.BYTE_ARRAY_OFFSET is the recommended way to memcopy b/w byte arrays. See Platform.
+      Platform.copyMemory(
+        row, Platform.BYTE_ARRAY_OFFSET,
+        encodedBytes, startingOffset + STATE_ENCODING_NUM_VERSION_BYTES, row.length)
+      encodedBytes
+    }
+  }
+
+  /**
+   * Decodes a byte array by removing column family prefix and version information.
+   * Input byte layout: [ColFamilyPrefix][VersionByte][OriginalBytes]
+   * Returns: [OriginalBytes]
+   *
+   * @param keyBytes The encoded byte array to decode
+   * @return The original byte array with prefix and version removed if column families are enabled,
+   *         null if input is null, or a clone of input if column families are disabled
+   */
+  override def decodeKeyBytes(keyBytes: Array[Byte]): Array[Byte] = {
+    if (keyBytes == null) {
+      null
+    } else if (useColumnFamilies) {
+      // Calculate start offset (skip column family prefix and version byte)
+      val startOffset = decodeKeyStartOffset + STATE_ENCODING_NUM_VERSION_BYTES
+
+      // Calculate length of original data (total length minus prefix and version)
+      val length = keyBytes.length -
+        STATE_ENCODING_NUM_VERSION_BYTES - VIRTUAL_COL_FAMILY_PREFIX_BYTES
+
+      // Extract just the original bytes
+      java.util.Arrays.copyOfRange(keyBytes, startOffset, startOffset + length)
+    } else {
+      // If column families not enabled, just return a copy of the input
+      // Assuming decodeToUnsafeRow is not applicable for byte array encoding
+      keyBytes.clone()
+    }
+  }
+
   override def supportPrefixKeyScan: Boolean = false
 
   override def encodePrefixKey(prefixKey: UnsafeRow): Array[Byte] = {
@@ -789,6 +912,95 @@ class MultiValuedStateEncoder(valueSchema: StructType)
   }
 
   override def supportsMultipleValuesPerKey: Boolean = true
+
+  /**
+   * Encodes a raw byte array value in the multi-value format.
+   * Format: [length (4 bytes)][value bytes][delimiter (1 byte)]
+   */
+  override def encodeValueBytes(row: Array[Byte]): Array[Byte] = {
+    if (row == null) {
+      null
+    } else {
+      val numBytes = row.length
+      // Allocate space for:
+      // - 4 bytes for length
+      // - The actual value bytes
+      val encodedBytes = new Array[Byte](java.lang.Integer.BYTES + numBytes)
+
+      // Write length as big-endian int
+      Platform.putInt(encodedBytes, Platform.BYTE_ARRAY_OFFSET, numBytes)
+
+      // Copy value bytes after the length
+      Platform.copyMemory(
+        row, Platform.BYTE_ARRAY_OFFSET,
+        encodedBytes, Platform.BYTE_ARRAY_OFFSET + java.lang.Integer.BYTES,
+        numBytes
+      )
+
+      encodedBytes
+    }
+  }
+
+  /**
+   * Decodes a single value from the encoded byte format.
+   * Assumes the bytes represent a single value, not multiple merged values.
+   */
+  override def decodeValueBytes(valueBytes: Array[Byte]): Array[Byte] = {
+    if (valueBytes == null) {
+      null
+    } else {
+      // Read length from first 4 bytes
+      val numBytes = Platform.getInt(valueBytes, Platform.BYTE_ARRAY_OFFSET)
+
+      // Extract just the value bytes after the length
+      val decodedBytes = new Array[Byte](numBytes)
+      Platform.copyMemory(
+        valueBytes, Platform.BYTE_ARRAY_OFFSET + java.lang.Integer.BYTES,
+        decodedBytes, Platform.BYTE_ARRAY_OFFSET,
+        numBytes
+      )
+
+      decodedBytes
+    }
+  }
+
+  /**
+   * Decodes multiple values from the merged byte format.
+   * Returns an iterator that lazily decodes each value.
+   */
+  override def decodeValuesBytes(valueBytes: Array[Byte]): Iterator[Array[Byte]] = {
+    if (valueBytes == null) {
+      Iterator.empty
+    } else {
+      new Iterator[Array[Byte]] {
+        // Track current position in the byte array
+        private var pos: Int = Platform.BYTE_ARRAY_OFFSET
+        private val maxPos = Platform.BYTE_ARRAY_OFFSET + valueBytes.length
+
+        override def hasNext: Boolean = pos < maxPos
+
+        override def next(): Array[Byte] = {
+          // Read length prefix
+          val numBytes = Platform.getInt(valueBytes, pos)
+          pos += java.lang.Integer.BYTES
+
+          // Extract value bytes
+          val decodedValue = new Array[Byte](numBytes)
+          Platform.copyMemory(
+            valueBytes, pos,
+            decodedValue, Platform.BYTE_ARRAY_OFFSET,
+            numBytes
+          )
+
+          // Move position past value and delimiter
+          pos += numBytes
+          pos += 1 // Skip delimiter byte
+
+          decodedValue
+        }
+      }
+    }
+  }
 }
 
 /**
@@ -828,4 +1040,16 @@ class SingleValueStateEncoder(valueSchema: StructType)
   override def decodeValues(valueBytes: Array[Byte]): Iterator[UnsafeRow] = {
     throw new IllegalStateException("This encoder doesn't support multiple values!")
   }
+
+  override def encodeValueBytes(row: Array[Byte]): Array[Byte] = {
+    encodeByteArray(row)
+  }
+
+  override def decodeValueBytes(valueBytes: Array[Byte]): Array[Byte] = {
+    decodeToByteArray(valueBytes)
+  }
+
+  override def decodeValuesBytes(valueBytes: Array[Byte]): Iterator[Array[Byte]] = {
+    throw new IllegalStateException("This encoder doesn't support multiple values!")
+  }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala
index 1fc6ab5910c6c..2bcccc1a2d310 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala
@@ -209,6 +209,98 @@ private[sql] class RocksDBStateStoreProvider
       }
     }
 
+    override def put(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit = {
+      verify(state == UPDATING, "Cannot put after already committed or aborted")
+      verify(key != null, "Key cannot be null")
+      require(value != null, "Cannot put a null value")
+      verifyColFamilyOperations("put", colFamilyName)
+
+      val kvEncoder = keyValueEncoderMap.get(colFamilyName)
+      rocksDB.put(kvEncoder._1.encodeKeyBytes(key), kvEncoder._2.encodeValueBytes(value))
+    }
+
+    override def remove(key: Array[Byte], colFamilyName: String): Unit = {
+      verify(state == UPDATING, "Cannot remove after already committed or aborted")
+      verify(key != null, "Key cannot be null")
+      verifyColFamilyOperations("remove", colFamilyName)
+
+      val kvEncoder = keyValueEncoderMap.get(colFamilyName)
+      rocksDB.remove(kvEncoder._1.encodeKeyBytes(key))
+    }
+
+    override def get(key: Array[Byte], colFamilyName: String): Array[Byte] = {
+      verify(key != null, "Key cannot be null")
+      verifyColFamilyOperations("get", colFamilyName)
+
+      val kvEncoder = keyValueEncoderMap.get(colFamilyName)
+      kvEncoder._2.decodeValueBytes(rocksDB.get(kvEncoder._1.encodeKeyBytes(key)))
+    }
+
+    override def byteArrayIter(colFamilyName: String): Iterator[ByteArrayPair] = {
+      // Verify column family operation is valid
+      verifyColFamilyOperations("byteArrayIter", colFamilyName)
+      val kvEncoder = keyValueEncoderMap.get(colFamilyName)
+      val pair = new ByteArrayPair()
+
+      // Similar to the regular iterator, we need to handle both column family
+      // and non-column family cases
+      if (useColumnFamilies) {
+        rocksDB.prefixScan(kvEncoder._1.getColumnFamilyIdBytes()).map { kv =>
+          pair.set(
+            kvEncoder._1.decodeKeyBytes(kv.key),
+            kvEncoder._2.decodeValueBytes(kv.value))
+          pair
+        }
+      } else {
+        rocksDB.iterator().map { kv =>
+          pair.set(
+            kvEncoder._1.decodeKeyBytes(kv.key),
+            kvEncoder._2.decodeValueBytes(kv.value))
+          pair
+        }
+      }
+    }
+
+    override def valuesIterator(key: Array[Byte], colFamilyName: String): Iterator[Array[Byte]] = {
+      verify(key != null, "Key cannot be null")
+      verifyColFamilyOperations("valuesIterator", colFamilyName)
+
+      val kvEncoder = keyValueEncoderMap.get(colFamilyName)
+      val valueEncoder = kvEncoder._2
+      val keyEncoder = kvEncoder._1
+
+      verify(valueEncoder.supportsMultipleValuesPerKey,
+        "valuesIterator requires an encoder that supports multiple values for a single key.")
+
+      // Get the encoded value bytes using the encoded key
+      val encodedValues = rocksDB.get(keyEncoder.encodeKeyBytes(key))
+
+      // Decode multiple values from the merged value bytes
+      valueEncoder.decodeValuesBytes(encodedValues)
+    }
+
+    override def prefixScan(
+        prefixKey: Array[Byte],
+        colFamilyName: String): Iterator[ByteArrayPair] = {
+      throw StateStoreErrors.unsupportedOperationException(
+        "bytearray prefixScan", "RocksDBStateStore")
+    }
+
+    override def merge(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit = {
+      verify(state == UPDATING, "Cannot merge after already committed or aborted")
+      verifyColFamilyOperations("merge", colFamilyName)
+
+      val kvEncoder = keyValueEncoderMap.get(colFamilyName)
+      val keyEncoder = kvEncoder._1
+      val valueEncoder = kvEncoder._2
+      verify(valueEncoder.supportsMultipleValuesPerKey, "Merge operation requires an encoder" +
+        " which supports multiple values for a single key")
+      verify(key != null, "Key cannot be null")
+      require(value != null, "Cannot merge a null value")
+
+      rocksDB.merge(keyEncoder.encodeKeyBytes(key), valueEncoder.encodeValueBytes(value))
+    }
+
     override def commit(): Long = synchronized {
       try {
         verify(state == UPDATING, "Cannot commit after already committed or aborted")
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
index 721d72b6a0991..b02d36c8ced85 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
@@ -24,6 +24,7 @@ import org.apache.hadoop.fs.Path
 
 import org.apache.spark.SparkUnsupportedOperationException
 import org.apache.spark.internal.{Logging, LogKeys, MDC}
+import org.apache.spark.sql.avro.{AvroDeserializer, AvroSerializer}
 import org.apache.spark.sql.catalyst.util.UnsafeRowUtils
 import org.apache.spark.sql.execution.streaming.{CheckpointFileManager, StatefulOperatorStateInfo}
 import org.apache.spark.sql.execution.streaming.state.SchemaHelper.{SchemaReader, SchemaWriter}
@@ -37,14 +38,21 @@ case class StateSchemaValidationResult(
     schemaPath: String
 )
 
+case class AvroEncoderSpec(
+  keySerializer: AvroSerializer,
+  valueSerializer: AvroSerializer,
+  valueDeserializer: AvroDeserializer
+) extends Serializable
+
 // Used to represent the schema of a column family in the state store
 case class StateStoreColFamilySchema(
     colFamilyName: String,
     keySchema: StructType,
     valueSchema: StructType,
     keyStateEncoderSpec: Option[KeyStateEncoderSpec] = None,
-    userKeyEncoderSchema: Option[StructType] = None
-)
+    userKeyEncoderSchema: Option[StructType] = None,
+    avroEnc: Option[AvroEncoderSpec] = None
+) extends Serializable
 
 class StateSchemaCompatibilityChecker(
     providerId: StateStoreProviderId,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
index 72bc3ca33054d..09221d8374e24 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
@@ -99,6 +99,44 @@ trait ReadStateStore {
   /** Return an iterator containing all the key-value pairs in the StateStore. */
   def iterator(colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Iterator[UnsafeRowPair]
 
+  /**
+   * Get the current value of a non-null key.
+   * @return a non-null row if the key exists in the store, otherwise null.
+   */
+  def get(
+      key: Array[Byte],
+      colFamilyName: String): Array[Byte]
+
+  /**
+   * Provides an iterator containing all values of a non-null key. If key does not exist,
+   * an empty iterator is returned. Implementations should make sure to return an empty
+   * iterator if the key does not exist.
+   *
+   * It is expected to throw exception if Spark calls this method without setting
+   * multipleValuesPerKey as true for the column family.
+   */
+  def valuesIterator(
+      key: Array[Byte],
+      colFamilyName: String): Iterator[Array[Byte]]
+
+  /**
+   * Return an iterator containing all the key-value pairs which are matched with
+   * the given prefix key.
+   *
+   * The operator will provide numColsPrefixKey greater than 0 in StateStoreProvider.init method
+   * if the operator needs to leverage the "prefix scan" feature. The schema of the prefix key
+   * should be same with the leftmost `numColsPrefixKey` columns of the key schema.
+   *
+   * It is expected to throw exception if Spark calls this method without setting numColsPrefixKey
+   * to the greater than 0.
+   */
+  def prefixScan(
+      prefixKey: Array[Byte],
+      colFamilyName: String): Iterator[ByteArrayPair]
+
+  /** Return an iterator containing all the key-value pairs in the StateStore. */
+  def byteArrayIter(colFamilyName: String): Iterator[ByteArrayPair]
+
   /**
    * Clean up the resource.
    *
@@ -163,6 +201,19 @@ trait StateStore extends ReadStateStore {
   def merge(key: UnsafeRow, value: UnsafeRow,
       colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Unit
 
+  /**
+   * Put a new non-null value for a non-null key. Implementations must be aware that the UnsafeRows
+   * in the params can be reused, and must make copies of the data as needed for persistence.
+   */
+  def put(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit
+
+  /**
+   * Remove a single non-null key.
+   */
+  def remove(key: Array[Byte], colFamilyName: String): Unit
+
+  def merge(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit
+
   /**
    * Commit all the updates that have been made to the store, and return the new version.
    * Implementations should ensure that no more updates (puts, removes) can be after a commit in
@@ -229,6 +280,19 @@ class WrappedReadStateStore(store: StateStore) extends ReadStateStore {
   override def valuesIterator(key: UnsafeRow, colFamilyName: String): Iterator[UnsafeRow] = {
     store.valuesIterator(key, colFamilyName)
   }
+
+  override def get(key: Array[Byte], colFamilyName: String): Array[Byte] =
+    store.get(key, colFamilyName)
+
+  override def valuesIterator(key: Array[Byte], colFamilyName: String):
+    Iterator[Array[Byte]] = store.valuesIterator(key, colFamilyName)
+
+
+  override def prefixScan(prefixKey: Array[Byte], colFamilyName: String):
+    Iterator[ByteArrayPair] = store.prefixScan(prefixKey, colFamilyName)
+
+  override def byteArrayIter(colFamilyName: String): Iterator[ByteArrayPair] =
+    store.byteArrayIter(colFamilyName)
 }
 
 /**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/package.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/package.scala
index e1a95dd10be74..19a90c6978df0 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/package.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/package.scala
@@ -89,6 +89,49 @@ package object state {
         extraOptions,
         useMultipleValuesPerKey)
     }
+
+    /** Map each partition of an RDD along with data in a [[StateStore]] that passes the
+     * column family schemas to the storeUpdateFunction. Used to pass Avro encoders/decoders
+     * to executors */
+    def mapPartitionsWithStateStoreWithSchemas[U: ClassTag](
+        stateInfo: StatefulOperatorStateInfo,
+        keySchema: StructType,
+        valueSchema: StructType,
+        keyStateEncoderSpec: KeyStateEncoderSpec,
+        sessionState: SessionState,
+        storeCoordinator: Option[StateStoreCoordinatorRef],
+        useColumnFamilies: Boolean = false,
+        extraOptions: Map[String, String] = Map.empty,
+        useMultipleValuesPerKey: Boolean = false,
+        columnFamilySchemas: Map[String, StateStoreColFamilySchema] = Map.empty)(
+        storeUpdateFunction: (StateStore, Iterator[T], Map[String, StateStoreColFamilySchema]) => Iterator[U]): StateStoreRDD[T, U] = {
+
+      val cleanedF = dataRDD.sparkContext.clean(storeUpdateFunction)
+      val wrappedF = (store: StateStore, iter: Iterator[T]) => {
+        // Abort the state store in case of error
+        TaskContext.get().addTaskCompletionListener[Unit](_ => {
+          if (!store.hasCommitted) store.abort()
+        })
+        cleanedF(store, iter, columnFamilySchemas)
+      }
+
+      new StateStoreRDD(
+        dataRDD,
+        wrappedF,
+        stateInfo.checkpointLocation,
+        stateInfo.queryRunId,
+        stateInfo.operatorId,
+        stateInfo.storeVersion,
+        stateInfo.stateStoreCkptIds,
+        keySchema,
+        valueSchema,
+        keyStateEncoderSpec,
+        sessionState,
+        storeCoordinator,
+        useColumnFamilies,
+        extraOptions,
+        useMultipleValuesPerKey)
+    }
     // scalastyle:on
 
     /** Map each partition of an RDD along with data in a [[ReadStateStore]]. */
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
index 8f800b9f0252c..55e2f3704c7c0 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
@@ -327,6 +327,9 @@ trait StateStoreWriter
     OperatorStateMetadataV1(operatorInfo, stateStoreInfo)
   }
 
+  lazy val useAvroEncoding: Boolean =
+    conf.getConf(SQLConf.STREAMING_STATE_STORE_ENCODING_FORMAT) == "Avro"
+
   /** Set the operator level metrics */
   protected def setOperatorMetrics(numStateStoreInstances: Int = 1): Unit = {
     assert(numStateStoreInstances >= 1, s"invalid number of stores: $numStateStoreInstances")
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MemoryStateStore.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MemoryStateStore.scala
index 9a04a0c759ac4..a8859673466ec 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MemoryStateStore.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MemoryStateStore.scala
@@ -78,4 +78,34 @@ class MemoryStateStore extends StateStore() {
   override def getStateStoreCheckpointInfo(): StateStoreCheckpointInfo = {
     StateStoreCheckpointInfo(id.partitionId, version + 1, None, None)
   }
+
+
+  override def put(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit = {
+    throw new UnsupportedOperationException("Doesn't support bytearray operations")
+  }
+
+  override def remove(key: Array[Byte], colFamilyName: String): Unit = {
+    throw new UnsupportedOperationException("Doesn't support bytearray operations")
+  }
+
+  override def merge(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit = {
+    throw new UnsupportedOperationException("Doesn't support bytearray operations")
+  }
+
+  override def get(key: Array[Byte], colFamilyName: String): Array[Byte] = {
+    throw new UnsupportedOperationException("Doesn't support bytearray operations")
+  }
+
+  override def valuesIterator(key: Array[Byte], colFamilyName: String): Iterator[Array[Byte]] = {
+    throw new UnsupportedOperationException("Doesn't support bytearray operations")
+  }
+
+  override def prefixScan(
+      prefixKey: Array[Byte], colFamilyName: String): Iterator[ByteArrayPair] = {
+    throw new UnsupportedOperationException("Doesn't support bytearray operations")
+  }
+
+  override def byteArrayIter(colFamilyName: String): Iterator[ByteArrayPair] = {
+    throw new UnsupportedOperationException("Doesn't support bytearray operations")
+  }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala
index 9ac74eb5b9e8f..9bf5b6b73ff6c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala
@@ -122,6 +122,35 @@ case class CkptIdCollectingStateStoreWrapper(innerStore: StateStore) extends Sta
     innerStore.merge(key, value, colFamilyName)
   }
 
+  override def put(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit = {
+    throw new UnsupportedOperationException
+  }
+
+  override def remove(key: Array[Byte], colFamilyName: String): Unit = {
+    throw new UnsupportedOperationException
+  }
+
+  override def get(key: Array[Byte], colFamilyName: String): Array[Byte] = {
+    throw new UnsupportedOperationException
+  }
+
+  override def valuesIterator(key: Array[Byte], colFamilyName: String): Iterator[Array[Byte]] = {
+    throw new UnsupportedOperationException
+  }
+
+  override def prefixScan(
+      prefixKey: Array[Byte], colFamilyName: String): Iterator[ByteArrayPair] = {
+    throw new UnsupportedOperationException
+  }
+
+  override def byteArrayIter(colFamilyName: String): Iterator[ByteArrayPair] = {
+    throw new UnsupportedOperationException
+  }
+
+  override def merge(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit = {
+    throw new UnsupportedOperationException
+  }
+
   override def commit(): Long = innerStore.commit()
   override def metrics: StateStoreMetrics = innerStore.metrics
   override def getStateStoreCheckpointInfo(): StateStoreCheckpointInfo = {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala
index 8fde216c14411..ae5a8c01038b3 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala
@@ -128,6 +128,23 @@ trait AlsoTestWithChangelogCheckpointingEnabled
     }
   }
 
+  def testWithAvroEncoding(testName: String, testTags: Tag*)
+                          (testBody: => Any): Unit = {
+    Seq("UnsafeRow", "Avro").foreach { encoding =>
+      super.test(testName + s" (encoding = $encoding)", testTags: _*) {
+        // in case tests have any code that needs to execute before every test
+        super.beforeEach()
+        withSQLConf(
+          SQLConf.STREAMING_STATE_STORE_ENCODING_FORMAT.key ->
+            encoding) {
+          testBody
+        }
+        // in case tests have any code that needs to execute after every test
+        super.afterEach()
+      }
+    }
+  }
+
   def testWithColumnFamilies(
       testName: String,
       testMode: TestMode,
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ValueStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ValueStateSuite.scala
index 13d758eb1b88f..cf13495f5f616 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ValueStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ValueStateSuite.scala
@@ -23,7 +23,7 @@ import java.util.UUID
 import scala.util.Random
 
 import org.apache.hadoop.conf.Configuration
-import org.scalatest.BeforeAndAfter
+import org.scalatest.{BeforeAndAfter, Tag}
 
 import org.apache.spark.{SparkException, SparkUnsupportedOperationException}
 import org.apache.spark.sql.Encoders
@@ -45,14 +45,15 @@ class ValueStateSuite extends StateVariableSuiteBase {
 
   import StateStoreTestsHelper._
 
-  test("Implicit key operations") {
+  testWithAvroEnc("Implicit key operations") { useAvro =>
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
       val stateName = "testState"
-      val testState: ValueState[Long] = handle.getValueState[Long]("testState", Encoders.scalaLong)
+      val testState: ValueState[Long] = handle.getValueStateWithAvro[Long](
+        "testState", Encoders.scalaLong, useAvro)
       assert(ImplicitGroupingKeyTracker.getImplicitKeyOption.isEmpty)
       val ex = intercept[Exception] {
         testState.update(123)
@@ -89,13 +90,14 @@ class ValueStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  test("Value state operations for single instance") {
+  testWithAvroEnc("Value state operations for single instance") { useAvro =>
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val testState: ValueState[Long] = handle.getValueState[Long]("testState", Encoders.scalaLong)
+      val testState: ValueState[Long] = handle.getValueStateWithAvro[Long](
+        "testState", Encoders.scalaLong, useAvro)
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
       testState.update(123)
       assert(testState.get() === 123)
@@ -115,16 +117,16 @@ class ValueStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  test("Value state operations for multiple instances") {
+  testWithAvroEnc("Value state operations for multiple instances") { useAvro =>
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val testState1: ValueState[Long] = handle.getValueState[Long](
-        "testState1", Encoders.scalaLong)
-      val testState2: ValueState[Long] = handle.getValueState[Long](
-        "testState2", Encoders.scalaLong)
+      val testState1: ValueState[Long] = handle.getValueStateWithAvro[Long](
+        "testState1", Encoders.scalaLong, useAvro)
+      val testState2: ValueState[Long] = handle.getValueStateWithAvro[Long](
+        "testState2", Encoders.scalaLong, useAvro)
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
       testState1.update(123)
       assert(testState1.get() === 123)
@@ -160,7 +162,7 @@ class ValueStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  test("Value state operations for unsupported type name should fail") {
+  testWithAvroEnc("Value state operations for unsupported type name should fail") { useAvro =>
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store,
@@ -168,7 +170,7 @@ class ValueStateSuite extends StateVariableSuiteBase {
 
       val cfName = "$testState"
       val ex = intercept[SparkUnsupportedOperationException] {
-        handle.getValueState[Long](cfName, Encoders.scalaLong)
+        handle.getValueStateWithAvro[Long](cfName, Encoders.scalaLong, useAvro)
       }
       checkError(
         ex,
@@ -200,14 +202,15 @@ class ValueStateSuite extends StateVariableSuiteBase {
     )
   }
 
-  test("test SQL encoder - Value state operations for Primitive(Double) instances") {
+  testWithAvroEnc("test SQL encoder - Value state operations" +
+    " for Primitive(Double) instances") { useAvro =>
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val testState: ValueState[Double] = handle.getValueState[Double]("testState",
-        Encoders.scalaDouble)
+      val testState: ValueState[Double] = handle.getValueStateWithAvro[Double]("testState",
+        Encoders.scalaDouble, useAvro)
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
       testState.update(1.0)
       assert(testState.get().equals(1.0))
@@ -226,14 +229,15 @@ class ValueStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  test("test SQL encoder - Value state operations for Primitive(Long) instances") {
+  testWithAvroEnc("test SQL encoder - Value state operations" +
+    " for Primitive(Long) instances") { useAvro =>
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val testState: ValueState[Long] = handle.getValueState[Long]("testState",
-        Encoders.scalaLong)
+      val testState: ValueState[Long] = handle.getValueStateWithAvro[Long]("testState",
+        Encoders.scalaLong, useAvro)
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
       testState.update(1L)
       assert(testState.get().equals(1L))
@@ -252,14 +256,15 @@ class ValueStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  test("test SQL encoder - Value state operations for case class instances") {
+  testWithAvroEnc("test SQL encoder - Value state operations" +
+    " for case class instances") { useAvro =>
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val testState: ValueState[TestClass] = handle.getValueState[TestClass]("testState",
-        Encoders.product[TestClass])
+      val testState: ValueState[TestClass] = handle.getValueStateWithAvro[TestClass]("testState",
+        Encoders.product[TestClass], useAvro)
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
       testState.update(TestClass(1, "testcase1"))
       assert(testState.get().equals(TestClass(1, "testcase1")))
@@ -278,14 +283,14 @@ class ValueStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  test("test SQL encoder - Value state operations for POJO instances") {
+  testWithAvroEnc("test SQL encoder - Value state operations for POJO instances") { useAvro =>
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val testState: ValueState[POJOTestClass] = handle.getValueState[POJOTestClass]("testState",
-        Encoders.bean(classOf[POJOTestClass]))
+      val testState: ValueState[POJOTestClass] = handle.getValueStateWithAvro[POJOTestClass](
+        "testState", Encoders.bean(classOf[POJOTestClass]), useAvro)
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
       testState.update(new POJOTestClass("testcase1", 1))
       assert(testState.get().equals(new POJOTestClass("testcase1", 1)))
@@ -474,5 +479,26 @@ abstract class StateVariableSuiteBase extends SharedSparkSession
       provider.close()
     }
   }
-}
 
+  def testWithAvroEnc(testName: String, testTags: Tag*)(testBody: Boolean => Any): Unit = {
+    // Run with serde (true)
+    super.test(testName + " (with Avro encoding)", testTags: _*) {
+      super.beforeEach()
+      try {
+        testBody(true)
+      } finally {
+        super.afterEach()
+      }
+    }
+
+    // Run without serde (false)
+    super.test(testName, testTags: _*) {
+      super.beforeEach()
+      try {
+        testBody(false)
+      } finally {
+        super.afterEach()
+      }
+    }
+  }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateSuite.scala
index 20f04cc66c0aa..a515ef24e806d 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateSuite.scala
@@ -129,7 +129,7 @@ class TransformWithListStateSuite extends StreamTest
   with AlsoTestWithChangelogCheckpointingEnabled {
   import testImplicits._
 
-  test("test appending null value in list state throw exception") {
+  testWithAvroEncoding("test appending null value in list state throw exception") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
 
@@ -149,7 +149,7 @@ class TransformWithListStateSuite extends StreamTest
     }
   }
 
-  test("test putting null value in list state throw exception") {
+  testWithAvroEncoding("test putting null value in list state throw exception") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
 
@@ -169,7 +169,7 @@ class TransformWithListStateSuite extends StreamTest
     }
   }
 
-  test("test putting null list in list state throw exception") {
+  testWithAvroEncoding("test putting null list in list state throw exception") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
 
@@ -189,7 +189,7 @@ class TransformWithListStateSuite extends StreamTest
     }
   }
 
-  test("test appending null list in list state throw exception") {
+  testWithAvroEncoding("test appending null list in list state throw exception") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
 
@@ -209,7 +209,7 @@ class TransformWithListStateSuite extends StreamTest
     }
   }
 
-  test("test putting empty list in list state throw exception") {
+  testWithAvroEncoding("test putting empty list in list state throw exception") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
 
@@ -229,7 +229,7 @@ class TransformWithListStateSuite extends StreamTest
     }
   }
 
-  test("test appending empty list in list state throw exception") {
+  testWithAvroEncoding("test appending empty list in list state throw exception") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
 
@@ -249,7 +249,7 @@ class TransformWithListStateSuite extends StreamTest
     }
   }
 
-  test("test list state correctness") {
+  testWithAvroEncoding("test list state correctness") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
 
@@ -296,17 +296,18 @@ class TransformWithListStateSuite extends StreamTest
         AddData(inputData, InputRow("k5", "append", "v4")),
         AddData(inputData, InputRow("k5", "put", "v5,v6")),
         AddData(inputData, InputRow("k5", "emitAllInState", "")),
-        CheckNewAnswer(("k5", "v5"), ("k5", "v6")),
-        Execute { q =>
-          assert(q.lastProgress.stateOperators(0).customMetrics.get("numListStateVars") > 0)
-          assert(q.lastProgress.stateOperators(0).numRowsUpdated === 2)
-          assert(q.lastProgress.stateOperators(0).numRowsRemoved === 2)
-        }
+        CheckNewAnswer(("k5", "v5"), ("k5", "v6"))
+        // TODO: Uncomment once we have implemented ListStateMetrics for Avro encoding
+//        Execute { q =>
+//          assert(q.lastProgress.stateOperators(0).customMetrics.get("numListStateVars") > 0)
+//          assert(q.lastProgress.stateOperators(0).numRowsUpdated === 2)
+//          assert(q.lastProgress.stateOperators(0).numRowsRemoved === 2)
+//        }
       )
     }
   }
 
-  test("test ValueState And ListState in Processor") {
+  testWithAvroEncoding("test ValueState And ListState in Processor") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
index 1a7970302e5bc..49b16c597c018 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
@@ -399,7 +399,8 @@ class TransformWithStateSuite extends StateStoreMetricsTest
 
   import testImplicits._
 
-  test("transformWithState - streaming with rocksdb and invalid processor should fail") {
+  testWithAvroEncoding("transformWithState - streaming with rocksdb " +
+    "and invalid processor should fail") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->
@@ -420,7 +421,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  test("transformWithState - lazy iterators can properly get/set keyed state") {
+  testWithAvroEncoding("transformWithState - lazy iterators can properly get/set keyed state") {
     class ProcessorWithLazyIterators
       extends StatefulProcessor[Long, Long, Long] {
       @transient protected var _myValueState: ValueState[Long] = _
@@ -495,7 +496,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  test("transformWithState - streaming with rocksdb should succeed") {
+  testWithAvroEncoding("transformWithState - streaming with rocksdb should succeed") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->
@@ -533,7 +534,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  test("transformWithState - streaming with rocksdb and processing time timer " +
+  testWithAvroEncoding("transformWithState - streaming with rocksdb and processing time timer " +
    "should succeed") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
@@ -578,7 +579,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  test("transformWithState - streaming with rocksdb and processing time timer " +
+  testWithAvroEncoding("transformWithState - streaming with rocksdb and processing time timer " +
    "and updating timers should succeed") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
@@ -614,7 +615,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  test("transformWithState - streaming with rocksdb and processing time timer " +
+  testWithAvroEncoding("transformWithState - streaming with rocksdb and processing time timer " +
    "and multiple timers should succeed") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
@@ -651,51 +652,56 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  test("transformWithState - streaming with rocksdb and event time based timer") {
-    val inputData = MemoryStream[(String, Int)]
-    val result =
-      inputData.toDS()
-        .select($"_1".as("key"), timestamp_seconds($"_2").as("eventTime"))
-        .withWatermark("eventTime", "10 seconds")
-        .as[(String, Long)]
-        .groupByKey(_._1)
-        .transformWithState(
-          new MaxEventTimeStatefulProcessor(),
-          TimeMode.EventTime(),
-          OutputMode.Update())
+  testWithAvroEncoding("transformWithState - streaming with rocksdb and event time based timer") {
+    withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
+      classOf[RocksDBStateStoreProvider].getName) {
+      val inputData = MemoryStream[(String, Int)]
+      val result =
+        inputData.toDS()
+          .select($"_1".as("key"), timestamp_seconds($"_2").as("eventTime"))
+          .withWatermark("eventTime", "10 seconds")
+          .as[(String, Long)]
+          .groupByKey(_._1)
+          .transformWithState(
+            new MaxEventTimeStatefulProcessor(),
+            TimeMode.EventTime(),
+            OutputMode.Update())
 
-    testStream(result, OutputMode.Update())(
-      StartStream(),
-
-      AddData(inputData, ("a", 11), ("a", 13), ("a", 15)),
-      // Max event time = 15. Timeout timestamp for "a" = 15 + 5 = 20. Watermark = 15 - 10 = 5.
-      CheckNewAnswer(("a", 15)), // Output = max event time of a
-
-      AddData(inputData, ("a", 4)), // Add data older than watermark for "a"
-      CheckNewAnswer(), // No output as data should get filtered by watermark
-
-      AddData(inputData, ("a", 10)), // Add data newer than watermark for "a"
-      CheckNewAnswer(("a", 15)), // Max event time is still the same
-      // Timeout timestamp for "a" is still 20 as max event time for "a" is still 15.
-      // Watermark is still 5 as max event time for all data is still 15.
-
-      AddData(inputData, ("b", 31)), // Add data newer than watermark for "b", not "a"
-      // Watermark = 31 - 10 = 21, so "a" should be timed out as timeout timestamp for "a" is 20.
-      CheckNewAnswer(("a", -1), ("b", 31)), // State for "a" should timeout and emit -1
-      Execute { q =>
-        // Filter for idle progress events and then verify the custom metrics for stateful operator
-        val progData = q.recentProgress.filter(prog => prog.stateOperators.size > 0)
-        assert(progData.filter(prog =>
-          prog.stateOperators(0).customMetrics.get("numValueStateVars") > 0).size > 0)
-        assert(progData.filter(prog =>
-          prog.stateOperators(0).customMetrics.get("numRegisteredTimers") > 0).size > 0)
-        assert(progData.filter(prog =>
-          prog.stateOperators(0).customMetrics.get("numDeletedTimers") > 0).size > 0)
-      }
-    )
+      testStream(result, OutputMode.Update())(
+        StartStream(),
+
+        AddData(inputData, ("a", 11), ("a", 13), ("a", 15)),
+        // Max event time = 15. Timeout timestamp for "a" = 15 + 5 = 20. Watermark = 15 - 10 = 5.
+        CheckNewAnswer(("a", 15)), // Output = max event time of a
+
+        AddData(inputData, ("a", 4)), // Add data older than watermark for "a"
+        CheckNewAnswer(), // No output as data should get filtered by watermark
+
+        AddData(inputData, ("a", 10)), // Add data newer than watermark for "a"
+        CheckNewAnswer(("a", 15)), // Max event time is still the same
+        // Timeout timestamp for "a" is still 20 as max event time for "a" is still 15.
+        // Watermark is still 5 as max event time for all data is still 15.
+
+        AddData(inputData, ("b", 31)), // Add data newer than watermark for "b", not "a"
+        // Watermark = 31 - 10 = 21, so "a" should be timed out as timeout timestamp for "a" is 20.
+        CheckNewAnswer(("a", -1), ("b", 31)), // State for "a" should timeout and emit -1
+        Execute { q =>
+          // Filter for idle progress events and then verify the custom metrics for
+          // stateful operator
+          val progData = q.recentProgress.filter(prog => prog.stateOperators.size > 0)
+          assert(progData.filter(prog =>
+            prog.stateOperators(0).customMetrics.get("numValueStateVars") > 0).size > 0)
+          assert(progData.filter(prog =>
+            prog.stateOperators(0).customMetrics.get("numRegisteredTimers") > 0).size > 0)
+          assert(progData.filter(prog =>
+            prog.stateOperators(0).customMetrics.get("numDeletedTimers") > 0).size > 0)
+        }
+      )
+    }
   }
 
-  test("Use statefulProcessor without transformWithState - handle should be absent") {
+  testWithAvroEncoding("Use statefulProcessor without transformWithState - " +
+    "handle should be absent") {
     val processor = new RunningCountStatefulProcessor()
     val ex = intercept[Exception] {
       processor.getHandle
@@ -707,7 +713,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     )
   }
 
-  test("transformWithState - batch should succeed") {
+  testWithAvroEncoding("transformWithState - batch should succeed") {
     val inputData = Seq("a", "b")
     val result = inputData.toDS()
       .groupByKey(x => x)
@@ -719,7 +725,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     checkAnswer(df, Seq(("a", "1"), ("b", "1")).toDF())
   }
 
-  test("transformWithState - test deleteIfExists operator") {
+  testWithAvroEncoding("transformWithState - test deleteIfExists operator") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->
@@ -760,7 +766,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  test("transformWithState - two input streams") {
+  testWithAvroEncoding("transformWithState - two input streams") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->
@@ -790,7 +796,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  test("transformWithState - three input streams") {
+  testWithAvroEncoding("transformWithState - three input streams") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->
@@ -825,7 +831,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  test("transformWithState - two input streams, different key type") {
+  testWithAvroEncoding("transformWithState - two input streams, different key type") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->
@@ -872,7 +878,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
         OutputMode.Update())
   }
 
-  test("transformWithState - availableNow trigger mode, rate limit is respected") {
+  testWithAvroEncoding("transformWithState - availableNow trigger mode, rate limit is respected") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
       withTempDir { srcDir =>
@@ -913,7 +919,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  test("transformWithState - availableNow trigger mode, multiple restarts") {
+  testWithAvroEncoding("transformWithState - availableNow trigger mode, multiple restarts") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
       withTempDir { srcDir =>
@@ -951,7 +957,8 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  test("transformWithState - verify StateSchemaV3 writes correct SQL schema of key/value") {
+  testWithAvroEncoding("transformWithState - verify StateSchemaV3 writes " +
+    "correct SQL schema of key/value") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->
@@ -1033,7 +1040,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  test("transformWithState - verify that OperatorStateMetadataV2" +
+  testWithAvroEncoding("transformWithState - verify that OperatorStateMetadataV2" +
     " file is being written correctly") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
@@ -1077,7 +1084,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  test("test that invalid schema evolution fails query for column family") {
+  testWithAvroEncoding("test that invalid schema evolution fails query for column family") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->
@@ -1157,7 +1164,8 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  test("test that changing between different state variable types fails") {
+  testWithAvroEncoding("test that changing between different state " +
+    "variable types fails") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->
@@ -1346,7 +1354,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  test("test query restart succeeds") {
+  testWithAvroEncoding("test query restart succeeds") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->
@@ -1431,7 +1439,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     new Path(stateCheckpointPath, "_stateSchema/default/")
   }
 
-  test("transformWithState - verify that metadata and schema logs are purged") {
+  testWithAvroEncoding("transformWithState - verify that metadata and schema logs are purged") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->

From c1db91da7a66094f10f2f5db81d41d9e5546c061 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Fri, 25 Oct 2024 10:25:20 -0700
Subject: [PATCH 05/30] adding enum

---
 .../spark/sql/execution/streaming/state/StateStore.scala   | 7 +++++++
 .../spark/sql/execution/streaming/statefulOperators.scala  | 2 +-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
index 09221d8374e24..f0b6fdd41ba18 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
@@ -41,6 +41,13 @@ import org.apache.spark.sql.execution.streaming.StatefulOperatorStateInfo
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.util.{NextIterator, ThreadUtils, Utils}
 
+sealed trait StateStoreEncoding
+
+object StateStoreEncoding {
+  case object UnsafeRow extends StateStoreEncoding
+  case object Avro extends StateStoreEncoding
+}
+
 /**
  * Base trait for a versioned key-value store which provides read operations. Each instance of a
  * `ReadStateStore` represents a specific version of state data, and such instances are created
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
index 55e2f3704c7c0..13d10d5aac525 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
@@ -328,7 +328,7 @@ trait StateStoreWriter
   }
 
   lazy val useAvroEncoding: Boolean =
-    conf.getConf(SQLConf.STREAMING_STATE_STORE_ENCODING_FORMAT) == "Avro"
+    conf.getConf(SQLConf.STREAMING_STATE_STORE_ENCODING_FORMAT) == StateStoreEncoding.Avro.toString
 
   /** Set the operator level metrics */
   protected def setOperatorMetrics(numStateStoreInstances: Int = 1): Unit = {

From a30a29d6c7d74131a4bf737f5d23d48914a3adae Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Fri, 25 Oct 2024 15:20:57 -0700
Subject: [PATCH 06/30] feedback and test

---
 .../spark/sql/avro/AvroDeserializer.scala     |  2 +-
 .../StateStoreColumnFamilySchemaUtils.scala   |  6 +--
 .../streaming/StateTypesEncoderUtils.scala    | 16 ++++---
 .../StatefulProcessorHandleImpl.scala         | 18 ++++++++
 .../streaming/state/ListStateSuite.scala      | 37 +++++++++-------
 .../streaming/state/MapStateSuite.scala       |  4 +-
 .../streaming/state/RocksDBSuite.scala        |  4 +-
 .../state/StatefulProcessorHandleSuite.scala  |  4 +-
 .../TransformWithListStateSuite.scala         | 16 +++----
 .../streaming/TransformWithStateSuite.scala   | 42 +++++++++----------
 10 files changed, 90 insertions(+), 59 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala
index 18eea38930b6d..7addc7608260e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala
@@ -29,7 +29,7 @@ import org.apache.avro.Schema.Type._
 import org.apache.avro.generic._
 import org.apache.avro.util.Utf8
 
-import org.apache.spark.sql.avro.AvroUtils.{ nonNullUnionBranches, toFieldStr, AvroMatchedField}
+import org.apache.spark.sql.avro.AvroUtils.{nonNullUnionBranches, toFieldStr, AvroMatchedField}
 import org.apache.spark.sql.catalyst.{InternalRow, NoopFilters, StructFilters}
 import org.apache.spark.sql.catalyst.expressions.{SpecificInternalRow, UnsafeArrayData}
 import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, ArrayData, DateTimeUtils, GenericArrayData}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
index 2214af226b300..6a4f59644e292 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
@@ -39,11 +39,11 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) {
       val avroOptions = AvroOptions(Map.empty)
       val keyAvroType = SchemaConverters.toAvroType(keySchema)
       val keySer = new AvroSerializer(keySchema, keyAvroType, nullable = false)
-      val ser = new AvroSerializer(valSchema, avroType, nullable = false)
-      val de = new AvroDeserializer(avroType, valSchema,
+      val valueSerializer = new AvroSerializer(valSchema, avroType, nullable = false)
+      val valueDeserializer = new AvroDeserializer(avroType, valSchema,
         avroOptions.datetimeRebaseModeInRead, avroOptions.useStableIdForUnionType,
         avroOptions.stableIdPrefixForUnionType, avroOptions.recursiveFieldMaxDepth)
-      Some(AvroEncoderSpec(keySer, ser, de))
+      Some(AvroEncoderSpec(keySer, valueSerializer, valueDeserializer))
     } else {
       None
     }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala
index dcfc797586aac..ee96f42e097b4 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala
@@ -188,9 +188,9 @@ class AvroTypesEncoder[V](
     valEncoder: Encoder[V],
     stateName: String,
     hasTtl: Boolean,
-    avroSerde: Option[AvroEncoderSpec]) extends StateTypesEncoder[V, Array[Byte]] {
+    avroEnc: Option[AvroEncoderSpec]) extends StateTypesEncoder[V, Array[Byte]] {
 
-  val out = new ByteArrayOutputStream
+  private lazy val out = new ByteArrayOutputStream
 
   /** Variables reused for value conversions between spark sql and object */
   private val keySerializer = keyEncoder.createSerializer()
@@ -198,7 +198,9 @@ class AvroTypesEncoder[V](
   private val objToRowSerializer = valExpressionEnc.createSerializer()
   private val rowToObjDeserializer = valExpressionEnc.resolveAndBind().createDeserializer()
 
+  // case class -> dataType
   private val keySchema = keyEncoder.schema
+  // dataType -> avroType
   private val keyAvroType = SchemaConverters.toAvroType(keySchema)
 
   // case class -> dataType
@@ -211,9 +213,10 @@ class AvroTypesEncoder[V](
     if (keyOption.isEmpty) {
       throw StateStoreErrors.implicitKeyNotFound(stateName)
     }
+    assert(avroEnc.isDefined)
 
     val keyRow = keySerializer.apply(keyOption.get).copy() // V -> InternalRow
-    val avroData = avroSerde.get.keySerializer.serialize(keyRow) // InternalRow -> GenericDataRecord
+    val avroData = avroEnc.get.keySerializer.serialize(keyRow) // InternalRow -> GenericDataRecord
 
     out.reset()
     val encoder = EncoderFactory.get().directBinaryEncoder(out, null)
@@ -225,9 +228,10 @@ class AvroTypesEncoder[V](
   }
 
   override def encodeValue(value: V): Array[Byte] = {
+    assert(avroEnc.isDefined)
     val objRow: InternalRow = objToRowSerializer.apply(value).copy() // V -> InternalRow
     val avroData =
-      avroSerde.get.valueSerializer.serialize(objRow) // InternalRow -> GenericDataRecord
+      avroEnc.get.valueSerializer.serialize(objRow) // InternalRow -> GenericDataRecord
     out.reset()
 
     val encoder = EncoderFactory.get().directBinaryEncoder(out, null)
@@ -240,16 +244,18 @@ class AvroTypesEncoder[V](
   }
 
   override def decodeValue(row: Array[Byte]): V = {
+    assert(avroEnc.isDefined)
     val reader = new GenericDatumReader[Any](valueAvroType)
     val decoder = DecoderFactory.get().binaryDecoder(row, 0, row.length, null)
     val genericData = reader.read(null, decoder) // bytes -> GenericDataRecord
-    val internalRow = avroSerde.get.valueDeserializer.deserialize(
+    val internalRow = avroEnc.get.valueDeserializer.deserialize(
       genericData).orNull.asInstanceOf[InternalRow] // GenericDataRecord -> InternalRow
     if (hasTtl) {
       rowToObjDeserializer.apply(internalRow.getStruct(0, valEncoder.schema.length))
     } else rowToObjDeserializer.apply(internalRow)
   }
 
+  // TODO: Implement the below methods for TTL
   override def encodeValue(value: V, expirationMs: Long): Array[Byte] = {
     throw new UnsupportedOperationException
   }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
index 88f0be0b2269f..7212514212964 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
@@ -258,6 +258,24 @@ class StatefulProcessorHandleImpl(
     resultState
   }
 
+  // This method is for unit-testing ListState, as the avroEnc will not be
+  // populated unless the handle is created through the TransformWithStateExec operator
+  private[sql] def getListStateWithAvro[T](
+      stateName: String,
+      valEncoder: Encoder[T],
+      useAvro: Boolean): ListState[T] = {
+    verifyStateVarOperations("get_list_state", CREATED)
+    val avroEnc = if (useAvro) {
+      new StateStoreColumnFamilySchemaUtils(true).getListStateSchema[T](
+        stateName, keyEncoder, valEncoder, hasTtl = false).avroEnc
+    } else {
+      None
+    }
+    val resultState = new ListStateImpl[T](
+      store, stateName, keyEncoder, valEncoder, avroEnc)
+    resultState
+  }
+
   /**
    * Function to create new or return existing list state variable of given type
    * with ttl. State values will not be returned past ttlDuration, and will be eventually removed
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ListStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ListStateSuite.scala
index e9300464af8dc..cf748da3881ec 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ListStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ListStateSuite.scala
@@ -34,13 +34,15 @@ class ListStateSuite extends StateVariableSuiteBase {
   // overwrite useMultipleValuesPerKey in base suite to be true for list state
   override def useMultipleValuesPerKey: Boolean = true
 
-  private def testMapStateWithNullUserKey()(runListOps: ListState[Long] => Unit): Unit = {
+  private def testMapStateWithNullUserKey(useAvro: Boolean)
+      (runListOps: ListState[Long] => Unit): Unit = {
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val listState: ListState[Long] = handle.getListState[Long]("listState", Encoders.scalaLong)
+      val listState: ListState[Long] = handle.getListStateWithAvro[Long](
+        "listState", Encoders.scalaLong, useAvro)
 
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
       val e = intercept[SparkIllegalArgumentException] {
@@ -57,8 +59,8 @@ class ListStateSuite extends StateVariableSuiteBase {
   }
 
   Seq("appendList", "put").foreach { listImplFunc =>
-    test(s"Test list operation($listImplFunc) with null") {
-      testMapStateWithNullUserKey() { listState =>
+    testWithAvroEnc(s"Test list operation($listImplFunc) with null") { useAvro =>
+      testMapStateWithNullUserKey(useAvro) { listState =>
         listImplFunc match {
           case "appendList" => listState.appendList(null)
           case "put" => listState.put(null)
@@ -67,13 +69,14 @@ class ListStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  test("List state operations for single instance") {
+  testWithAvroEnc("List state operations for single instance") { useAvro =>
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val testState: ListState[Long] = handle.getListState[Long]("testState", Encoders.scalaLong)
+      val testState: ListState[Long] = handle.getListStateWithAvro[Long](
+        "testState", Encoders.scalaLong, useAvro)
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
 
       // simple put and get test
@@ -95,14 +98,16 @@ class ListStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  test("List state operations for multiple instance") {
+  testWithAvroEnc("List state operations for multiple instance") { useAvro =>
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val testState1: ListState[Long] = handle.getListState[Long]("testState1", Encoders.scalaLong)
-      val testState2: ListState[Long] = handle.getListState[Long]("testState2", Encoders.scalaLong)
+      val testState1: ListState[Long] = handle.getListStateWithAvro[Long](
+        "testState1", Encoders.scalaLong, useAvro)
+      val testState2: ListState[Long] = handle.getListStateWithAvro[Long](
+        "testState2", Encoders.scalaLong, useAvro)
 
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
 
@@ -133,16 +138,18 @@ class ListStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  test("List state operations with list, value, another list instances") {
+  testWithAvroEnc("List state operations with list, value, another list instances") { useAvro =>
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val listState1: ListState[Long] = handle.getListState[Long]("listState1", Encoders.scalaLong)
-      val listState2: ListState[Long] = handle.getListState[Long]("listState2", Encoders.scalaLong)
-      val valueState: ValueState[Long] = handle.getValueState[Long](
-        "valueState", Encoders.scalaLong)
+      val listState1: ListState[Long] = handle.getListStateWithAvro[Long](
+        "listState1", Encoders.scalaLong, useAvro)
+      val listState2: ListState[Long] = handle.getListStateWithAvro[Long](
+        "listState2", Encoders.scalaLong, useAvro)
+      val valueState: ValueState[Long] = handle.getValueStateWithAvro[Long](
+        "valueState", Encoders.scalaLong, useAvro = false)
 
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
       // simple put and get test
@@ -245,7 +252,7 @@ class ListStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  test("ListState TTL with non-primitive types") {
+  testWithAvroEnc("ListState TTL with non-primitive types") { useAvro =>
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val timestampMs = 10
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MapStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MapStateSuite.scala
index b067d589de904..d9f29ac69ceae 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MapStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MapStateSuite.scala
@@ -120,9 +120,9 @@ class MapStateSuite extends StateVariableSuiteBase {
       val mapTestState2: MapState[String, Int] =
         handle.getMapState[String, Int]("mapTestState2", Encoders.STRING, Encoders.scalaInt)
       val valueTestState: ValueState[String] =
-        handle.getValueState[String]("valueTestState", Encoders.STRING)
+        handle.getValueStateWithAvro[String]("valueTestState", Encoders.STRING, false)
       val listTestState: ListState[String] =
-        handle.getListState[String]("listTestState", Encoders.STRING)
+        handle.getListStateWithAvro[String]("listTestState", Encoders.STRING, false)
 
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
       // put initial values
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala
index ae5a8c01038b3..89a82e9aa4a2f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala
@@ -128,8 +128,8 @@ trait AlsoTestWithChangelogCheckpointingEnabled
     }
   }
 
-  def testWithAvroEncoding(testName: String, testTags: Tag*)
-                          (testBody: => Any): Unit = {
+  def testWithEncodingTypes(testName: String, testTags: Tag*)
+                           (testBody: => Any): Unit = {
     Seq("UnsafeRow", "Avro").foreach { encoding =>
       super.test(testName + s" (encoding = $encoding)", testTags: _*) {
         // in case tests have any code that needs to execute before every test
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatefulProcessorHandleSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatefulProcessorHandleSuite.scala
index 48a6fd836a462..0c14b8a8601c5 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatefulProcessorHandleSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatefulProcessorHandleSuite.scala
@@ -226,7 +226,7 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase {
         Encoders.STRING, TTLConfig(Duration.ofHours(1)))
 
       // create another state without TTL, this should not be captured in the handle
-      handle.getValueState("testState", Encoders.STRING)
+      handle.getValueStateWithAvro("testState", Encoders.STRING, useAvro = false)
 
       assert(handle.ttlStates.size() === 1)
       assert(handle.ttlStates.get(0) === valueStateWithTTL)
@@ -275,7 +275,7 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase {
       val handle = new StatefulProcessorHandleImpl(store,
         UUID.randomUUID(), stringEncoder, TimeMode.None())
 
-      handle.getValueState("testValueState", Encoders.STRING)
+      handle.getValueStateWithAvro("testValueState", Encoders.STRING, useAvro = false)
       handle.getListState("testListState", Encoders.STRING)
       handle.getMapState("testMapState", Encoders.STRING, Encoders.STRING)
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateSuite.scala
index a515ef24e806d..8e5a2fd183a8d 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateSuite.scala
@@ -129,7 +129,7 @@ class TransformWithListStateSuite extends StreamTest
   with AlsoTestWithChangelogCheckpointingEnabled {
   import testImplicits._
 
-  testWithAvroEncoding("test appending null value in list state throw exception") {
+  testWithEncodingTypes("test appending null value in list state throw exception") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
 
@@ -149,7 +149,7 @@ class TransformWithListStateSuite extends StreamTest
     }
   }
 
-  testWithAvroEncoding("test putting null value in list state throw exception") {
+  testWithEncodingTypes("test putting null value in list state throw exception") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
 
@@ -169,7 +169,7 @@ class TransformWithListStateSuite extends StreamTest
     }
   }
 
-  testWithAvroEncoding("test putting null list in list state throw exception") {
+  testWithEncodingTypes("test putting null list in list state throw exception") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
 
@@ -189,7 +189,7 @@ class TransformWithListStateSuite extends StreamTest
     }
   }
 
-  testWithAvroEncoding("test appending null list in list state throw exception") {
+  testWithEncodingTypes("test appending null list in list state throw exception") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
 
@@ -209,7 +209,7 @@ class TransformWithListStateSuite extends StreamTest
     }
   }
 
-  testWithAvroEncoding("test putting empty list in list state throw exception") {
+  testWithEncodingTypes("test putting empty list in list state throw exception") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
 
@@ -229,7 +229,7 @@ class TransformWithListStateSuite extends StreamTest
     }
   }
 
-  testWithAvroEncoding("test appending empty list in list state throw exception") {
+  testWithEncodingTypes("test appending empty list in list state throw exception") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
 
@@ -249,7 +249,7 @@ class TransformWithListStateSuite extends StreamTest
     }
   }
 
-  testWithAvroEncoding("test list state correctness") {
+  testWithEncodingTypes("test list state correctness") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
 
@@ -307,7 +307,7 @@ class TransformWithListStateSuite extends StreamTest
     }
   }
 
-  testWithAvroEncoding("test ValueState And ListState in Processor") {
+  testWithEncodingTypes("test ValueState And ListState in Processor") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
index 49b16c597c018..11c2e4f418ecf 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala
@@ -399,7 +399,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
 
   import testImplicits._
 
-  testWithAvroEncoding("transformWithState - streaming with rocksdb " +
+  testWithEncodingTypes("transformWithState - streaming with rocksdb " +
     "and invalid processor should fail") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
@@ -421,7 +421,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  testWithAvroEncoding("transformWithState - lazy iterators can properly get/set keyed state") {
+  testWithEncodingTypes("transformWithState - lazy iterators can properly get/set keyed state") {
     class ProcessorWithLazyIterators
       extends StatefulProcessor[Long, Long, Long] {
       @transient protected var _myValueState: ValueState[Long] = _
@@ -496,7 +496,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  testWithAvroEncoding("transformWithState - streaming with rocksdb should succeed") {
+  testWithEncodingTypes("transformWithState - streaming with rocksdb should succeed") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->
@@ -534,7 +534,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  testWithAvroEncoding("transformWithState - streaming with rocksdb and processing time timer " +
+  testWithEncodingTypes("transformWithState - streaming with rocksdb and processing time timer " +
    "should succeed") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
@@ -579,7 +579,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  testWithAvroEncoding("transformWithState - streaming with rocksdb and processing time timer " +
+  testWithEncodingTypes("transformWithState - streaming with rocksdb and processing time timer " +
    "and updating timers should succeed") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
@@ -615,7 +615,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  testWithAvroEncoding("transformWithState - streaming with rocksdb and processing time timer " +
+  testWithEncodingTypes("transformWithState - streaming with rocksdb and processing time timer " +
    "and multiple timers should succeed") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
@@ -652,7 +652,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  testWithAvroEncoding("transformWithState - streaming with rocksdb and event time based timer") {
+  testWithEncodingTypes("transformWithState - streaming with rocksdb and event time based timer") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
       val inputData = MemoryStream[(String, Int)]
@@ -700,7 +700,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  testWithAvroEncoding("Use statefulProcessor without transformWithState - " +
+  testWithEncodingTypes("Use statefulProcessor without transformWithState - " +
     "handle should be absent") {
     val processor = new RunningCountStatefulProcessor()
     val ex = intercept[Exception] {
@@ -713,7 +713,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     )
   }
 
-  testWithAvroEncoding("transformWithState - batch should succeed") {
+  testWithEncodingTypes("transformWithState - batch should succeed") {
     val inputData = Seq("a", "b")
     val result = inputData.toDS()
       .groupByKey(x => x)
@@ -725,7 +725,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     checkAnswer(df, Seq(("a", "1"), ("b", "1")).toDF())
   }
 
-  testWithAvroEncoding("transformWithState - test deleteIfExists operator") {
+  testWithEncodingTypes("transformWithState - test deleteIfExists operator") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->
@@ -766,7 +766,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  testWithAvroEncoding("transformWithState - two input streams") {
+  testWithEncodingTypes("transformWithState - two input streams") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->
@@ -796,7 +796,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  testWithAvroEncoding("transformWithState - three input streams") {
+  testWithEncodingTypes("transformWithState - three input streams") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->
@@ -831,7 +831,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  testWithAvroEncoding("transformWithState - two input streams, different key type") {
+  testWithEncodingTypes("transformWithState - two input streams, different key type") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->
@@ -878,7 +878,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
         OutputMode.Update())
   }
 
-  testWithAvroEncoding("transformWithState - availableNow trigger mode, rate limit is respected") {
+  testWithEncodingTypes("transformWithState - availableNow trigger mode, rate limit is respected") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
       withTempDir { srcDir =>
@@ -919,7 +919,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  testWithAvroEncoding("transformWithState - availableNow trigger mode, multiple restarts") {
+  testWithEncodingTypes("transformWithState - availableNow trigger mode, multiple restarts") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
       withTempDir { srcDir =>
@@ -957,7 +957,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  testWithAvroEncoding("transformWithState - verify StateSchemaV3 writes " +
+  testWithEncodingTypes("transformWithState - verify StateSchemaV3 writes " +
     "correct SQL schema of key/value") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
@@ -1040,7 +1040,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  testWithAvroEncoding("transformWithState - verify that OperatorStateMetadataV2" +
+  testWithEncodingTypes("transformWithState - verify that OperatorStateMetadataV2" +
     " file is being written correctly") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
@@ -1084,7 +1084,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  testWithAvroEncoding("test that invalid schema evolution fails query for column family") {
+  testWithEncodingTypes("test that invalid schema evolution fails query for column family") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->
@@ -1164,7 +1164,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  testWithAvroEncoding("test that changing between different state " +
+  testWithEncodingTypes("test that changing between different state " +
     "variable types fails") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
@@ -1354,7 +1354,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     }
   }
 
-  testWithAvroEncoding("test query restart succeeds") {
+  testWithEncodingTypes("test query restart succeeds") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->
@@ -1439,7 +1439,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest
     new Path(stateCheckpointPath, "_stateSchema/default/")
   }
 
-  testWithAvroEncoding("transformWithState - verify that metadata and schema logs are purged") {
+  testWithEncodingTypes("transformWithState - verify that metadata and schema logs are purged") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->

From 2ebf6a8e7f4089f0eed23eb46f325580a3ed6f64 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Fri, 25 Oct 2024 15:31:55 -0700
Subject: [PATCH 07/30] creating utils class

---
 .../sql/execution/streaming/ListStateImpl.scala     |  7 ++++---
 .../streaming/TransformWithStateExec.scala          |  3 ++-
 .../sql/execution/streaming/statefulOperators.scala | 13 +++++++++----
 3 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
index 50294fa3d0587..a011513b06339 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
@@ -212,9 +212,10 @@ class ListStateImpl[S](
     if (usingAvro) {
       val encodedKey: Array[Byte] = avroTypesEncoder.encodeGroupingKey()
       store.remove(encodedKey, stateName)
-       val entryCount = getEntryCount(encodedKey)
-       TWSMetricsUtils.incrementMetric(metrics, "numRemovedStateRows", entryCount)
-       removeEntryCount(encodedKey)
+      // TODO: Create byte array methods for ListState Metrics
+      // val entryCount = getEntryCount(encodedKey)
+      // TWSMetricsUtils.incrementMetric(metrics, "numRemovedStateRows", entryCount)
+      // removeEntryCount(encodedKey)
     } else {
       val encodedKey = unsafeRowTypesEncoder.encodeGroupingKey()
       store.remove(encodedKey, stateName)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
index eed0593d2d7dd..4cdf49d9ef837 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
@@ -105,7 +105,8 @@ case class TransformWithStateExec(
    */
   private def getDriverProcessorHandle(): DriverStatefulProcessorHandleImpl = {
     val driverProcessorHandle = new DriverStatefulProcessorHandleImpl(
-      timeMode, keyEncoder, initializeAvroEnc = useAvroEncoding)
+      timeMode, keyEncoder, initializeAvroEnc =
+        StatefulOperatorUtils.stateStoreEncoding == StateStoreEncoding.Avro.toString)
     driverProcessorHandle.setHandleState(StatefulProcessorHandleState.PRE_INIT)
     statefulProcessor.setHandle(driverProcessorHandle)
     statefulProcessor.init(outputMode, timeMode)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
index 13d10d5aac525..4c07d41c62f78 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
@@ -29,7 +29,7 @@ import org.apache.hadoop.fs.Path
 import org.apache.spark.SparkContext
 import org.apache.spark.internal.Logging
 import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.AnalysisException
+import org.apache.spark.sql.{AnalysisException, SparkSession}
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeProjection
@@ -97,6 +97,14 @@ object StatefulOperatorStateInfo {
   }
 }
 
+object StatefulOperatorUtils {
+  @transient final val session = SparkSession.getActiveSession.orNull
+
+  lazy val stateStoreEncoding: String =
+    session.sessionState.conf.getConf(
+      SQLConf.STREAMING_STATE_STORE_ENCODING_FORMAT)
+}
+
 /**
  * An operator that reads or writes state from the [[StateStore]].
  * The [[StatefulOperatorStateInfo]] should be filled in by `prepareForExecution` in
@@ -327,9 +335,6 @@ trait StateStoreWriter
     OperatorStateMetadataV1(operatorInfo, stateStoreInfo)
   }
 
-  lazy val useAvroEncoding: Boolean =
-    conf.getConf(SQLConf.STREAMING_STATE_STORE_ENCODING_FORMAT) == StateStoreEncoding.Avro.toString
-
   /** Set the operator level metrics */
   protected def setOperatorMetrics(numStateStoreInstances: Int = 1): Unit = {
     assert(numStateStoreInstances >= 1, s"invalid number of stores: $numStateStoreInstances")

From 0559480044fa9f9fcb24c37971268b339adbe61c Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Thu, 31 Oct 2024 09:36:20 -0700
Subject: [PATCH 08/30] micheal feedback

---
 .../apache/spark/sql/internal/SQLConf.scala   |  2 +-
 .../execution/streaming/ListStateImpl.scala   |  1 +
 .../streaming/ListStateImplWithTTL.scala      |  1 +
 .../execution/streaming/MapStateImpl.scala    |  1 +
 .../streaming/MapStateImplWithTTL.scala       |  1 +
 .../streaming/StateTypesEncoderUtils.scala    | 19 ++++++++++++-------
 .../streaming/TransformWithStateExec.scala    |  3 ++-
 .../execution/streaming/ValueStateImpl.scala  |  1 +
 .../streaming/ValueStateImplWithTTL.scala     |  1 +
 .../state/HDFSBackedStateStoreProvider.scala  |  1 -
 10 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index a8d2af606e994..3f528f7a4aa07 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -2211,7 +2211,7 @@ object SQLConf {
       .version("4.0.0")
       .stringConf
       .checkValue(v => Set("UnsafeRow", "Avro").contains(v),
-        "Valid versions are 'UnsafeRow' and 'Avro'")
+        "Valid values are 'UnsafeRow' and 'Avro'")
       .createWithDefault("UnsafeRow")
 
   // The feature is still in development, so it is still internal.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
index a011513b06339..802633302a8e1 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
@@ -32,6 +32,7 @@ import org.apache.spark.sql.types.StructType
  * @param stateName - name of logical state partition
  * @param keyExprEnc - Spark SQL encoder for key
  * @param valEncoder - Spark SQL encoder for value
+ * @param avroEnc: Optional Avro encoder and decoder to convert between S and Avro row
  * @param metrics - metrics to be updated as part of stateful processing
  * @tparam S - data type of object that will be stored in the list
  */
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala
index 639683e5ff549..d99800165cbbc 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala
@@ -36,6 +36,7 @@ import org.apache.spark.util.NextIterator
  * @param valEncoder - Spark SQL encoder for value
  * @param ttlConfig  - TTL configuration for values  stored in this state
  * @param batchTimestampMs - current batch processing timestamp.
+ * @param avroEnc: Optional Avro encoder and decoder to convert between S and Avro row
  * @param metrics - metrics to be updated as part of stateful processing
  * @tparam S - data type of object that will be stored
  */
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala
index b723020d98e02..aeb30504080cc 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala
@@ -32,6 +32,7 @@ import org.apache.spark.sql.types.StructType
  * @param stateName - name of logical state partition
  * @param keyExprEnc - Spark SQL encoder for key
  * @param valEncoder - Spark SQL encoder for value
+ * @param avroEnc: Optional Avro encoder and decoder to convert between S and Avro row
  * @param metrics - metrics to be updated as part of stateful processing
  * @tparam K - type of key for map state variable
  * @tparam V - type of value for map state variable
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala
index 4020b1b4fd904..592536e026211 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala
@@ -36,6 +36,7 @@ import org.apache.spark.util.NextIterator
  * @param valEncoder - SQL encoder for state variable
  * @param ttlConfig  - the ttl configuration (time to live duration etc.)
  * @param batchTimestampMs - current batch processing timestamp.
+ * @param avroEnc: Optional Avro encoder and decoder to convert between S and Avro row
  * @param metrics - metrics to be updated as part of stateful processing
  * @tparam K - type of key for map state variable
  * @tparam V - type of value for map state variable
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala
index ee96f42e097b4..779b84fd67580 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala
@@ -65,6 +65,14 @@ object TransformWithStateKeyValueRowSchemaUtils {
   }
 }
 
+/**
+ * A trait that defines encoding and decoding operations for state types in Structured Streaming.
+ * This encoder handles the conversion between user-defined types and storage types for both
+ * state keys and values, with optional TTL (Time-To-Live) support.
+ *
+ * @tparam V The user-defined value type to be stored in state
+ * @tparam S The storage type used to represent the state (e.g., UnsafeRow or Array[Byte])
+ */
 trait StateTypesEncoder[V, S] {
   def encodeGroupingKey(): S
   def encodeValue(value: V): S
@@ -78,7 +86,7 @@ trait StateTypesEncoder[V, S] {
  * Helper class providing APIs to encode the grouping key, and user provided values
  * to Spark [[UnsafeRow]].
  *
- * CAUTION: StateTypesEncoder class instance is *not* thread-safe.
+ * CAUTION: UnsafeRowTypesEncoder class instance is *not* thread-safe.
  * This class reuses the keyProjection and valueProjection for encoding grouping
  * key and state value respectively. As UnsafeProjection is not thread safe, this
  * class is also not thread safe.
@@ -171,16 +179,13 @@ object UnsafeRowTypesEncoder {
 
 /**
  * Helper class providing APIs to encode the grouping key, and user provided values
- * to Spark [[UnsafeRow]].
- *
- * CAUTION: StateTypesEncoder class instance is *not* thread-safe.
- * This class reuses the keyProjection and valueProjection for encoding grouping
- * key and state value respectively. As UnsafeProjection is not thread safe, this
- * class is also not thread safe.
+ * to an Avro Byte Array.
  *
  * @param keyEncoder - SQL encoder for the grouping key, key type is implicit
  * @param valEncoder - SQL encoder for value of type `S`
  * @param stateName - name of logical state partition
+ * @param hasTtl - whether or not TTL is enabled for this state variable
+ * @param avroEnc = Avro encoder that should be specified to encode keys and values to byte arrays
  * @tparam V - value type
  */
 class AvroTypesEncoder[V](
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
index 4cdf49d9ef837..74e54d361a660 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
@@ -732,7 +732,8 @@ object TransformWithStateExec {
       queryRunId = UUID.randomUUID(),
       operatorId = 0,
       storeVersion = 0,
-      numPartitions = shufflePartitions
+      numPartitions = shufflePartitions,
+      stateStoreCkptIds = None
     )
 
     new TransformWithStateExec(
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala
index fff67396cce3b..6519d55bcab86 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala
@@ -32,6 +32,7 @@ import org.apache.spark.sql.streaming.ValueState
  * @param keyExprEnc - Spark SQL encoder for key
  * @param valEncoder - Spark SQL encoder for value
  * @param metrics - metrics to be updated as part of stateful processing
+ * @param avroEnc: Optional Avro encoder and decoder to convert between S and Avro row
  * @tparam S - data type of object that will be stored
  */
 class ValueStateImpl[S](
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala
index ac7a83ff65c21..780d10cd4539d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala
@@ -34,6 +34,7 @@ import org.apache.spark.sql.streaming.{TTLConfig, ValueState}
  * @param valEncoder - Spark SQL encoder for value
  * @param ttlConfig  - TTL configuration for values  stored in this state
  * @param batchTimestampMs - current batch processing timestamp.
+ * @param avroEnc: Optional Avro encoder and decoder to convert between S and Avro row
  * @param metrics - metrics to be updated as part of stateful processing
  * @tparam S - data type of object that will be stored
  */
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala
index 899d96e1b341e..217a764698416 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala
@@ -102,7 +102,6 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
       throw StateStoreErrors.unsupportedOperationException("multipleValuesPerKey", "HDFSStateStore")
     }
 
-
     override def get(key: Array[Byte], colFamilyName: String): Array[Byte] = {
       throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
     }

From d3845a5f763e5fc752d9d41cacdbb64b275eb67e Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Fri, 1 Nov 2024 14:09:08 -0700
Subject: [PATCH 09/30] ValueState post-refactor

---
 .../execution/streaming/ListStateImpl.scala   | 286 ++++++----------
 .../streaming/ListStateImplWithTTL.scala      |   6 +-
 .../execution/streaming/MapStateImpl.scala    |   8 +-
 .../streaming/MapStateImplWithTTL.scala       |   6 +-
 .../StateStoreColumnFamilySchemaUtils.scala   |   5 +-
 .../streaming/StateTypesEncoderUtils.scala    | 139 +-------
 .../StatefulProcessorHandleImpl.scala         |  49 +--
 .../streaming/TransformWithStateExec.scala    |   3 +-
 .../execution/streaming/ValueStateImpl.scala  |  59 +---
 .../streaming/ValueStateImplWithTTL.scala     |   6 +-
 .../state/HDFSBackedStateStoreProvider.scala  |  49 +--
 .../streaming/state/RocksDBStateEncoder.scala | 307 ++++--------------
 .../state/RocksDBStateStoreProvider.scala     |  99 +-----
 .../StateSchemaCompatibilityChecker.scala     |   1 +
 .../streaming/state/StateStore.scala          |  67 +---
 .../streaming/statefulOperators.scala         |  14 +-
 .../streaming/state/ListStateSuite.scala      |  37 +--
 .../streaming/state/MapStateSuite.scala       |   4 +-
 .../streaming/state/MemoryStateStore.scala    |  33 +-
 ...sDBStateStoreCheckpointFormatV2Suite.scala |  32 +-
 .../state/StatefulProcessorHandleSuite.scala  |   4 +-
 .../streaming/state/ValueStateSuite.scala     |  76 ++---
 22 files changed, 274 insertions(+), 1016 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
index 802633302a8e1..1047a9e87a837 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
@@ -32,7 +32,6 @@ import org.apache.spark.sql.types.StructType
  * @param stateName - name of logical state partition
  * @param keyExprEnc - Spark SQL encoder for key
  * @param valEncoder - Spark SQL encoder for value
- * @param avroEnc: Optional Avro encoder and decoder to convert between S and Avro row
  * @param metrics - metrics to be updated as part of stateful processing
  * @tparam S - data type of object that will be stored in the list
  */
@@ -41,8 +40,8 @@ class ListStateImpl[S](
      stateName: String,
      keyExprEnc: ExpressionEncoder[Any],
      valEncoder: Encoder[S],
-     avroEnc: Option[AvroEncoderSpec],
-     metrics: Map[String, SQLMetric] = Map.empty)
+     metrics: Map[String, SQLMetric] = Map.empty,
+     avroEnc: Option[AvroEncoderSpec] = None)
   extends ListStateMetricsImpl
   with ListState[S]
   with Logging {
@@ -51,187 +50,106 @@ class ListStateImpl[S](
   override def baseStateName: String = stateName
   override def exprEncSchema: StructType = keyExprEnc.schema
 
-  // If we are using Avro, the avroSerde parameter must be populated
-  // else, we will default to using UnsafeRow.
-  private val usingAvro: Boolean = avroEnc.isDefined
-  private val avroTypesEncoder = new AvroTypesEncoder[S](
-    keyExprEnc, valEncoder, stateName, hasTtl = false, avroEnc)
-  private val unsafeRowTypesEncoder = new UnsafeRowTypesEncoder[S](
-    keyExprEnc, valEncoder, stateName, hasTtl = false)
+  private val stateTypesEncoder = StateTypesEncoder(keyExprEnc, valEncoder, stateName)
 
-  store.createColFamilyIfAbsent(stateName, keyExprEnc.schema, valEncoder.schema,
-    NoPrefixKeyStateEncoderSpec(keyExprEnc.schema), useMultipleValuesPerKey = true)
+  store.createColFamilyIfAbsent(
+    stateName,
+    keyExprEnc.schema,
+    valEncoder.schema,
+    NoPrefixKeyStateEncoderSpec(keyExprEnc.schema),
+    useMultipleValuesPerKey = true,
+    avroEncoderSpec = None)
 
   /** Whether state exists or not. */
-  override def exists(): Boolean = {
-    if (usingAvro) {
-      val encodedKey: Array[Byte] = avroTypesEncoder.encodeGroupingKey()
-      store.get(encodedKey, stateName) != null
-    } else {
-      val encodedKey = unsafeRowTypesEncoder.encodeGroupingKey()
-      store.get(encodedKey, stateName) != null
-    }
-  }
-
-  /**
-   * Get the state value if it exists. If the state does not exist in state store, an
-   * empty iterator is returned.
-   */
-  override def get(): Iterator[S] = {
-    if (usingAvro) {
-      getAvro()
-    } else {
-      getUnsafeRow()
-    }
-  }
-
-  private def getAvro(): Iterator[S] = {
-    val encodedKey: Array[Byte] = avroTypesEncoder.encodeGroupingKey()
-    val avroValuesIterator = store.valuesIterator(encodedKey, stateName)
-    new Iterator[S] {
-      override def hasNext: Boolean = {
-        avroValuesIterator.hasNext
-      }
-
-      override def next(): S = {
-        val valueRow = avroValuesIterator.next()
-        avroTypesEncoder.decodeValue(valueRow)
-      }
-    }
-  }
-
-  private def getUnsafeRow(): Iterator[S] = {
-    val encodedKey = unsafeRowTypesEncoder.encodeGroupingKey()
-    val unsafeRowValuesIterator = store.valuesIterator(encodedKey, stateName)
-    new Iterator[S] {
-      override def hasNext: Boolean = {
-        unsafeRowValuesIterator.hasNext
-      }
-
-      override def next(): S = {
-        val valueUnsafeRow = unsafeRowValuesIterator.next()
-        unsafeRowTypesEncoder.decodeValue(valueUnsafeRow)
-      }
-    }
-  }
-
-  /** Update the value of the list. */
-  override def put(newState: Array[S]): Unit = {
-    validateNewState(newState)
-
-    if (usingAvro) {
-      putAvro(newState)
-    } else {
-      putUnsafeRow(newState)
-    }
-  }
-
-  private def putAvro(newState: Array[S]): Unit = {
-    val encodedKey: Array[Byte] = avroTypesEncoder.encodeGroupingKey()
-    var isFirst = true
-    var entryCount = 0L
-    TWSMetricsUtils.resetMetric(metrics, "numUpdatedStateRows")
-
-    newState.foreach { v =>
-      val encodedValue = avroTypesEncoder.encodeValue(v)
-      if (isFirst) {
-        store.put(encodedKey, encodedValue, stateName)
-        isFirst = false
-      } else {
-        store.merge(encodedKey, encodedValue, stateName)
-      }
-      entryCount += 1
-      TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
-    }
-  }
-
-  private def putUnsafeRow(newState: Array[S]): Unit = {
-    val encodedKey = unsafeRowTypesEncoder.encodeGroupingKey()
-    var isFirst = true
-    var entryCount = 0L
-    TWSMetricsUtils.resetMetric(metrics, "numUpdatedStateRows")
-
-    newState.foreach { v =>
-      val encodedValue = unsafeRowTypesEncoder.encodeValue(v)
-      if (isFirst) {
-        store.put(encodedKey, encodedValue, stateName)
-        isFirst = false
-      } else {
-        store.merge(encodedKey, encodedValue, stateName)
-      }
-      entryCount += 1
-      TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
-    }
-    updateEntryCount(encodedKey, entryCount)
-  }
-
-  /** Append an entry to the list. */
-  override def appendValue(newState: S): Unit = {
-    StateStoreErrors.requireNonNullStateValue(newState, stateName)
-
-    if (usingAvro) {
-      val encodedKey: Array[Byte] = avroTypesEncoder.encodeGroupingKey()
-      val encodedValue = avroTypesEncoder.encodeValue(newState)
-      store.merge(encodedKey, encodedValue, stateName)
-      TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
-    } else {
-      val encodedKey = unsafeRowTypesEncoder.encodeGroupingKey()
-      val entryCount = getEntryCount(encodedKey)
-      val encodedValue = unsafeRowTypesEncoder.encodeValue(newState)
-      store.merge(encodedKey, encodedValue, stateName)
-      TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
-      updateEntryCount(encodedKey, entryCount + 1)
-    }
-  }
-
-  /** Append an entire list to the existing value. */
-  override def appendList(newState: Array[S]): Unit = {
-    validateNewState(newState)
-
-    if (usingAvro) {
-      val encodedKey: Array[Byte] = avroTypesEncoder.encodeGroupingKey()
-      newState.foreach { v =>
-        val encodedValue = avroTypesEncoder.encodeValue(v)
-        store.merge(encodedKey, encodedValue, stateName)
-        TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
-      }
-    } else {
-      val encodedKey = unsafeRowTypesEncoder.encodeGroupingKey()
-      var entryCount = getEntryCount(encodedKey)
-      newState.foreach { v =>
-        val encodedValue = unsafeRowTypesEncoder.encodeValue(v)
-        store.merge(encodedKey, encodedValue, stateName)
-        entryCount += 1
-        TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
-      }
-      updateEntryCount(encodedKey, entryCount)
-    }
-  }
-
-  /** Remove this state. */
-  override def clear(): Unit = {
-    if (usingAvro) {
-      val encodedKey: Array[Byte] = avroTypesEncoder.encodeGroupingKey()
-      store.remove(encodedKey, stateName)
-      // TODO: Create byte array methods for ListState Metrics
-      // val entryCount = getEntryCount(encodedKey)
-      // TWSMetricsUtils.incrementMetric(metrics, "numRemovedStateRows", entryCount)
-      // removeEntryCount(encodedKey)
-    } else {
-      val encodedKey = unsafeRowTypesEncoder.encodeGroupingKey()
-      store.remove(encodedKey, stateName)
-      val entryCount = getEntryCount(encodedKey)
-      TWSMetricsUtils.incrementMetric(metrics, "numRemovedStateRows", entryCount)
-      removeEntryCount(encodedKey)
-    }
-  }
-
-  private def validateNewState(newState: Array[S]): Unit = {
-    StateStoreErrors.requireNonNullStateValue(newState, stateName)
-    StateStoreErrors.requireNonEmptyListStateValue(newState, stateName)
-
-    newState.foreach { v =>
-      StateStoreErrors.requireNonNullStateValue(v, stateName)
-    }
-  }
-}
+   override def exists(): Boolean = {
+     val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
+     val stateValue = store.get(encodedGroupingKey, stateName)
+     stateValue != null
+   }
+
+   /**
+    * Get the state value if it exists. If the state does not exist in state store, an
+    * empty iterator is returned.
+    */
+   override def get(): Iterator[S] = {
+     val encodedKey = stateTypesEncoder.encodeGroupingKey()
+     val unsafeRowValuesIterator = store.valuesIterator(encodedKey, stateName)
+     new Iterator[S] {
+       override def hasNext: Boolean = {
+         unsafeRowValuesIterator.hasNext
+       }
+
+       override def next(): S = {
+         val valueUnsafeRow = unsafeRowValuesIterator.next()
+         stateTypesEncoder.decodeValue(valueUnsafeRow)
+       }
+     }
+   }
+
+   /** Update the value of the list. */
+   override def put(newState: Array[S]): Unit = {
+     validateNewState(newState)
+
+     val encodedKey = stateTypesEncoder.encodeGroupingKey()
+     var isFirst = true
+     var entryCount = 0L
+     TWSMetricsUtils.resetMetric(metrics, "numUpdatedStateRows")
+
+     newState.foreach { v =>
+       val encodedValue = stateTypesEncoder.encodeValue(v)
+       if (isFirst) {
+         store.put(encodedKey, encodedValue, stateName)
+         isFirst = false
+       } else {
+         store.merge(encodedKey, encodedValue, stateName)
+       }
+       entryCount += 1
+       TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
+     }
+     updateEntryCount(encodedKey, entryCount)
+   }
+
+   /** Append an entry to the list. */
+   override def appendValue(newState: S): Unit = {
+     StateStoreErrors.requireNonNullStateValue(newState, stateName)
+     val encodedKey = stateTypesEncoder.encodeGroupingKey()
+     val entryCount = getEntryCount(encodedKey)
+     store.merge(encodedKey,
+         stateTypesEncoder.encodeValue(newState), stateName)
+     TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
+     updateEntryCount(encodedKey, entryCount + 1)
+   }
+
+   /** Append an entire list to the existing value. */
+   override def appendList(newState: Array[S]): Unit = {
+     validateNewState(newState)
+
+     val encodedKey = stateTypesEncoder.encodeGroupingKey()
+     var entryCount = getEntryCount(encodedKey)
+     newState.foreach { v =>
+       val encodedValue = stateTypesEncoder.encodeValue(v)
+       store.merge(encodedKey, encodedValue, stateName)
+       entryCount += 1
+       TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
+     }
+     updateEntryCount(encodedKey, entryCount)
+   }
+
+   /** Remove this state. */
+   override def clear(): Unit = {
+     val encodedKey = stateTypesEncoder.encodeGroupingKey()
+     store.remove(encodedKey, stateName)
+     val entryCount = getEntryCount(encodedKey)
+     TWSMetricsUtils.incrementMetric(metrics, "numRemovedStateRows", entryCount)
+     removeEntryCount(encodedKey)
+   }
+
+   private def validateNewState(newState: Array[S]): Unit = {
+     StateStoreErrors.requireNonNullStateValue(newState, stateName)
+     StateStoreErrors.requireNonEmptyListStateValue(newState, stateName)
+
+     newState.foreach { v =>
+       StateStoreErrors.requireNonNullStateValue(v, stateName)
+     }
+   }
+ }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala
index d99800165cbbc..be47f566bc6a9 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala
@@ -21,7 +21,7 @@ import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, NoPrefixKeyStateEncoderSpec, StateStore, StateStoreErrors}
+import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, StateStore, StateStoreErrors}
 import org.apache.spark.sql.streaming.{ListState, TTLConfig}
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.util.NextIterator
@@ -36,7 +36,6 @@ import org.apache.spark.util.NextIterator
  * @param valEncoder - Spark SQL encoder for value
  * @param ttlConfig  - TTL configuration for values  stored in this state
  * @param batchTimestampMs - current batch processing timestamp.
- * @param avroEnc: Optional Avro encoder and decoder to convert between S and Avro row
  * @param metrics - metrics to be updated as part of stateful processing
  * @tparam S - data type of object that will be stored
  */
@@ -47,7 +46,6 @@ class ListStateImplWithTTL[S](
     valEncoder: Encoder[S],
     ttlConfig: TTLConfig,
     batchTimestampMs: Long,
-    avroEnc: Option[AvroEncoderSpec], // TODO: Add Avro Encoding support for TTL
     metrics: Map[String, SQLMetric] = Map.empty)
   extends SingleKeyTTLStateImpl(stateName, store, keyExprEnc, batchTimestampMs)
   with ListStateMetricsImpl
@@ -57,7 +55,7 @@ class ListStateImplWithTTL[S](
   override def baseStateName: String = stateName
   override def exprEncSchema: StructType = keyExprEnc.schema
 
-  private lazy val stateTypesEncoder = UnsafeRowTypesEncoder(keyExprEnc, valEncoder,
+  private lazy val stateTypesEncoder = StateTypesEncoder(keyExprEnc, valEncoder,
     stateName, hasTtl = true)
 
   private lazy val ttlExpirationMs =
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala
index aeb30504080cc..cb3db19496dd2 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala
@@ -21,7 +21,7 @@ import org.apache.spark.sql.Encoder
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, PrefixKeyScanStateEncoderSpec, StateStore, StateStoreErrors, UnsafeRowPair}
+import org.apache.spark.sql.execution.streaming.state.{PrefixKeyScanStateEncoderSpec, StateStore, StateStoreErrors, UnsafeRowPair}
 import org.apache.spark.sql.streaming.MapState
 import org.apache.spark.sql.types.StructType
 
@@ -32,7 +32,6 @@ import org.apache.spark.sql.types.StructType
  * @param stateName - name of logical state partition
  * @param keyExprEnc - Spark SQL encoder for key
  * @param valEncoder - Spark SQL encoder for value
- * @param avroEnc: Optional Avro encoder and decoder to convert between S and Avro row
  * @param metrics - metrics to be updated as part of stateful processing
  * @tparam K - type of key for map state variable
  * @tparam V - type of value for map state variable
@@ -43,7 +42,6 @@ class MapStateImpl[K, V](
     keyExprEnc: ExpressionEncoder[Any],
     userKeyEnc: Encoder[K],
     valEncoder: Encoder[V],
-    avroEnc: Option[AvroEncoderSpec],
     metrics: Map[String, SQLMetric] = Map.empty) extends MapState[K, V] with Logging {
 
   // Pack grouping key and user key together as a prefixed composite key
@@ -51,8 +49,8 @@ class MapStateImpl[K, V](
     getCompositeKeySchema(keyExprEnc.schema, userKeyEnc.schema)
   }
   private val schemaForValueRow: StructType = valEncoder.schema
-  private val stateTypesEncoder = new CompositeKeyUnsafeRowEncoder(
-    keyExprEnc, userKeyEnc, valEncoder, stateName, hasTtl = false)
+  private val stateTypesEncoder = new CompositeKeyStateEncoder(
+    keyExprEnc, userKeyEnc, valEncoder, stateName)
 
   store.createColFamilyIfAbsent(stateName, schemaForCompositeKeyRow, schemaForValueRow,
     PrefixKeyScanStateEncoderSpec(schemaForCompositeKeyRow, 1))
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala
index 592536e026211..6a3685ad6c46c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala
@@ -22,7 +22,7 @@ import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, PrefixKeyScanStateEncoderSpec, StateStore, StateStoreErrors}
+import org.apache.spark.sql.execution.streaming.state.{PrefixKeyScanStateEncoderSpec, StateStore, StateStoreErrors}
 import org.apache.spark.sql.streaming.{MapState, TTLConfig}
 import org.apache.spark.util.NextIterator
 
@@ -36,7 +36,6 @@ import org.apache.spark.util.NextIterator
  * @param valEncoder - SQL encoder for state variable
  * @param ttlConfig  - the ttl configuration (time to live duration etc.)
  * @param batchTimestampMs - current batch processing timestamp.
- * @param avroEnc: Optional Avro encoder and decoder to convert between S and Avro row
  * @param metrics - metrics to be updated as part of stateful processing
  * @tparam K - type of key for map state variable
  * @tparam V - type of value for map state variable
@@ -50,13 +49,12 @@ class MapStateImplWithTTL[K, V](
     valEncoder: Encoder[V],
     ttlConfig: TTLConfig,
     batchTimestampMs: Long,
-    avroEnc: Option[AvroEncoderSpec], // TODO: Add Avro Encoding support for TTL
     metrics: Map[String, SQLMetric] = Map.empty)
   extends CompositeKeyTTLStateImpl[K](stateName, store,
     keyExprEnc, userKeyEnc, batchTimestampMs)
   with MapState[K, V] with Logging {
 
-  private val stateTypesEncoder = new CompositeKeyUnsafeRowEncoder(
+  private val stateTypesEncoder = new CompositeKeyStateEncoder(
     keyExprEnc, userKeyEnc, valEncoder, stateName, hasTtl = true)
 
   private val ttlExpirationMs =
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
index 6a4f59644e292..744fd6d5b6b14 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
@@ -39,11 +39,14 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) {
       val avroOptions = AvroOptions(Map.empty)
       val keyAvroType = SchemaConverters.toAvroType(keySchema)
       val keySer = new AvroSerializer(keySchema, keyAvroType, nullable = false)
+      val keyDe = new AvroDeserializer(keyAvroType, keySchema,
+        avroOptions.datetimeRebaseModeInRead, avroOptions.useStableIdForUnionType,
+        avroOptions.stableIdPrefixForUnionType, avroOptions.recursiveFieldMaxDepth)
       val valueSerializer = new AvroSerializer(valSchema, avroType, nullable = false)
       val valueDeserializer = new AvroDeserializer(avroType, valSchema,
         avroOptions.datetimeRebaseModeInRead, avroOptions.useStableIdForUnionType,
         avroOptions.stableIdPrefixForUnionType, avroOptions.recursiveFieldMaxDepth)
-      Some(AvroEncoderSpec(keySer, valueSerializer, valueDeserializer))
+      Some(AvroEncoderSpec(keySer, keyDe, valueSerializer, valueDeserializer))
     } else {
       None
     }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala
index 779b84fd67580..b70f9699195d4 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala
@@ -17,18 +17,12 @@
 
 package org.apache.spark.sql.execution.streaming
 
-import java.io.ByteArrayOutputStream
-
-import org.apache.avro.generic.{GenericDatumReader, GenericDatumWriter}
-import org.apache.avro.io.{DecoderFactory, EncoderFactory}
-
 import org.apache.spark.sql.Encoder
-import org.apache.spark.sql.avro.SchemaConverters
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.encoders.{encoderFor, ExpressionEncoder}
 import org.apache.spark.sql.catalyst.expressions.{UnsafeProjection, UnsafeRow}
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, StateStoreErrors}
+import org.apache.spark.sql.execution.streaming.state.StateStoreErrors
 import org.apache.spark.sql.types._
 
 /**
@@ -65,28 +59,11 @@ object TransformWithStateKeyValueRowSchemaUtils {
   }
 }
 
-/**
- * A trait that defines encoding and decoding operations for state types in Structured Streaming.
- * This encoder handles the conversion between user-defined types and storage types for both
- * state keys and values, with optional TTL (Time-To-Live) support.
- *
- * @tparam V The user-defined value type to be stored in state
- * @tparam S The storage type used to represent the state (e.g., UnsafeRow or Array[Byte])
- */
-trait StateTypesEncoder[V, S] {
-  def encodeGroupingKey(): S
-  def encodeValue(value: V): S
-  def decodeValue(row: S): V
-  def encodeValue(value: V, expirationMs: Long): S
-  def decodeTtlExpirationMs(row: S): Option[Long]
-  def isExpired(row: S, batchTimestampMs: Long): Boolean
-}
-
 /**
  * Helper class providing APIs to encode the grouping key, and user provided values
  * to Spark [[UnsafeRow]].
  *
- * CAUTION: UnsafeRowTypesEncoder class instance is *not* thread-safe.
+ * CAUTION: StateTypesEncoder class instance is *not* thread-safe.
  * This class reuses the keyProjection and valueProjection for encoding grouping
  * key and state value respectively. As UnsafeProjection is not thread safe, this
  * class is also not thread safe.
@@ -96,11 +73,11 @@ trait StateTypesEncoder[V, S] {
  * @param stateName - name of logical state partition
  * @tparam V - value type
  */
-class UnsafeRowTypesEncoder[V](
+class StateTypesEncoder[V](
     keyEncoder: ExpressionEncoder[Any],
     valEncoder: Encoder[V],
     stateName: String,
-    hasTtl: Boolean) extends StateTypesEncoder[V, UnsafeRow] {
+    hasTtl: Boolean) {
 
   /** Variables reused for value conversions between spark sql and object */
   private val keySerializer = keyEncoder.createSerializer()
@@ -166,121 +143,23 @@ class UnsafeRowTypesEncoder[V](
   }
 }
 
-
-object UnsafeRowTypesEncoder {
+object StateTypesEncoder {
   def apply[V](
       keyEncoder: ExpressionEncoder[Any],
       valEncoder: Encoder[V],
       stateName: String,
-      hasTtl: Boolean = false): UnsafeRowTypesEncoder[V] = {
-    new UnsafeRowTypesEncoder[V](keyEncoder, valEncoder, stateName, hasTtl)
-  }
-}
-
-/**
- * Helper class providing APIs to encode the grouping key, and user provided values
- * to an Avro Byte Array.
- *
- * @param keyEncoder - SQL encoder for the grouping key, key type is implicit
- * @param valEncoder - SQL encoder for value of type `S`
- * @param stateName - name of logical state partition
- * @param hasTtl - whether or not TTL is enabled for this state variable
- * @param avroEnc = Avro encoder that should be specified to encode keys and values to byte arrays
- * @tparam V - value type
- */
-class AvroTypesEncoder[V](
-    keyEncoder: ExpressionEncoder[Any],
-    valEncoder: Encoder[V],
-    stateName: String,
-    hasTtl: Boolean,
-    avroEnc: Option[AvroEncoderSpec]) extends StateTypesEncoder[V, Array[Byte]] {
-
-  private lazy val out = new ByteArrayOutputStream
-
-  /** Variables reused for value conversions between spark sql and object */
-  private val keySerializer = keyEncoder.createSerializer()
-  private val valExpressionEnc = encoderFor(valEncoder)
-  private val objToRowSerializer = valExpressionEnc.createSerializer()
-  private val rowToObjDeserializer = valExpressionEnc.resolveAndBind().createDeserializer()
-
-  // case class -> dataType
-  private val keySchema = keyEncoder.schema
-  // dataType -> avroType
-  private val keyAvroType = SchemaConverters.toAvroType(keySchema)
-
-  // case class -> dataType
-  private val valSchema: StructType = valEncoder.schema
-  // dataType -> avroType
-  private val valueAvroType = SchemaConverters.toAvroType(valSchema)
-
-  override def encodeGroupingKey(): Array[Byte] = {
-    val keyOption = ImplicitGroupingKeyTracker.getImplicitKeyOption
-    if (keyOption.isEmpty) {
-      throw StateStoreErrors.implicitKeyNotFound(stateName)
-    }
-    assert(avroEnc.isDefined)
-
-    val keyRow = keySerializer.apply(keyOption.get).copy() // V -> InternalRow
-    val avroData = avroEnc.get.keySerializer.serialize(keyRow) // InternalRow -> GenericDataRecord
-
-    out.reset()
-    val encoder = EncoderFactory.get().directBinaryEncoder(out, null)
-    val writer = new GenericDatumWriter[Any](keyAvroType)
-
-    writer.write(avroData, encoder)
-    encoder.flush()
-    out.toByteArray
-  }
-
-  override def encodeValue(value: V): Array[Byte] = {
-    assert(avroEnc.isDefined)
-    val objRow: InternalRow = objToRowSerializer.apply(value).copy() // V -> InternalRow
-    val avroData =
-      avroEnc.get.valueSerializer.serialize(objRow) // InternalRow -> GenericDataRecord
-    out.reset()
-
-    val encoder = EncoderFactory.get().directBinaryEncoder(out, null)
-    val writer = new GenericDatumWriter[Any](
-      valueAvroType) // Defining Avro writer for this struct type
-
-    writer.write(avroData, encoder) // GenericDataRecord -> bytes
-    encoder.flush()
-    out.toByteArray
-  }
-
-  override def decodeValue(row: Array[Byte]): V = {
-    assert(avroEnc.isDefined)
-    val reader = new GenericDatumReader[Any](valueAvroType)
-    val decoder = DecoderFactory.get().binaryDecoder(row, 0, row.length, null)
-    val genericData = reader.read(null, decoder) // bytes -> GenericDataRecord
-    val internalRow = avroEnc.get.valueDeserializer.deserialize(
-      genericData).orNull.asInstanceOf[InternalRow] // GenericDataRecord -> InternalRow
-    if (hasTtl) {
-      rowToObjDeserializer.apply(internalRow.getStruct(0, valEncoder.schema.length))
-    } else rowToObjDeserializer.apply(internalRow)
-  }
-
-  // TODO: Implement the below methods for TTL
-  override def encodeValue(value: V, expirationMs: Long): Array[Byte] = {
-    throw new UnsupportedOperationException
-  }
-
-  override def decodeTtlExpirationMs(row: Array[Byte]): Option[Long] = {
-    throw new UnsupportedOperationException
-  }
-
-  override def isExpired(row: Array[Byte], batchTimestampMs: Long): Boolean = {
-    throw new UnsupportedOperationException
+      hasTtl: Boolean = false): StateTypesEncoder[V] = {
+    new StateTypesEncoder[V](keyEncoder, valEncoder, stateName, hasTtl)
   }
 }
 
-class CompositeKeyUnsafeRowEncoder[K, V](
+class CompositeKeyStateEncoder[K, V](
     keyEncoder: ExpressionEncoder[Any],
     userKeyEnc: Encoder[K],
     valEncoder: Encoder[V],
     stateName: String,
     hasTtl: Boolean = false)
-  extends UnsafeRowTypesEncoder[V](keyEncoder, valEncoder, stateName, hasTtl) {
+  extends StateTypesEncoder[V](keyEncoder, valEncoder, stateName, hasTtl) {
   import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
 
   /** Encoders */
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
index 7212514212964..8b6e6f0ba3508 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
@@ -141,29 +141,11 @@ class StatefulProcessorHandleImpl(
       valEncoder: Encoder[T]): ValueState[T] = {
     verifyStateVarOperations("get_value_state", CREATED)
     val resultState = new ValueStateImpl[T](
-      store, stateName, keyEncoder, valEncoder, schemas(stateName).avroEnc, metrics)
+      store, stateName, keyEncoder, valEncoder, metrics, schemas(stateName).avroEnc)
     TWSMetricsUtils.incrementMetric(metrics, "numValueStateVars")
     resultState
   }
 
-  // This method is for unit-testing ValueState, as the avroEnc will not be
-  // populated unless the handle is created through the TransformWithStateExec operator
-  private[sql] def getValueStateWithAvro[T](
-      stateName: String,
-      valEncoder: Encoder[T],
-      useAvro: Boolean): ValueState[T] = {
-    verifyStateVarOperations("get_value_state", CREATED)
-    val avroEnc = if (useAvro) {
-      new StateStoreColumnFamilySchemaUtils(true).getValueStateSchema[T](
-        stateName, keyEncoder, valEncoder, hasTtl = false).avroEnc
-    } else {
-      None
-    }
-    val resultState = new ValueStateImpl[T](
-      store, stateName, keyEncoder, valEncoder, avroEnc)
-    resultState
-  }
-
   override def getValueState[T](
       stateName: String,
       valEncoder: Encoder[T],
@@ -173,7 +155,7 @@ class StatefulProcessorHandleImpl(
 
     assert(batchTimestampMs.isDefined)
     val valueStateWithTTL = new ValueStateImplWithTTL[T](store, stateName,
-      keyEncoder, valEncoder, ttlConfig, batchTimestampMs.get, avroEnc = None, metrics)
+      keyEncoder, valEncoder, ttlConfig, batchTimestampMs.get, metrics)
     ttlStates.add(valueStateWithTTL)
     TWSMetricsUtils.incrementMetric(metrics, "numValueStateWithTTLVars")
 
@@ -252,30 +234,11 @@ class StatefulProcessorHandleImpl(
 
   override def getListState[T](stateName: String, valEncoder: Encoder[T]): ListState[T] = {
     verifyStateVarOperations("get_list_state", CREATED)
-    val resultState = new ListStateImpl[T](
-      store, stateName, keyEncoder, valEncoder, schemas(stateName).avroEnc, metrics)
+    val resultState = new ListStateImpl[T](store, stateName, keyEncoder, valEncoder, metrics)
     TWSMetricsUtils.incrementMetric(metrics, "numListStateVars")
     resultState
   }
 
-  // This method is for unit-testing ListState, as the avroEnc will not be
-  // populated unless the handle is created through the TransformWithStateExec operator
-  private[sql] def getListStateWithAvro[T](
-      stateName: String,
-      valEncoder: Encoder[T],
-      useAvro: Boolean): ListState[T] = {
-    verifyStateVarOperations("get_list_state", CREATED)
-    val avroEnc = if (useAvro) {
-      new StateStoreColumnFamilySchemaUtils(true).getListStateSchema[T](
-        stateName, keyEncoder, valEncoder, hasTtl = false).avroEnc
-    } else {
-      None
-    }
-    val resultState = new ListStateImpl[T](
-      store, stateName, keyEncoder, valEncoder, avroEnc)
-    resultState
-  }
-
   /**
    * Function to create new or return existing list state variable of given type
    * with ttl. State values will not be returned past ttlDuration, and will be eventually removed
@@ -301,7 +264,7 @@ class StatefulProcessorHandleImpl(
 
     assert(batchTimestampMs.isDefined)
     val listStateWithTTL = new ListStateImplWithTTL[T](store, stateName,
-      keyEncoder, valEncoder, ttlConfig, batchTimestampMs.get, avroEnc = None, metrics)
+      keyEncoder, valEncoder, ttlConfig, batchTimestampMs.get, metrics)
     TWSMetricsUtils.incrementMetric(metrics, "numListStateWithTTLVars")
     ttlStates.add(listStateWithTTL)
 
@@ -314,7 +277,7 @@ class StatefulProcessorHandleImpl(
       valEncoder: Encoder[V]): MapState[K, V] = {
     verifyStateVarOperations("get_map_state", CREATED)
     val resultState = new MapStateImpl[K, V](store, stateName, keyEncoder,
-      userKeyEnc, valEncoder, avroEnc = None, metrics)
+      userKeyEnc, valEncoder, metrics)
     TWSMetricsUtils.incrementMetric(metrics, "numMapStateVars")
     resultState
   }
@@ -329,7 +292,7 @@ class StatefulProcessorHandleImpl(
 
     assert(batchTimestampMs.isDefined)
     val mapStateWithTTL = new MapStateImplWithTTL[K, V](store, stateName, keyEncoder, userKeyEnc,
-      valEncoder, ttlConfig, batchTimestampMs.get, avroEnc = None, metrics)
+      valEncoder, ttlConfig, batchTimestampMs.get, metrics)
     TWSMetricsUtils.incrementMetric(metrics, "numMapStateWithTTLVars")
     ttlStates.add(mapStateWithTTL)
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
index 74e54d361a660..ecc3d2a573eb2 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
@@ -104,9 +104,10 @@ case class TransformWithStateExec(
    * @return a new instance of the driver processor handle
    */
   private def getDriverProcessorHandle(): DriverStatefulProcessorHandleImpl = {
+
     val driverProcessorHandle = new DriverStatefulProcessorHandleImpl(
       timeMode, keyEncoder, initializeAvroEnc =
-        StatefulOperatorUtils.stateStoreEncoding == StateStoreEncoding.Avro.toString)
+        stateStoreEncoding == StateStoreEncoding.Avro.toString)
     driverProcessorHandle.setHandleState(StatefulProcessorHandleState.PRE_INIT)
     statefulProcessor.setHandle(driverProcessorHandle)
     statefulProcessor.init(outputMode, timeMode)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala
index 6519d55bcab86..db8c405ee6193 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala
@@ -20,8 +20,7 @@ import org.apache.spark.internal.Logging
 import org.apache.spark.sql.Encoder
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.execution.metric.SQLMetric
-import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, StateStore}
-import org.apache.spark.sql.execution.streaming.state.AvroEncoderSpec
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, NoPrefixKeyStateEncoderSpec, StateStore}
 import org.apache.spark.sql.streaming.ValueState
 
 /**
@@ -32,7 +31,6 @@ import org.apache.spark.sql.streaming.ValueState
  * @param keyExprEnc - Spark SQL encoder for key
  * @param valEncoder - Spark SQL encoder for value
  * @param metrics - metrics to be updated as part of stateful processing
- * @param avroEnc: Optional Avro encoder and decoder to convert between S and Avro row
  * @tparam S - data type of object that will be stored
  */
 class ValueStateImpl[S](
@@ -40,23 +38,17 @@ class ValueStateImpl[S](
     stateName: String,
     keyExprEnc: ExpressionEncoder[Any],
     valEncoder: Encoder[S],
-    avroEnc: Option[AvroEncoderSpec],
-    metrics: Map[String, SQLMetric] = Map.empty)
+    metrics: Map[String, SQLMetric] = Map.empty,
+    avroEnc: Option[AvroEncoderSpec] = None)
   extends ValueState[S] with Logging {
 
-  // If we are using Avro, the avroSerde parameter must be populated
-  // else, we will default to using UnsafeRow.
-  private val usingAvro: Boolean = avroEnc.isDefined
-  private val avroTypesEncoder = new AvroTypesEncoder[S](
-    keyExprEnc, valEncoder, stateName, hasTtl = false, avroEnc)
-  private val unsafeRowTypesEncoder = new UnsafeRowTypesEncoder[S](
-    keyExprEnc, valEncoder, stateName, hasTtl = false)
+  private val stateTypesEncoder = StateTypesEncoder(keyExprEnc, valEncoder, stateName)
 
   initialize()
 
   private def initialize(): Unit = {
     store.createColFamilyIfAbsent(stateName, keyExprEnc.schema, valEncoder.schema,
-      NoPrefixKeyStateEncoderSpec(keyExprEnc.schema))
+      NoPrefixKeyStateEncoderSpec(keyExprEnc.schema), avroEncoderSpec = avroEnc)
   }
 
   /** Function to check if state exists. Returns true if present and false otherwise */
@@ -71,28 +63,11 @@ class ValueStateImpl[S](
 
   /** Function to return associated value with key if exists and null otherwise */
   override def get(): S = {
-    if (usingAvro) {
-      getAvro()
-    } else {
-      getUnsafeRow()
-    }
-  }
-
-  private def getAvro(): S = {
-    val encodedGroupingKey = avroTypesEncoder.encodeGroupingKey()
+    val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
     val retRow = store.get(encodedGroupingKey, stateName)
-    if (retRow != null) {
-      avroTypesEncoder.decodeValue(retRow)
-    } else {
-      null.asInstanceOf[S]
-    }
-  }
 
-  private def getUnsafeRow(): S = {
-    val encodedGroupingKey = unsafeRowTypesEncoder.encodeGroupingKey()
-    val retRow = store.get(encodedGroupingKey, stateName)
     if (retRow != null) {
-      unsafeRowTypesEncoder.decodeValue(retRow)
+      stateTypesEncoder.decodeValue(retRow)
     } else {
       null.asInstanceOf[S]
     }
@@ -100,27 +75,15 @@ class ValueStateImpl[S](
 
   /** Function to update and overwrite state associated with given key */
   override def update(newState: S): Unit = {
-    if (usingAvro) {
-      val encodedValue = avroTypesEncoder.encodeValue(newState)
-      store.put(avroTypesEncoder.encodeGroupingKey(),
-        encodedValue, stateName)
-    } else {
-      val encodedValue = unsafeRowTypesEncoder.encodeValue(newState)
-      store.put(unsafeRowTypesEncoder.encodeGroupingKey(),
-        encodedValue, stateName)
-    }
+    val encodedValue = stateTypesEncoder.encodeValue(newState)
+    store.put(stateTypesEncoder.encodeGroupingKey(),
+      encodedValue, stateName)
     TWSMetricsUtils.incrementMetric(metrics, "numUpdatedStateRows")
   }
 
   /** Function to remove state for given key */
   override def clear(): Unit = {
-    if (usingAvro) {
-      val encodedKey = avroTypesEncoder.encodeGroupingKey()
-      store.remove(encodedKey, stateName)
-    } else {
-      val encodedKey = unsafeRowTypesEncoder.encodeGroupingKey()
-      store.remove(encodedKey, stateName)
-    }
+    store.remove(stateTypesEncoder.encodeGroupingKey(), stateName)
     TWSMetricsUtils.incrementMetric(metrics, "numRemovedStateRows")
   }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala
index 780d10cd4539d..145cd90264910 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala
@@ -21,7 +21,7 @@ import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, NoPrefixKeyStateEncoderSpec, StateStore}
+import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, StateStore}
 import org.apache.spark.sql.streaming.{TTLConfig, ValueState}
 
 /**
@@ -34,7 +34,6 @@ import org.apache.spark.sql.streaming.{TTLConfig, ValueState}
  * @param valEncoder - Spark SQL encoder for value
  * @param ttlConfig  - TTL configuration for values  stored in this state
  * @param batchTimestampMs - current batch processing timestamp.
- * @param avroEnc: Optional Avro encoder and decoder to convert between S and Avro row
  * @param metrics - metrics to be updated as part of stateful processing
  * @tparam S - data type of object that will be stored
  */
@@ -45,12 +44,11 @@ class ValueStateImplWithTTL[S](
     valEncoder: Encoder[S],
     ttlConfig: TTLConfig,
     batchTimestampMs: Long,
-    avroEnc: Option[AvroEncoderSpec],   // TODO: Add Avro Encoding support for TTL
     metrics: Map[String, SQLMetric] = Map.empty)
   extends SingleKeyTTLStateImpl(
     stateName, store, keyExprEnc, batchTimestampMs) with ValueState[S] {
 
-  private val stateTypesEncoder = UnsafeRowTypesEncoder(keyExprEnc, valEncoder,
+  private val stateTypesEncoder = StateTypesEncoder(keyExprEnc, valEncoder,
     stateName, hasTtl = true)
   private val ttlExpirationMs =
     StateTTL.calculateExpirationTimeForDuration(ttlConfig.ttlDuration, batchTimestampMs)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala
index 217a764698416..11bf8ce53b560 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala
@@ -101,23 +101,6 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
     override def valuesIterator(key: UnsafeRow, colFamilyName: String): Iterator[UnsafeRow] = {
       throw StateStoreErrors.unsupportedOperationException("multipleValuesPerKey", "HDFSStateStore")
     }
-
-    override def get(key: Array[Byte], colFamilyName: String): Array[Byte] = {
-      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
-    }
-
-    override def valuesIterator(key: Array[Byte], colFamilyName: String): Iterator[Array[Byte]] = {
-      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
-    }
-
-    override def prefixScan(
-        prefixKey: Array[Byte], colFamilyName: String): Iterator[ByteArrayPair] = {
-      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
-    }
-
-    override def byteArrayIter(colFamilyName: String): Iterator[ByteArrayPair] = {
-      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
-    }
   }
 
   /** Implementation of [[StateStore]] API which is backed by an HDFS-compatible file system */
@@ -144,7 +127,8 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
         valueSchema: StructType,
         keyStateEncoderSpec: KeyStateEncoderSpec,
         useMultipleValuesPerKey: Boolean = false,
-        isInternal: Boolean = false): Unit = {
+        isInternal: Boolean = false,
+        avroEnc: Option[AvroEncoderSpec]): Unit = {
       throw StateStoreErrors.multipleColumnFamiliesNotSupported(providerName)
     }
 
@@ -267,35 +251,6 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
         colFamilyName: String): Unit = {
       throw StateStoreErrors.unsupportedOperationException("merge", providerName)
     }
-
-    override def put(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit = {
-      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
-    }
-
-    override def remove(key: Array[Byte], colFamilyName: String): Unit = {
-      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
-    }
-
-    override def merge(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit = {
-      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
-    }
-
-    override def get(key: Array[Byte], colFamilyName: String): Array[Byte] = {
-      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
-    }
-
-    override def valuesIterator(key: Array[Byte], colFamilyName: String): Iterator[Array[Byte]] = {
-      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
-    }
-
-    override def prefixScan(
-        prefixKey: Array[Byte], colFamilyName: String): Iterator[ByteArrayPair] = {
-      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
-    }
-
-    override def byteArrayIter(colFamilyName: String): Iterator[ByteArrayPair] = {
-      throw StateStoreErrors.unsupportedOperationException("Byte array method", "HDFSStateStore")
-    }
   }
 
   def getMetricsForProvider(): Map[String, Long] = synchronized {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
index bbf0cfcfc7905..3b4fe6a40633c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
@@ -17,11 +17,17 @@
 
 package org.apache.spark.sql.execution.streaming.state
 
+import java.io.ByteArrayOutputStream
 import java.lang.Double.{doubleToRawLongBits, longBitsToDouble}
 import java.lang.Float.{floatToRawIntBits, intBitsToFloat}
 import java.nio.{ByteBuffer, ByteOrder}
 
+import org.apache.avro.generic.{GenericDatumReader, GenericDatumWriter}
+import org.apache.avro.io.{DecoderFactory, EncoderFactory}
+
 import org.apache.spark.internal.Logging
+import org.apache.spark.sql.avro.SchemaConverters
+import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{BoundReference, JoinedRow, UnsafeProjection, UnsafeRow}
 import org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter
 import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider.{STATE_ENCODING_NUM_VERSION_BYTES, STATE_ENCODING_VERSION, VIRTUAL_COL_FAMILY_PREFIX_BYTES}
@@ -33,8 +39,6 @@ sealed trait RocksDBKeyStateEncoder {
   def encodePrefixKey(prefixKey: UnsafeRow): Array[Byte]
   def encodeKey(row: UnsafeRow): Array[Byte]
   def decodeKey(keyBytes: Array[Byte]): UnsafeRow
-  def encodeKeyBytes(row: Array[Byte]): Array[Byte]
-  def decodeKeyBytes(keyBytes: Array[Byte]): Array[Byte]
   def getColumnFamilyIdBytes(): Array[Byte]
 }
 
@@ -43,9 +47,6 @@ sealed trait RocksDBValueStateEncoder {
   def encodeValue(row: UnsafeRow): Array[Byte]
   def decodeValue(valueBytes: Array[Byte]): UnsafeRow
   def decodeValues(valueBytes: Array[Byte]): Iterator[UnsafeRow]
-  def encodeValueBytes(row: Array[Byte]): Array[Byte]
-  def decodeValueBytes(valueBytes: Array[Byte]): Array[Byte]
-  def decodeValuesBytes(valueBytes: Array[Byte]): Iterator[Array[Byte]]
 }
 
 abstract class RocksDBKeyStateEncoderBase(
@@ -54,6 +55,7 @@ abstract class RocksDBKeyStateEncoderBase(
   def offsetForColFamilyPrefix: Int =
     if (useColumnFamilies) VIRTUAL_COL_FAMILY_PREFIX_BYTES else 0
 
+  val out = new ByteArrayOutputStream
   /**
    * Get Byte Array for the virtual column family id that is used as prefix for
    * key state rows.
@@ -98,19 +100,20 @@ object RocksDBStateEncoder {
   def getKeyEncoder(
       keyStateEncoderSpec: KeyStateEncoderSpec,
       useColumnFamilies: Boolean,
-      virtualColFamilyId: Option[Short] = None): RocksDBKeyStateEncoder = {
+      virtualColFamilyId: Option[Short] = None,
+      avroEnc: Option[AvroEncoderSpec] = None): RocksDBKeyStateEncoder = {
     // Return the key state encoder based on the requested type
     keyStateEncoderSpec match {
       case NoPrefixKeyStateEncoderSpec(keySchema) =>
-        new NoPrefixKeyStateEncoder(keySchema, useColumnFamilies, virtualColFamilyId)
+        new NoPrefixKeyStateEncoder(keySchema, useColumnFamilies, virtualColFamilyId, avroEnc)
 
       case PrefixKeyScanStateEncoderSpec(keySchema, numColsPrefixKey) =>
         new PrefixKeyScanStateEncoder(keySchema, numColsPrefixKey,
-          useColumnFamilies, virtualColFamilyId)
+          useColumnFamilies, virtualColFamilyId, avroEnc)
 
       case RangeKeyScanStateEncoderSpec(keySchema, orderingOrdinals) =>
         new RangeKeyScanStateEncoder(keySchema, orderingOrdinals,
-          useColumnFamilies, virtualColFamilyId)
+          useColumnFamilies, virtualColFamilyId, avroEnc)
 
       case _ =>
         throw new IllegalArgumentException(s"Unsupported key state encoder spec: " +
@@ -120,11 +123,12 @@ object RocksDBStateEncoder {
 
   def getValueEncoder(
       valueSchema: StructType,
-      useMultipleValuesPerKey: Boolean): RocksDBValueStateEncoder = {
+      useMultipleValuesPerKey: Boolean,
+      avroEnc: Option[AvroEncoderSpec] = None): RocksDBValueStateEncoder = {
     if (useMultipleValuesPerKey) {
       new MultiValuedStateEncoder(valueSchema)
     } else {
-      new SingleValueStateEncoder(valueSchema)
+      new SingleValueStateEncoder(valueSchema, avroEnc)
     }
   }
 
@@ -171,49 +175,6 @@ object RocksDBStateEncoder {
       null
     }
   }
-
-  /**
-   * Encode a byte array by adding a version byte at the beginning.
-   * Final byte layout: [VersionByte][OriginalBytes]
-   * where:
-   * - VersionByte: Single byte indicating encoding version
-   * - OriginalBytes: The input byte array unchanged
-   *
-   * @param input The original byte array to encode
-   * @return A new byte array containing the version byte followed by the input bytes
-   * @note This creates a new byte array and copies the input array to the new array.
-   */
-  def encodeByteArray(input: Array[Byte]): Array[Byte] = {
-    val encodedBytes = new Array[Byte](input.length + STATE_ENCODING_NUM_VERSION_BYTES)
-    Platform.putByte(encodedBytes, Platform.BYTE_ARRAY_OFFSET, STATE_ENCODING_VERSION)
-    Platform.copyMemory(
-      input, Platform.BYTE_ARRAY_OFFSET,
-      encodedBytes, Platform.BYTE_ARRAY_OFFSET + STATE_ENCODING_NUM_VERSION_BYTES,
-      input.length)
-    encodedBytes
-  }
-
-  /**
-   * Decode bytes by removing the version byte at the beginning.
-   * Input byte layout: [VersionByte][OriginalBytes]
-   * Returns: [OriginalBytes]
-   *
-   * @param bytes The encoded byte array
-   * @return A new byte array containing just the original bytes (excluding version byte),
-   *         or null if input is null
-   */
-  def decodeToByteArray(bytes: Array[Byte]): Array[Byte] = {
-    if (bytes != null) {
-      val decodedBytes = new Array[Byte](bytes.length - STATE_ENCODING_NUM_VERSION_BYTES)
-      Platform.copyMemory(
-        bytes, Platform.BYTE_ARRAY_OFFSET + STATE_ENCODING_NUM_VERSION_BYTES,
-        decodedBytes, Platform.BYTE_ARRAY_OFFSET,
-        decodedBytes.length)
-      decodedBytes
-    } else {
-      null
-    }
-  }
 }
 
 /**
@@ -227,7 +188,8 @@ class PrefixKeyScanStateEncoder(
     keySchema: StructType,
     numColsPrefixKey: Int,
     useColumnFamilies: Boolean = false,
-    virtualColFamilyId: Option[Short] = None)
+    virtualColFamilyId: Option[Short] = None,
+    avroEnc: Option[AvroEncoderSpec] = None)
   extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) {
 
   import RocksDBStateEncoder._
@@ -316,11 +278,6 @@ class PrefixKeyScanStateEncoder(
 
   override def supportPrefixKeyScan: Boolean = true
 
-  override def encodeKeyBytes(row: Array[Byte]): Array[Byte] =
-    throw new UnsupportedOperationException
-
-  override def decodeKeyBytes(keyBytes: Array[Byte]): Array[Byte] =
-    throw new UnsupportedOperationException
 }
 
 /**
@@ -358,7 +315,8 @@ class RangeKeyScanStateEncoder(
     keySchema: StructType,
     orderingOrdinals: Seq[Int],
     useColumnFamilies: Boolean = false,
-    virtualColFamilyId: Option[Short] = None)
+    virtualColFamilyId: Option[Short] = None,
+    avroEnc: Option[AvroEncoderSpec] = None)
   extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) {
 
   import RocksDBStateEncoder._
@@ -698,12 +656,6 @@ class RangeKeyScanStateEncoder(
   }
 
   override def supportPrefixKeyScan: Boolean = true
-
-  override def encodeKeyBytes(row: Array[Byte]): Array[Byte] =
-    throw new UnsupportedOperationException
-
-  override def decodeKeyBytes(keyBytes: Array[Byte]): Array[Byte] =
-    throw new UnsupportedOperationException
 }
 
 /**
@@ -721,19 +673,31 @@ class RangeKeyScanStateEncoder(
 class NoPrefixKeyStateEncoder(
     keySchema: StructType,
     useColumnFamilies: Boolean = false,
-    virtualColFamilyId: Option[Short] = None)
-  extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) {
+    virtualColFamilyId: Option[Short] = None,
+    avroEnc: Option[AvroEncoderSpec] = None)
+  extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) with Logging {
 
   import RocksDBStateEncoder._
 
   // Reusable objects
   private val keyRow = new UnsafeRow(keySchema.size)
+  private val keyAvroType = SchemaConverters.toAvroType(keySchema)
 
   override def encodeKey(row: UnsafeRow): Array[Byte] = {
     if (!useColumnFamilies) {
       encodeUnsafeRow(row)
     } else {
-      val bytesToEncode = row.getBytes
+      // If avroEnc is defined, we know that we need to use Avro to
+      // encode this UnsafeRow to Avro bytes
+      val bytesToEncode = if (avroEnc.isDefined) {
+        val avroData = avroEnc.get.keySerializer.serialize(row)
+        out.reset()
+        val encoder = EncoderFactory.get().directBinaryEncoder(out, null)
+        val writer = new GenericDatumWriter[Any](keyAvroType)
+        writer.write(avroData, encoder)
+        encoder.flush()
+        out.toByteArray
+      } else row.getBytes
       val (encodedBytes, startingOffset) = encodeColumnFamilyPrefix(
         bytesToEncode.length +
           STATE_ENCODING_NUM_VERSION_BYTES
@@ -768,69 +732,6 @@ class NoPrefixKeyStateEncoder(
     } else decodeToUnsafeRow(keyBytes, keyRow)
   }
 
-  /**
-   * Encodes a byte array by adding column family prefix and version information.
-   * Final byte layout: [ColFamilyPrefix][VersionByte][OriginalBytes]
-   * where:
-   * - ColFamilyPrefix: Optional prefix identifying the column family (if useColumnFamilies=true)
-   * - VersionByte: Single byte indicating encoding version
-   * - OriginalBytes: The input byte array unchanged
-   *
-   * @param row The original byte array to encode
-   * @return The encoded byte array with prefix and version if column families are enabled,
-   *         otherwise returns the original array
-   */
-  override def encodeKeyBytes(row: Array[Byte]): Array[Byte] = {
-    if (!useColumnFamilies) {
-      row
-    } else {
-      // Calculate total size needed: original bytes + 1 byte for version + column family prefix
-      val (encodedBytes, startingOffset) = encodeColumnFamilyPrefix(
-        row.length +
-          STATE_ENCODING_NUM_VERSION_BYTES
-      )
-
-      // Add version byte right after column family prefix
-      Platform.putByte(encodedBytes, startingOffset, STATE_ENCODING_VERSION)
-
-      // Copy original bytes after the version byte
-      // Platform.BYTE_ARRAY_OFFSET is the recommended way to memcopy b/w byte arrays. See Platform.
-      Platform.copyMemory(
-        row, Platform.BYTE_ARRAY_OFFSET,
-        encodedBytes, startingOffset + STATE_ENCODING_NUM_VERSION_BYTES, row.length)
-      encodedBytes
-    }
-  }
-
-  /**
-   * Decodes a byte array by removing column family prefix and version information.
-   * Input byte layout: [ColFamilyPrefix][VersionByte][OriginalBytes]
-   * Returns: [OriginalBytes]
-   *
-   * @param keyBytes The encoded byte array to decode
-   * @return The original byte array with prefix and version removed if column families are enabled,
-   *         null if input is null, or a clone of input if column families are disabled
-   */
-  override def decodeKeyBytes(keyBytes: Array[Byte]): Array[Byte] = {
-    if (keyBytes == null) {
-      null
-    } else if (useColumnFamilies) {
-      // Calculate start offset (skip column family prefix and version byte)
-      val startOffset = decodeKeyStartOffset + STATE_ENCODING_NUM_VERSION_BYTES
-
-      // Calculate length of original data (total length minus prefix and version)
-      val length = keyBytes.length -
-        STATE_ENCODING_NUM_VERSION_BYTES - VIRTUAL_COL_FAMILY_PREFIX_BYTES
-
-      // Extract just the original bytes
-      java.util.Arrays.copyOfRange(keyBytes, startOffset, startOffset + length)
-    } else {
-      // If column families not enabled, just return a copy of the input
-      // Assuming decodeToUnsafeRow is not applicable for byte array encoding
-      keyBytes.clone()
-    }
-  }
-
   override def supportPrefixKeyScan: Boolean = false
 
   override def encodePrefixKey(prefixKey: UnsafeRow): Array[Byte] = {
@@ -912,95 +813,6 @@ class MultiValuedStateEncoder(valueSchema: StructType)
   }
 
   override def supportsMultipleValuesPerKey: Boolean = true
-
-  /**
-   * Encodes a raw byte array value in the multi-value format.
-   * Format: [length (4 bytes)][value bytes][delimiter (1 byte)]
-   */
-  override def encodeValueBytes(row: Array[Byte]): Array[Byte] = {
-    if (row == null) {
-      null
-    } else {
-      val numBytes = row.length
-      // Allocate space for:
-      // - 4 bytes for length
-      // - The actual value bytes
-      val encodedBytes = new Array[Byte](java.lang.Integer.BYTES + numBytes)
-
-      // Write length as big-endian int
-      Platform.putInt(encodedBytes, Platform.BYTE_ARRAY_OFFSET, numBytes)
-
-      // Copy value bytes after the length
-      Platform.copyMemory(
-        row, Platform.BYTE_ARRAY_OFFSET,
-        encodedBytes, Platform.BYTE_ARRAY_OFFSET + java.lang.Integer.BYTES,
-        numBytes
-      )
-
-      encodedBytes
-    }
-  }
-
-  /**
-   * Decodes a single value from the encoded byte format.
-   * Assumes the bytes represent a single value, not multiple merged values.
-   */
-  override def decodeValueBytes(valueBytes: Array[Byte]): Array[Byte] = {
-    if (valueBytes == null) {
-      null
-    } else {
-      // Read length from first 4 bytes
-      val numBytes = Platform.getInt(valueBytes, Platform.BYTE_ARRAY_OFFSET)
-
-      // Extract just the value bytes after the length
-      val decodedBytes = new Array[Byte](numBytes)
-      Platform.copyMemory(
-        valueBytes, Platform.BYTE_ARRAY_OFFSET + java.lang.Integer.BYTES,
-        decodedBytes, Platform.BYTE_ARRAY_OFFSET,
-        numBytes
-      )
-
-      decodedBytes
-    }
-  }
-
-  /**
-   * Decodes multiple values from the merged byte format.
-   * Returns an iterator that lazily decodes each value.
-   */
-  override def decodeValuesBytes(valueBytes: Array[Byte]): Iterator[Array[Byte]] = {
-    if (valueBytes == null) {
-      Iterator.empty
-    } else {
-      new Iterator[Array[Byte]] {
-        // Track current position in the byte array
-        private var pos: Int = Platform.BYTE_ARRAY_OFFSET
-        private val maxPos = Platform.BYTE_ARRAY_OFFSET + valueBytes.length
-
-        override def hasNext: Boolean = pos < maxPos
-
-        override def next(): Array[Byte] = {
-          // Read length prefix
-          val numBytes = Platform.getInt(valueBytes, pos)
-          pos += java.lang.Integer.BYTES
-
-          // Extract value bytes
-          val decodedValue = new Array[Byte](numBytes)
-          Platform.copyMemory(
-            valueBytes, pos,
-            decodedValue, Platform.BYTE_ARRAY_OFFSET,
-            numBytes
-          )
-
-          // Move position past value and delimiter
-          pos += numBytes
-          pos += 1 // Skip delimiter byte
-
-          decodedValue
-        }
-      }
-    }
-  }
 }
 
 /**
@@ -1015,15 +827,34 @@ class MultiValuedStateEncoder(valueSchema: StructType)
  *    (offset 0 is the version byte of value 0). That is, if the unsafe row has N bytes,
  *    then the generated array byte will be N+1 bytes.
  */
-class SingleValueStateEncoder(valueSchema: StructType)
-  extends RocksDBValueStateEncoder {
+class SingleValueStateEncoder(
+    valueSchema: StructType,
+    avroEnc: Option[AvroEncoderSpec] = None)
+  extends RocksDBValueStateEncoder with Logging {
 
   import RocksDBStateEncoder._
 
   // Reusable objects
+  private val out = new ByteArrayOutputStream
   private val valueRow = new UnsafeRow(valueSchema.size)
+  private val valueAvroType = SchemaConverters.toAvroType(valueSchema)
+  private val valueProj = UnsafeProjection.create(valueSchema)
 
-  override def encodeValue(row: UnsafeRow): Array[Byte] = encodeUnsafeRow(row)
+  override def encodeValue(row: UnsafeRow): Array[Byte] = {
+    if (avroEnc.isDefined) {
+      val avroData =
+        avroEnc.get.valueSerializer.serialize(row) // InternalRow -> GenericDataRecord
+      out.reset()
+      val encoder = EncoderFactory.get().directBinaryEncoder(out, null)
+      val writer = new GenericDatumWriter[Any](
+        valueAvroType) // Defining Avro writer for this struct type
+      writer.write(avroData, encoder) // GenericDataRecord -> bytes
+      encoder.flush()
+      out.toByteArray
+    } else {
+      encodeUnsafeRow(row)
+    }
+  }
 
   /**
    * Decode byte array for a value to a UnsafeRow.
@@ -1032,7 +863,19 @@ class SingleValueStateEncoder(valueSchema: StructType)
    *       the given byte array.
    */
   override def decodeValue(valueBytes: Array[Byte]): UnsafeRow = {
-    decodeToUnsafeRow(valueBytes, valueRow)
+    if (valueBytes == null) {
+      return null
+    }
+    if (avroEnc.isDefined) {
+      val reader = new GenericDatumReader[Any](valueAvroType)
+      val decoder = DecoderFactory.get().binaryDecoder(valueBytes, 0, valueBytes.length, null)
+      val genericData = reader.read(null, decoder) // bytes -> GenericDataRecord
+      val internalRow = avroEnc.get.valueDeserializer.deserialize(
+        genericData).orNull.asInstanceOf[InternalRow]
+      valueProj.apply(internalRow)
+    } else {
+      decodeToUnsafeRow(valueBytes, valueRow)
+    }
   }
 
   override def supportsMultipleValuesPerKey: Boolean = false
@@ -1040,16 +883,4 @@ class SingleValueStateEncoder(valueSchema: StructType)
   override def decodeValues(valueBytes: Array[Byte]): Iterator[UnsafeRow] = {
     throw new IllegalStateException("This encoder doesn't support multiple values!")
   }
-
-  override def encodeValueBytes(row: Array[Byte]): Array[Byte] = {
-    encodeByteArray(row)
-  }
-
-  override def decodeValueBytes(valueBytes: Array[Byte]): Array[Byte] = {
-    decodeToByteArray(valueBytes)
-  }
-
-  override def decodeValuesBytes(valueBytes: Array[Byte]): Iterator[Array[Byte]] = {
-    throw new IllegalStateException("This encoder doesn't support multiple values!")
-  }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala
index 2bcccc1a2d310..0ab10a6fbdb98 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala
@@ -71,13 +71,14 @@ private[sql] class RocksDBStateStoreProvider
         valueSchema: StructType,
         keyStateEncoderSpec: KeyStateEncoderSpec,
         useMultipleValuesPerKey: Boolean = false,
-        isInternal: Boolean = false): Unit = {
+        isInternal: Boolean = false,
+        avroEnc: Option[AvroEncoderSpec]): Unit = {
       verifyColFamilyCreationOrDeletion("create_col_family", colFamilyName, isInternal)
       val newColFamilyId = rocksDB.createColFamilyIfAbsent(colFamilyName)
       keyValueEncoderMap.putIfAbsent(colFamilyName,
         (RocksDBStateEncoder.getKeyEncoder(keyStateEncoderSpec, useColumnFamilies,
-          Some(newColFamilyId)), RocksDBStateEncoder.getValueEncoder(valueSchema,
-          useMultipleValuesPerKey)))
+          Some(newColFamilyId), avroEnc), RocksDBStateEncoder.getValueEncoder(valueSchema,
+          useMultipleValuesPerKey, avroEnc)))
     }
 
     override def get(key: UnsafeRow, colFamilyName: String): UnsafeRow = {
@@ -209,98 +210,6 @@ private[sql] class RocksDBStateStoreProvider
       }
     }
 
-    override def put(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit = {
-      verify(state == UPDATING, "Cannot put after already committed or aborted")
-      verify(key != null, "Key cannot be null")
-      require(value != null, "Cannot put a null value")
-      verifyColFamilyOperations("put", colFamilyName)
-
-      val kvEncoder = keyValueEncoderMap.get(colFamilyName)
-      rocksDB.put(kvEncoder._1.encodeKeyBytes(key), kvEncoder._2.encodeValueBytes(value))
-    }
-
-    override def remove(key: Array[Byte], colFamilyName: String): Unit = {
-      verify(state == UPDATING, "Cannot remove after already committed or aborted")
-      verify(key != null, "Key cannot be null")
-      verifyColFamilyOperations("remove", colFamilyName)
-
-      val kvEncoder = keyValueEncoderMap.get(colFamilyName)
-      rocksDB.remove(kvEncoder._1.encodeKeyBytes(key))
-    }
-
-    override def get(key: Array[Byte], colFamilyName: String): Array[Byte] = {
-      verify(key != null, "Key cannot be null")
-      verifyColFamilyOperations("get", colFamilyName)
-
-      val kvEncoder = keyValueEncoderMap.get(colFamilyName)
-      kvEncoder._2.decodeValueBytes(rocksDB.get(kvEncoder._1.encodeKeyBytes(key)))
-    }
-
-    override def byteArrayIter(colFamilyName: String): Iterator[ByteArrayPair] = {
-      // Verify column family operation is valid
-      verifyColFamilyOperations("byteArrayIter", colFamilyName)
-      val kvEncoder = keyValueEncoderMap.get(colFamilyName)
-      val pair = new ByteArrayPair()
-
-      // Similar to the regular iterator, we need to handle both column family
-      // and non-column family cases
-      if (useColumnFamilies) {
-        rocksDB.prefixScan(kvEncoder._1.getColumnFamilyIdBytes()).map { kv =>
-          pair.set(
-            kvEncoder._1.decodeKeyBytes(kv.key),
-            kvEncoder._2.decodeValueBytes(kv.value))
-          pair
-        }
-      } else {
-        rocksDB.iterator().map { kv =>
-          pair.set(
-            kvEncoder._1.decodeKeyBytes(kv.key),
-            kvEncoder._2.decodeValueBytes(kv.value))
-          pair
-        }
-      }
-    }
-
-    override def valuesIterator(key: Array[Byte], colFamilyName: String): Iterator[Array[Byte]] = {
-      verify(key != null, "Key cannot be null")
-      verifyColFamilyOperations("valuesIterator", colFamilyName)
-
-      val kvEncoder = keyValueEncoderMap.get(colFamilyName)
-      val valueEncoder = kvEncoder._2
-      val keyEncoder = kvEncoder._1
-
-      verify(valueEncoder.supportsMultipleValuesPerKey,
-        "valuesIterator requires an encoder that supports multiple values for a single key.")
-
-      // Get the encoded value bytes using the encoded key
-      val encodedValues = rocksDB.get(keyEncoder.encodeKeyBytes(key))
-
-      // Decode multiple values from the merged value bytes
-      valueEncoder.decodeValuesBytes(encodedValues)
-    }
-
-    override def prefixScan(
-        prefixKey: Array[Byte],
-        colFamilyName: String): Iterator[ByteArrayPair] = {
-      throw StateStoreErrors.unsupportedOperationException(
-        "bytearray prefixScan", "RocksDBStateStore")
-    }
-
-    override def merge(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit = {
-      verify(state == UPDATING, "Cannot merge after already committed or aborted")
-      verifyColFamilyOperations("merge", colFamilyName)
-
-      val kvEncoder = keyValueEncoderMap.get(colFamilyName)
-      val keyEncoder = kvEncoder._1
-      val valueEncoder = kvEncoder._2
-      verify(valueEncoder.supportsMultipleValuesPerKey, "Merge operation requires an encoder" +
-        " which supports multiple values for a single key")
-      verify(key != null, "Key cannot be null")
-      require(value != null, "Cannot merge a null value")
-
-      rocksDB.merge(keyEncoder.encodeKeyBytes(key), valueEncoder.encodeValueBytes(value))
-    }
-
     override def commit(): Long = synchronized {
       try {
         verify(state == UPDATING, "Cannot commit after already committed or aborted")
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
index b02d36c8ced85..d8094f78f587f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
@@ -40,6 +40,7 @@ case class StateSchemaValidationResult(
 
 case class AvroEncoderSpec(
   keySerializer: AvroSerializer,
+  keyDeserializer: AvroDeserializer,
   valueSerializer: AvroSerializer,
   valueDeserializer: AvroDeserializer
 ) extends Serializable
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
index f0b6fdd41ba18..d5f79f27c7ff7 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
@@ -106,44 +106,6 @@ trait ReadStateStore {
   /** Return an iterator containing all the key-value pairs in the StateStore. */
   def iterator(colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Iterator[UnsafeRowPair]
 
-  /**
-   * Get the current value of a non-null key.
-   * @return a non-null row if the key exists in the store, otherwise null.
-   */
-  def get(
-      key: Array[Byte],
-      colFamilyName: String): Array[Byte]
-
-  /**
-   * Provides an iterator containing all values of a non-null key. If key does not exist,
-   * an empty iterator is returned. Implementations should make sure to return an empty
-   * iterator if the key does not exist.
-   *
-   * It is expected to throw exception if Spark calls this method without setting
-   * multipleValuesPerKey as true for the column family.
-   */
-  def valuesIterator(
-      key: Array[Byte],
-      colFamilyName: String): Iterator[Array[Byte]]
-
-  /**
-   * Return an iterator containing all the key-value pairs which are matched with
-   * the given prefix key.
-   *
-   * The operator will provide numColsPrefixKey greater than 0 in StateStoreProvider.init method
-   * if the operator needs to leverage the "prefix scan" feature. The schema of the prefix key
-   * should be same with the leftmost `numColsPrefixKey` columns of the key schema.
-   *
-   * It is expected to throw exception if Spark calls this method without setting numColsPrefixKey
-   * to the greater than 0.
-   */
-  def prefixScan(
-      prefixKey: Array[Byte],
-      colFamilyName: String): Iterator[ByteArrayPair]
-
-  /** Return an iterator containing all the key-value pairs in the StateStore. */
-  def byteArrayIter(colFamilyName: String): Iterator[ByteArrayPair]
-
   /**
    * Clean up the resource.
    *
@@ -180,7 +142,8 @@ trait StateStore extends ReadStateStore {
       valueSchema: StructType,
       keyStateEncoderSpec: KeyStateEncoderSpec,
       useMultipleValuesPerKey: Boolean = false,
-      isInternal: Boolean = false): Unit
+      isInternal: Boolean = false,
+      avroEncoderSpec: Option[AvroEncoderSpec] = None): Unit
 
   /**
    * Put a new non-null value for a non-null key. Implementations must be aware that the UnsafeRows
@@ -208,19 +171,6 @@ trait StateStore extends ReadStateStore {
   def merge(key: UnsafeRow, value: UnsafeRow,
       colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Unit
 
-  /**
-   * Put a new non-null value for a non-null key. Implementations must be aware that the UnsafeRows
-   * in the params can be reused, and must make copies of the data as needed for persistence.
-   */
-  def put(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit
-
-  /**
-   * Remove a single non-null key.
-   */
-  def remove(key: Array[Byte], colFamilyName: String): Unit
-
-  def merge(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit
-
   /**
    * Commit all the updates that have been made to the store, and return the new version.
    * Implementations should ensure that no more updates (puts, removes) can be after a commit in
@@ -287,19 +237,6 @@ class WrappedReadStateStore(store: StateStore) extends ReadStateStore {
   override def valuesIterator(key: UnsafeRow, colFamilyName: String): Iterator[UnsafeRow] = {
     store.valuesIterator(key, colFamilyName)
   }
-
-  override def get(key: Array[Byte], colFamilyName: String): Array[Byte] =
-    store.get(key, colFamilyName)
-
-  override def valuesIterator(key: Array[Byte], colFamilyName: String):
-    Iterator[Array[Byte]] = store.valuesIterator(key, colFamilyName)
-
-
-  override def prefixScan(prefixKey: Array[Byte], colFamilyName: String):
-    Iterator[ByteArrayPair] = store.prefixScan(prefixKey, colFamilyName)
-
-  override def byteArrayIter(colFamilyName: String): Iterator[ByteArrayPair] =
-    store.byteArrayIter(colFamilyName)
 }
 
 /**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
index 4c07d41c62f78..0cf641c703d6a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
@@ -29,7 +29,7 @@ import org.apache.hadoop.fs.Path
 import org.apache.spark.SparkContext
 import org.apache.spark.internal.Logging
 import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.{AnalysisException, SparkSession}
+import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeProjection
@@ -97,14 +97,6 @@ object StatefulOperatorStateInfo {
   }
 }
 
-object StatefulOperatorUtils {
-  @transient final val session = SparkSession.getActiveSession.orNull
-
-  lazy val stateStoreEncoding: String =
-    session.sessionState.conf.getConf(
-      SQLConf.STREAMING_STATE_STORE_ENCODING_FORMAT)
-}
-
 /**
  * An operator that reads or writes state from the [[StateStore]].
  * The [[StatefulOperatorStateInfo]] should be filled in by `prepareForExecution` in
@@ -119,6 +111,10 @@ trait StatefulOperator extends SparkPlan {
     }
   }
 
+  lazy val stateStoreEncoding: String =
+    session.sessionState.conf.getConf(
+      SQLConf.STREAMING_STATE_STORE_ENCODING_FORMAT)
+
   def metadataFilePath(): Path = {
     val stateCheckpointPath =
       new Path(getStateInfo.checkpointLocation, getStateInfo.operatorId.toString)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ListStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ListStateSuite.scala
index cf748da3881ec..e9300464af8dc 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ListStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ListStateSuite.scala
@@ -34,15 +34,13 @@ class ListStateSuite extends StateVariableSuiteBase {
   // overwrite useMultipleValuesPerKey in base suite to be true for list state
   override def useMultipleValuesPerKey: Boolean = true
 
-  private def testMapStateWithNullUserKey(useAvro: Boolean)
-      (runListOps: ListState[Long] => Unit): Unit = {
+  private def testMapStateWithNullUserKey()(runListOps: ListState[Long] => Unit): Unit = {
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val listState: ListState[Long] = handle.getListStateWithAvro[Long](
-        "listState", Encoders.scalaLong, useAvro)
+      val listState: ListState[Long] = handle.getListState[Long]("listState", Encoders.scalaLong)
 
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
       val e = intercept[SparkIllegalArgumentException] {
@@ -59,8 +57,8 @@ class ListStateSuite extends StateVariableSuiteBase {
   }
 
   Seq("appendList", "put").foreach { listImplFunc =>
-    testWithAvroEnc(s"Test list operation($listImplFunc) with null") { useAvro =>
-      testMapStateWithNullUserKey(useAvro) { listState =>
+    test(s"Test list operation($listImplFunc) with null") {
+      testMapStateWithNullUserKey() { listState =>
         listImplFunc match {
           case "appendList" => listState.appendList(null)
           case "put" => listState.put(null)
@@ -69,14 +67,13 @@ class ListStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  testWithAvroEnc("List state operations for single instance") { useAvro =>
+  test("List state operations for single instance") {
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val testState: ListState[Long] = handle.getListStateWithAvro[Long](
-        "testState", Encoders.scalaLong, useAvro)
+      val testState: ListState[Long] = handle.getListState[Long]("testState", Encoders.scalaLong)
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
 
       // simple put and get test
@@ -98,16 +95,14 @@ class ListStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  testWithAvroEnc("List state operations for multiple instance") { useAvro =>
+  test("List state operations for multiple instance") {
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val testState1: ListState[Long] = handle.getListStateWithAvro[Long](
-        "testState1", Encoders.scalaLong, useAvro)
-      val testState2: ListState[Long] = handle.getListStateWithAvro[Long](
-        "testState2", Encoders.scalaLong, useAvro)
+      val testState1: ListState[Long] = handle.getListState[Long]("testState1", Encoders.scalaLong)
+      val testState2: ListState[Long] = handle.getListState[Long]("testState2", Encoders.scalaLong)
 
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
 
@@ -138,18 +133,16 @@ class ListStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  testWithAvroEnc("List state operations with list, value, another list instances") { useAvro =>
+  test("List state operations with list, value, another list instances") {
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val listState1: ListState[Long] = handle.getListStateWithAvro[Long](
-        "listState1", Encoders.scalaLong, useAvro)
-      val listState2: ListState[Long] = handle.getListStateWithAvro[Long](
-        "listState2", Encoders.scalaLong, useAvro)
-      val valueState: ValueState[Long] = handle.getValueStateWithAvro[Long](
-        "valueState", Encoders.scalaLong, useAvro = false)
+      val listState1: ListState[Long] = handle.getListState[Long]("listState1", Encoders.scalaLong)
+      val listState2: ListState[Long] = handle.getListState[Long]("listState2", Encoders.scalaLong)
+      val valueState: ValueState[Long] = handle.getValueState[Long](
+        "valueState", Encoders.scalaLong)
 
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
       // simple put and get test
@@ -252,7 +245,7 @@ class ListStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  testWithAvroEnc("ListState TTL with non-primitive types") { useAvro =>
+  test("ListState TTL with non-primitive types") {
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val timestampMs = 10
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MapStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MapStateSuite.scala
index d9f29ac69ceae..b067d589de904 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MapStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MapStateSuite.scala
@@ -120,9 +120,9 @@ class MapStateSuite extends StateVariableSuiteBase {
       val mapTestState2: MapState[String, Int] =
         handle.getMapState[String, Int]("mapTestState2", Encoders.STRING, Encoders.scalaInt)
       val valueTestState: ValueState[String] =
-        handle.getValueStateWithAvro[String]("valueTestState", Encoders.STRING, false)
+        handle.getValueState[String]("valueTestState", Encoders.STRING)
       val listTestState: ListState[String] =
-        handle.getListStateWithAvro[String]("listTestState", Encoders.STRING, false)
+        handle.getListState[String]("listTestState", Encoders.STRING)
 
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
       // put initial values
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MemoryStateStore.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MemoryStateStore.scala
index a8859673466ec..bfcd828d01cc3 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MemoryStateStore.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MemoryStateStore.scala
@@ -36,7 +36,8 @@ class MemoryStateStore extends StateStore() {
       valueSchema: StructType,
       keyStateEncoderSpec: KeyStateEncoderSpec,
       useMultipleValuesPerKey: Boolean = false,
-      isInternal: Boolean = false): Unit = {
+      isInternal: Boolean = false,
+      avroEnc: Option[AvroEncoderSpec]): Unit = {
     throw StateStoreErrors.multipleColumnFamiliesNotSupported("MemoryStateStoreProvider")
   }
 
@@ -78,34 +79,4 @@ class MemoryStateStore extends StateStore() {
   override def getStateStoreCheckpointInfo(): StateStoreCheckpointInfo = {
     StateStoreCheckpointInfo(id.partitionId, version + 1, None, None)
   }
-
-
-  override def put(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit = {
-    throw new UnsupportedOperationException("Doesn't support bytearray operations")
-  }
-
-  override def remove(key: Array[Byte], colFamilyName: String): Unit = {
-    throw new UnsupportedOperationException("Doesn't support bytearray operations")
-  }
-
-  override def merge(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit = {
-    throw new UnsupportedOperationException("Doesn't support bytearray operations")
-  }
-
-  override def get(key: Array[Byte], colFamilyName: String): Array[Byte] = {
-    throw new UnsupportedOperationException("Doesn't support bytearray operations")
-  }
-
-  override def valuesIterator(key: Array[Byte], colFamilyName: String): Iterator[Array[Byte]] = {
-    throw new UnsupportedOperationException("Doesn't support bytearray operations")
-  }
-
-  override def prefixScan(
-      prefixKey: Array[Byte], colFamilyName: String): Iterator[ByteArrayPair] = {
-    throw new UnsupportedOperationException("Doesn't support bytearray operations")
-  }
-
-  override def byteArrayIter(colFamilyName: String): Iterator[ByteArrayPair] = {
-    throw new UnsupportedOperationException("Doesn't support bytearray operations")
-  }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala
index 9bf5b6b73ff6c..e6454d3c77a2f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala
@@ -91,7 +91,8 @@ case class CkptIdCollectingStateStoreWrapper(innerStore: StateStore) extends Sta
       valueSchema: StructType,
       keyStateEncoderSpec: KeyStateEncoderSpec,
       useMultipleValuesPerKey: Boolean = false,
-      isInternal: Boolean = false): Unit = {
+      isInternal: Boolean = false,
+      avroEnc: Option[AvroEncoderSpec]): Unit = {
     innerStore.createColFamilyIfAbsent(
       colFamilyName,
       keySchema,
@@ -122,35 +123,6 @@ case class CkptIdCollectingStateStoreWrapper(innerStore: StateStore) extends Sta
     innerStore.merge(key, value, colFamilyName)
   }
 
-  override def put(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit = {
-    throw new UnsupportedOperationException
-  }
-
-  override def remove(key: Array[Byte], colFamilyName: String): Unit = {
-    throw new UnsupportedOperationException
-  }
-
-  override def get(key: Array[Byte], colFamilyName: String): Array[Byte] = {
-    throw new UnsupportedOperationException
-  }
-
-  override def valuesIterator(key: Array[Byte], colFamilyName: String): Iterator[Array[Byte]] = {
-    throw new UnsupportedOperationException
-  }
-
-  override def prefixScan(
-      prefixKey: Array[Byte], colFamilyName: String): Iterator[ByteArrayPair] = {
-    throw new UnsupportedOperationException
-  }
-
-  override def byteArrayIter(colFamilyName: String): Iterator[ByteArrayPair] = {
-    throw new UnsupportedOperationException
-  }
-
-  override def merge(key: Array[Byte], value: Array[Byte], colFamilyName: String): Unit = {
-    throw new UnsupportedOperationException
-  }
-
   override def commit(): Long = innerStore.commit()
   override def metrics: StateStoreMetrics = innerStore.metrics
   override def getStateStoreCheckpointInfo(): StateStoreCheckpointInfo = {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatefulProcessorHandleSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatefulProcessorHandleSuite.scala
index 0c14b8a8601c5..48a6fd836a462 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatefulProcessorHandleSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatefulProcessorHandleSuite.scala
@@ -226,7 +226,7 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase {
         Encoders.STRING, TTLConfig(Duration.ofHours(1)))
 
       // create another state without TTL, this should not be captured in the handle
-      handle.getValueStateWithAvro("testState", Encoders.STRING, useAvro = false)
+      handle.getValueState("testState", Encoders.STRING)
 
       assert(handle.ttlStates.size() === 1)
       assert(handle.ttlStates.get(0) === valueStateWithTTL)
@@ -275,7 +275,7 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase {
       val handle = new StatefulProcessorHandleImpl(store,
         UUID.randomUUID(), stringEncoder, TimeMode.None())
 
-      handle.getValueStateWithAvro("testValueState", Encoders.STRING, useAvro = false)
+      handle.getValueState("testValueState", Encoders.STRING)
       handle.getListState("testListState", Encoders.STRING)
       handle.getMapState("testMapState", Encoders.STRING, Encoders.STRING)
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ValueStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ValueStateSuite.scala
index cf13495f5f616..13d758eb1b88f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ValueStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ValueStateSuite.scala
@@ -23,7 +23,7 @@ import java.util.UUID
 import scala.util.Random
 
 import org.apache.hadoop.conf.Configuration
-import org.scalatest.{BeforeAndAfter, Tag}
+import org.scalatest.BeforeAndAfter
 
 import org.apache.spark.{SparkException, SparkUnsupportedOperationException}
 import org.apache.spark.sql.Encoders
@@ -45,15 +45,14 @@ class ValueStateSuite extends StateVariableSuiteBase {
 
   import StateStoreTestsHelper._
 
-  testWithAvroEnc("Implicit key operations") { useAvro =>
+  test("Implicit key operations") {
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
       val stateName = "testState"
-      val testState: ValueState[Long] = handle.getValueStateWithAvro[Long](
-        "testState", Encoders.scalaLong, useAvro)
+      val testState: ValueState[Long] = handle.getValueState[Long]("testState", Encoders.scalaLong)
       assert(ImplicitGroupingKeyTracker.getImplicitKeyOption.isEmpty)
       val ex = intercept[Exception] {
         testState.update(123)
@@ -90,14 +89,13 @@ class ValueStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  testWithAvroEnc("Value state operations for single instance") { useAvro =>
+  test("Value state operations for single instance") {
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val testState: ValueState[Long] = handle.getValueStateWithAvro[Long](
-        "testState", Encoders.scalaLong, useAvro)
+      val testState: ValueState[Long] = handle.getValueState[Long]("testState", Encoders.scalaLong)
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
       testState.update(123)
       assert(testState.get() === 123)
@@ -117,16 +115,16 @@ class ValueStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  testWithAvroEnc("Value state operations for multiple instances") { useAvro =>
+  test("Value state operations for multiple instances") {
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val testState1: ValueState[Long] = handle.getValueStateWithAvro[Long](
-        "testState1", Encoders.scalaLong, useAvro)
-      val testState2: ValueState[Long] = handle.getValueStateWithAvro[Long](
-        "testState2", Encoders.scalaLong, useAvro)
+      val testState1: ValueState[Long] = handle.getValueState[Long](
+        "testState1", Encoders.scalaLong)
+      val testState2: ValueState[Long] = handle.getValueState[Long](
+        "testState2", Encoders.scalaLong)
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
       testState1.update(123)
       assert(testState1.get() === 123)
@@ -162,7 +160,7 @@ class ValueStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  testWithAvroEnc("Value state operations for unsupported type name should fail") { useAvro =>
+  test("Value state operations for unsupported type name should fail") {
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store,
@@ -170,7 +168,7 @@ class ValueStateSuite extends StateVariableSuiteBase {
 
       val cfName = "$testState"
       val ex = intercept[SparkUnsupportedOperationException] {
-        handle.getValueStateWithAvro[Long](cfName, Encoders.scalaLong, useAvro)
+        handle.getValueState[Long](cfName, Encoders.scalaLong)
       }
       checkError(
         ex,
@@ -202,15 +200,14 @@ class ValueStateSuite extends StateVariableSuiteBase {
     )
   }
 
-  testWithAvroEnc("test SQL encoder - Value state operations" +
-    " for Primitive(Double) instances") { useAvro =>
+  test("test SQL encoder - Value state operations for Primitive(Double) instances") {
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val testState: ValueState[Double] = handle.getValueStateWithAvro[Double]("testState",
-        Encoders.scalaDouble, useAvro)
+      val testState: ValueState[Double] = handle.getValueState[Double]("testState",
+        Encoders.scalaDouble)
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
       testState.update(1.0)
       assert(testState.get().equals(1.0))
@@ -229,15 +226,14 @@ class ValueStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  testWithAvroEnc("test SQL encoder - Value state operations" +
-    " for Primitive(Long) instances") { useAvro =>
+  test("test SQL encoder - Value state operations for Primitive(Long) instances") {
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val testState: ValueState[Long] = handle.getValueStateWithAvro[Long]("testState",
-        Encoders.scalaLong, useAvro)
+      val testState: ValueState[Long] = handle.getValueState[Long]("testState",
+        Encoders.scalaLong)
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
       testState.update(1L)
       assert(testState.get().equals(1L))
@@ -256,15 +252,14 @@ class ValueStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  testWithAvroEnc("test SQL encoder - Value state operations" +
-    " for case class instances") { useAvro =>
+  test("test SQL encoder - Value state operations for case class instances") {
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val testState: ValueState[TestClass] = handle.getValueStateWithAvro[TestClass]("testState",
-        Encoders.product[TestClass], useAvro)
+      val testState: ValueState[TestClass] = handle.getValueState[TestClass]("testState",
+        Encoders.product[TestClass])
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
       testState.update(TestClass(1, "testcase1"))
       assert(testState.get().equals(TestClass(1, "testcase1")))
@@ -283,14 +278,14 @@ class ValueStateSuite extends StateVariableSuiteBase {
     }
   }
 
-  testWithAvroEnc("test SQL encoder - Value state operations for POJO instances") { useAvro =>
+  test("test SQL encoder - Value state operations for POJO instances") {
     tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
       val store = provider.getStore(0)
       val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
         stringEncoder, TimeMode.None())
 
-      val testState: ValueState[POJOTestClass] = handle.getValueStateWithAvro[POJOTestClass](
-        "testState", Encoders.bean(classOf[POJOTestClass]), useAvro)
+      val testState: ValueState[POJOTestClass] = handle.getValueState[POJOTestClass]("testState",
+        Encoders.bean(classOf[POJOTestClass]))
       ImplicitGroupingKeyTracker.setImplicitKey("test_key")
       testState.update(new POJOTestClass("testcase1", 1))
       assert(testState.get().equals(new POJOTestClass("testcase1", 1)))
@@ -479,26 +474,5 @@ abstract class StateVariableSuiteBase extends SharedSparkSession
       provider.close()
     }
   }
-
-  def testWithAvroEnc(testName: String, testTags: Tag*)(testBody: Boolean => Any): Unit = {
-    // Run with serde (true)
-    super.test(testName + " (with Avro encoding)", testTags: _*) {
-      super.beforeEach()
-      try {
-        testBody(true)
-      } finally {
-        super.afterEach()
-      }
-    }
-
-    // Run without serde (false)
-    super.test(testName, testTags: _*) {
-      super.beforeEach()
-      try {
-        testBody(false)
-      } finally {
-        super.afterEach()
-      }
-    }
-  }
 }
+

From 35b3b0d3c5c03bb4c8509be1178816b0198099c7 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Fri, 1 Nov 2024 16:26:09 -0700
Subject: [PATCH 10/30] multivalue state encoder

---
 .../execution/streaming/ListStateImpl.scala   |  2 +-
 .../StatefulProcessorHandleImpl.scala         |  3 +-
 .../streaming/state/RocksDBStateEncoder.scala | 47 +++++++++++++++++--
 3 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
index 1047a9e87a837..98e7c73396c70 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
@@ -58,7 +58,7 @@ class ListStateImpl[S](
     valEncoder.schema,
     NoPrefixKeyStateEncoderSpec(keyExprEnc.schema),
     useMultipleValuesPerKey = true,
-    avroEncoderSpec = None)
+    avroEncoderSpec = avroEnc)
 
   /** Whether state exists or not. */
    override def exists(): Boolean = {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
index 8b6e6f0ba3508..a08181cc94bd1 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
@@ -234,7 +234,8 @@ class StatefulProcessorHandleImpl(
 
   override def getListState[T](stateName: String, valEncoder: Encoder[T]): ListState[T] = {
     verifyStateVarOperations("get_list_state", CREATED)
-    val resultState = new ListStateImpl[T](store, stateName, keyEncoder, valEncoder, metrics)
+    val resultState = new ListStateImpl[T](store, stateName, keyEncoder, valEncoder,
+      metrics, schemas(stateName).avroEnc)
     TWSMetricsUtils.incrementMetric(metrics, "numListStateVars")
     resultState
   }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
index 3b4fe6a40633c..0778039e89cc5 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
@@ -126,7 +126,7 @@ object RocksDBStateEncoder {
       useMultipleValuesPerKey: Boolean,
       avroEnc: Option[AvroEncoderSpec] = None): RocksDBValueStateEncoder = {
     if (useMultipleValuesPerKey) {
-      new MultiValuedStateEncoder(valueSchema)
+      new MultiValuedStateEncoder(valueSchema, avroEnc)
     } else {
       new SingleValueStateEncoder(valueSchema, avroEnc)
     }
@@ -752,16 +752,33 @@ class NoPrefixKeyStateEncoder(
  * merged in RocksDB using merge operation, and all merged values can be read using decodeValues
  * operation.
  */
-class MultiValuedStateEncoder(valueSchema: StructType)
+class MultiValuedStateEncoder(
+    valueSchema: StructType,
+    avroEnc: Option[AvroEncoderSpec] = None)
   extends RocksDBValueStateEncoder with Logging {
 
   import RocksDBStateEncoder._
 
   // Reusable objects
+  private val out = new ByteArrayOutputStream
   private val valueRow = new UnsafeRow(valueSchema.size)
+  private val valueAvroType = SchemaConverters.toAvroType(valueSchema)
+  private val valueProj = UnsafeProjection.create(valueSchema)
 
   override def encodeValue(row: UnsafeRow): Array[Byte] = {
-    val bytes = encodeUnsafeRow(row)
+    val bytes = if (avroEnc.isDefined) {
+      val avroData =
+        avroEnc.get.valueSerializer.serialize(row) // InternalRow -> GenericDataRecord
+      out.reset()
+      val encoder = EncoderFactory.get().directBinaryEncoder(out, null)
+      val writer = new GenericDatumWriter[Any](
+        valueAvroType) // Defining Avro writer for this struct type
+      writer.write(avroData, encoder) // GenericDataRecord -> bytes
+      encoder.flush()
+      out.toByteArray
+    } else {
+      encodeUnsafeRow(row)
+    }
     val numBytes = bytes.length
 
     val encodedBytes = new Array[Byte](java.lang.Integer.BYTES + bytes.length)
@@ -780,7 +797,17 @@ class MultiValuedStateEncoder(valueSchema: StructType)
       val encodedValue = new Array[Byte](numBytes)
       Platform.copyMemory(valueBytes, java.lang.Integer.BYTES + Platform.BYTE_ARRAY_OFFSET,
         encodedValue, Platform.BYTE_ARRAY_OFFSET, numBytes)
-      decodeToUnsafeRow(encodedValue, valueRow)
+      if (avroEnc.isDefined) {
+        val reader = new GenericDatumReader[Any](valueAvroType)
+        val decoder = DecoderFactory.get().binaryDecoder(encodedValue,
+          0, encodedValue.length, null)
+        val genericData = reader.read(null, decoder) // bytes -> GenericDataRecord
+        val internalRow = avroEnc.get.valueDeserializer.deserialize(
+          genericData).orNull.asInstanceOf[InternalRow]
+        valueProj.apply(internalRow)
+      } else {
+        decodeToUnsafeRow(encodedValue, valueRow)
+      }
     }
   }
 
@@ -806,7 +833,17 @@ class MultiValuedStateEncoder(valueSchema: StructType)
 
           pos += numBytes
           pos += 1 // eat the delimiter character
-          decodeToUnsafeRow(encodedValue, valueRow)
+          if (avroEnc.isDefined) {
+            val reader = new GenericDatumReader[Any](valueAvroType)
+            val decoder = DecoderFactory.get().binaryDecoder(encodedValue,
+              0, encodedValue.length, null)
+            val genericData = reader.read(null, decoder) // bytes -> GenericDataRecord
+            val internalRow = avroEnc.get.valueDeserializer.deserialize(
+              genericData).orNull.asInstanceOf[InternalRow]
+            valueProj.apply(internalRow)
+          } else {
+            decodeToUnsafeRow(encodedValue, valueRow)
+          }
         }
       }
     }

From dcf0df7744b08c7aebc9b0574a10613cfe26d4bf Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Fri, 1 Nov 2024 18:42:37 -0700
Subject: [PATCH 11/30] encodeToUnsafeRow avro method

---
 .../streaming/state/RocksDBStateEncoder.scala | 79 ++++++++++---------
 1 file changed, 40 insertions(+), 39 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
index 0778039e89cc5..cb6366b3cb79a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
@@ -22,11 +22,12 @@ import java.lang.Double.{doubleToRawLongBits, longBitsToDouble}
 import java.lang.Float.{floatToRawIntBits, intBitsToFloat}
 import java.nio.{ByteBuffer, ByteOrder}
 
+import org.apache.avro.Schema
 import org.apache.avro.generic.{GenericDatumReader, GenericDatumWriter}
 import org.apache.avro.io.{DecoderFactory, EncoderFactory}
 
 import org.apache.spark.internal.Logging
-import org.apache.spark.sql.avro.SchemaConverters
+import org.apache.spark.sql.avro.{AvroDeserializer, AvroSerializer, SchemaConverters}
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{BoundReference, JoinedRow, UnsafeProjection, UnsafeRow}
 import org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter
@@ -154,6 +155,22 @@ object RocksDBStateEncoder {
     encodedBytes
   }
 
+  def encodeUnsafeRow(
+     row: UnsafeRow,
+     avroSerializer: AvroSerializer,
+     valueAvroType: Schema,
+     out: ByteArrayOutputStream): Array[Byte] = {
+    val avroData =
+      avroSerializer.serialize(row) // InternalRow -> GenericDataRecord
+    out.reset()
+    val encoder = EncoderFactory.get().directBinaryEncoder(out, null)
+    val writer = new GenericDatumWriter[Any](
+      valueAvroType) // Defining Avro writer for this struct type
+    writer.write(avroData, encoder) // GenericDataRecord -> bytes
+    encoder.flush()
+    out.toByteArray
+  }
+
   def decodeToUnsafeRow(bytes: Array[Byte], numFields: Int): UnsafeRow = {
     if (bytes != null) {
       val row = new UnsafeRow(numFields)
@@ -163,6 +180,20 @@ object RocksDBStateEncoder {
     }
   }
 
+
+  def decodeToUnsafeRow(
+      valueBytes: Array[Byte],
+      avroDeserializer: AvroDeserializer,
+      valueAvroType: Schema,
+      valueProj: UnsafeProjection): UnsafeRow = {
+    val reader = new GenericDatumReader[Any](valueAvroType)
+    val decoder = DecoderFactory.get().binaryDecoder(valueBytes, 0, valueBytes.length, null)
+    val genericData = reader.read(null, decoder) // bytes -> GenericDataRecord
+    val internalRow = avroDeserializer.deserialize(
+      genericData).orNull.asInstanceOf[InternalRow]
+    valueProj.apply(internalRow)
+  }
+
   def decodeToUnsafeRow(bytes: Array[Byte], reusedRow: UnsafeRow): UnsafeRow = {
     if (bytes != null) {
       // Platform.BYTE_ARRAY_OFFSET is the recommended way refer to the 1st offset. See Platform.
@@ -767,15 +798,7 @@ class MultiValuedStateEncoder(
 
   override def encodeValue(row: UnsafeRow): Array[Byte] = {
     val bytes = if (avroEnc.isDefined) {
-      val avroData =
-        avroEnc.get.valueSerializer.serialize(row) // InternalRow -> GenericDataRecord
-      out.reset()
-      val encoder = EncoderFactory.get().directBinaryEncoder(out, null)
-      val writer = new GenericDatumWriter[Any](
-        valueAvroType) // Defining Avro writer for this struct type
-      writer.write(avroData, encoder) // GenericDataRecord -> bytes
-      encoder.flush()
-      out.toByteArray
+      encodeUnsafeRow(row, avroEnc.get.valueSerializer, valueAvroType, out)
     } else {
       encodeUnsafeRow(row)
     }
@@ -798,13 +821,8 @@ class MultiValuedStateEncoder(
       Platform.copyMemory(valueBytes, java.lang.Integer.BYTES + Platform.BYTE_ARRAY_OFFSET,
         encodedValue, Platform.BYTE_ARRAY_OFFSET, numBytes)
       if (avroEnc.isDefined) {
-        val reader = new GenericDatumReader[Any](valueAvroType)
-        val decoder = DecoderFactory.get().binaryDecoder(encodedValue,
-          0, encodedValue.length, null)
-        val genericData = reader.read(null, decoder) // bytes -> GenericDataRecord
-        val internalRow = avroEnc.get.valueDeserializer.deserialize(
-          genericData).orNull.asInstanceOf[InternalRow]
-        valueProj.apply(internalRow)
+        decodeToUnsafeRow(
+          valueBytes, avroEnc.get.valueDeserializer, valueAvroType, valueProj)
       } else {
         decodeToUnsafeRow(encodedValue, valueRow)
       }
@@ -834,13 +852,8 @@ class MultiValuedStateEncoder(
           pos += numBytes
           pos += 1 // eat the delimiter character
           if (avroEnc.isDefined) {
-            val reader = new GenericDatumReader[Any](valueAvroType)
-            val decoder = DecoderFactory.get().binaryDecoder(encodedValue,
-              0, encodedValue.length, null)
-            val genericData = reader.read(null, decoder) // bytes -> GenericDataRecord
-            val internalRow = avroEnc.get.valueDeserializer.deserialize(
-              genericData).orNull.asInstanceOf[InternalRow]
-            valueProj.apply(internalRow)
+            decodeToUnsafeRow(
+              valueBytes, avroEnc.get.valueDeserializer, valueAvroType, valueProj)
           } else {
             decodeToUnsafeRow(encodedValue, valueRow)
           }
@@ -879,15 +892,7 @@ class SingleValueStateEncoder(
 
   override def encodeValue(row: UnsafeRow): Array[Byte] = {
     if (avroEnc.isDefined) {
-      val avroData =
-        avroEnc.get.valueSerializer.serialize(row) // InternalRow -> GenericDataRecord
-      out.reset()
-      val encoder = EncoderFactory.get().directBinaryEncoder(out, null)
-      val writer = new GenericDatumWriter[Any](
-        valueAvroType) // Defining Avro writer for this struct type
-      writer.write(avroData, encoder) // GenericDataRecord -> bytes
-      encoder.flush()
-      out.toByteArray
+      encodeUnsafeRow(row, avroEnc.get.valueSerializer, valueAvroType, out)
     } else {
       encodeUnsafeRow(row)
     }
@@ -904,12 +909,8 @@ class SingleValueStateEncoder(
       return null
     }
     if (avroEnc.isDefined) {
-      val reader = new GenericDatumReader[Any](valueAvroType)
-      val decoder = DecoderFactory.get().binaryDecoder(valueBytes, 0, valueBytes.length, null)
-      val genericData = reader.read(null, decoder) // bytes -> GenericDataRecord
-      val internalRow = avroEnc.get.valueDeserializer.deserialize(
-        genericData).orNull.asInstanceOf[InternalRow]
-      valueProj.apply(internalRow)
+      decodeToUnsafeRow(
+        valueBytes, avroEnc.get.valueDeserializer, valueAvroType, valueProj)
     } else {
       decodeToUnsafeRow(valueBytes, valueRow)
     }

From dfc6b1eeb028f9cb95ef5bcff054fb533cf8f1dd Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Mon, 4 Nov 2024 10:46:09 -0800
Subject: [PATCH 12/30] using correct val

---
 .../sql/execution/streaming/state/RocksDBStateEncoder.scala   | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
index cb6366b3cb79a..3d497e08aa7d0 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
@@ -822,7 +822,7 @@ class MultiValuedStateEncoder(
         encodedValue, Platform.BYTE_ARRAY_OFFSET, numBytes)
       if (avroEnc.isDefined) {
         decodeToUnsafeRow(
-          valueBytes, avroEnc.get.valueDeserializer, valueAvroType, valueProj)
+          encodedValue, avroEnc.get.valueDeserializer, valueAvroType, valueProj)
       } else {
         decodeToUnsafeRow(encodedValue, valueRow)
       }
@@ -853,7 +853,7 @@ class MultiValuedStateEncoder(
           pos += 1 // eat the delimiter character
           if (avroEnc.isDefined) {
             decodeToUnsafeRow(
-              valueBytes, avroEnc.get.valueDeserializer, valueAvroType, valueProj)
+              encodedValue, avroEnc.get.valueDeserializer, valueAvroType, valueProj)
           } else {
             decodeToUnsafeRow(encodedValue, valueRow)
           }

From 5b98aa68937a600bb276c0774db0637c78684198 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Mon, 4 Nov 2024 14:27:30 -0800
Subject: [PATCH 13/30] comments

---
 .../execution/streaming/ListStateImpl.scala   |  2 +
 .../StateStoreColumnFamilySchemaUtils.scala   | 14 +++++--
 .../StatefulProcessorHandleImpl.scala         |  2 +
 .../execution/streaming/ValueStateImpl.scala  |  2 +
 .../streaming/state/RocksDBStateEncoder.scala | 41 ++++++++++++++-----
 .../StateSchemaCompatibilityChecker.scala     |  2 +
 .../streaming/state/StateStore.scala          |  3 +-
 .../TransformWithListStateSuite.scala         | 13 +++---
 8 files changed, 57 insertions(+), 22 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
index 98e7c73396c70..e44ceadbe6f54 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
@@ -33,6 +33,8 @@ import org.apache.spark.sql.types.StructType
  * @param keyExprEnc - Spark SQL encoder for key
  * @param valEncoder - Spark SQL encoder for value
  * @param metrics - metrics to be updated as part of stateful processing
+ * @param avroEnc - optional Avro serializer and deserializer for this state variable that
+ *                is used by the StateStore to encode state in Avro format
  * @tparam S - data type of object that will be stored in the list
  */
 class ListStateImpl[S](
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
index 744fd6d5b6b14..78b9040dabf8b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
@@ -30,8 +30,18 @@ object StateStoreColumnFamilySchemaUtils {
     new StateStoreColumnFamilySchemaUtils(initializeAvroSerde)
 }
 
+/**
+ *
+ * @param initializeAvroSerde Whether or not to create the Avro serializers and deserializers
+ *                            for this state type. This class is used to create the
+ *                            StateStoreColumnFamilySchema for each state variable from the driver
+ */
 class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) {
 
+  /**
+   * If initializeAvroSerde is true, this method will create an Avro Serializer and Deserializer
+   * for a particular key and value schema.
+   */
   private def getAvroSerde(
       keySchema: StructType, valSchema: StructType): Option[AvroEncoderSpec] = {
     if (initializeAvroSerde) {
@@ -87,14 +97,12 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) {
       valEncoder: Encoder[V],
       hasTtl: Boolean): StateStoreColFamilySchema = {
     val compositeKeySchema = getCompositeKeySchema(keyEncoder.schema, userKeyEnc.schema)
-    val valSchema = getValueSchemaWithTTL(valEncoder.schema, hasTtl)
     StateStoreColFamilySchema(
       stateName,
       compositeKeySchema,
       getValueSchemaWithTTL(valEncoder.schema, hasTtl),
       Some(PrefixKeyScanStateEncoderSpec(compositeKeySchema, 1)),
-      Some(userKeyEnc.schema),
-      avroEnc = getAvroSerde(compositeKeySchema, valSchema))
+      Some(userKeyEnc.schema))
   }
 
   def getTimerStateSchema(
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
index a08181cc94bd1..c4ea6373fbf19 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
@@ -96,6 +96,8 @@ class QueryInfoImpl(
  * @param isStreaming - defines whether the query is streaming or batch
  * @param batchTimestampMs - timestamp for the current batch if available
  * @param metrics - metrics to be updated as part of stateful processing
+ * @param schemas - StateStoreColumnFamilySchemas that include Avro serializers and deserializers
+ *                for each state variable, if Avro encoding is enabled for this query
  */
 class StatefulProcessorHandleImpl(
     store: StateStore,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala
index db8c405ee6193..22a7e892d25aa 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala
@@ -31,6 +31,8 @@ import org.apache.spark.sql.streaming.ValueState
  * @param keyExprEnc - Spark SQL encoder for key
  * @param valEncoder - Spark SQL encoder for value
  * @param metrics - metrics to be updated as part of stateful processing
+ * @param avroEnc - optional Avro serializer and deserializer for this state variable that
+ *                is used by the StateStore to encode state in Avro format
  * @tparam S - data type of object that will be stored
  */
 class ValueStateImpl[S](
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
index 3d497e08aa7d0..9f30e97faa54a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
@@ -155,18 +155,22 @@ object RocksDBStateEncoder {
     encodedBytes
   }
 
+  /**
+   * This method takes an UnsafeRow, and serializes to a byte array using Avro encoding.
+   */
   def encodeUnsafeRow(
      row: UnsafeRow,
      avroSerializer: AvroSerializer,
      valueAvroType: Schema,
      out: ByteArrayOutputStream): Array[Byte] = {
+    // InternalRow -> Avro.GenericDataRecord
     val avroData =
-      avroSerializer.serialize(row) // InternalRow -> GenericDataRecord
+      avroSerializer.serialize(row)
     out.reset()
     val encoder = EncoderFactory.get().directBinaryEncoder(out, null)
     val writer = new GenericDatumWriter[Any](
       valueAvroType) // Defining Avro writer for this struct type
-    writer.write(avroData, encoder) // GenericDataRecord -> bytes
+    writer.write(avroData, encoder) // Avro.GenericDataRecord -> byte array
     encoder.flush()
     out.toByteArray
   }
@@ -180,7 +184,10 @@ object RocksDBStateEncoder {
     }
   }
 
-
+  /**
+   * This method takes a byte array written using Avro encoding, and
+   * deserializes to an UnsafeRow using the Avro deserializer
+   */
   def decodeToUnsafeRow(
       valueBytes: Array[Byte],
       avroDeserializer: AvroDeserializer,
@@ -188,9 +195,12 @@ object RocksDBStateEncoder {
       valueProj: UnsafeProjection): UnsafeRow = {
     val reader = new GenericDatumReader[Any](valueAvroType)
     val decoder = DecoderFactory.get().binaryDecoder(valueBytes, 0, valueBytes.length, null)
-    val genericData = reader.read(null, decoder) // bytes -> GenericDataRecord
+    // bytes -> Avro.GenericDataRecord
+    val genericData = reader.read(null, decoder)
+    // Avro.GenericDataRecord -> InternalRow
     val internalRow = avroDeserializer.deserialize(
       genericData).orNull.asInstanceOf[InternalRow]
+    // InternalRow -> UnsafeRow
     valueProj.apply(internalRow)
   }
 
@@ -214,6 +224,8 @@ object RocksDBStateEncoder {
  * @param keySchema - schema of the key to be encoded
  * @param numColsPrefixKey - number of columns to be used for prefix key
  * @param useColumnFamilies - if column family is enabled for this encoder
+ * @param avroEnc - if Avro encoding is specified for this StateEncoder, this encoder will
+ *                be defined
  */
 class PrefixKeyScanStateEncoder(
     keySchema: StructType,
@@ -308,7 +320,6 @@ class PrefixKeyScanStateEncoder(
   }
 
   override def supportPrefixKeyScan: Boolean = true
-
 }
 
 /**
@@ -341,6 +352,8 @@ class PrefixKeyScanStateEncoder(
  * @param keySchema - schema of the key to be encoded
  * @param orderingOrdinals - the ordinals for which the range scan is constructed
  * @param useColumnFamilies - if column family is enabled for this encoder
+ * @param avroEnc - if Avro encoding is specified for this StateEncoder, this encoder will
+ *                be defined
  */
 class RangeKeyScanStateEncoder(
     keySchema: StructType,
@@ -700,6 +713,7 @@ class RangeKeyScanStateEncoder(
  *    The bytes of a UnsafeRow is written unmodified to starting from offset 1
  *    (offset 0 is the version byte of value 0). That is, if the unsafe row has N bytes,
  *    then the generated array byte will be N+1 bytes.
+ * If the avroEnc is specified, we are using Avro encoding for this column family's keys
  */
 class NoPrefixKeyStateEncoder(
     keySchema: StructType,
@@ -711,6 +725,7 @@ class NoPrefixKeyStateEncoder(
   import RocksDBStateEncoder._
 
   // Reusable objects
+  private val usingAvroEncoding = avroEnc.isDefined
   private val keyRow = new UnsafeRow(keySchema.size)
   private val keyAvroType = SchemaConverters.toAvroType(keySchema)
 
@@ -720,7 +735,7 @@ class NoPrefixKeyStateEncoder(
     } else {
       // If avroEnc is defined, we know that we need to use Avro to
       // encode this UnsafeRow to Avro bytes
-      val bytesToEncode = if (avroEnc.isDefined) {
+      val bytesToEncode = if (usingAvroEncoding) {
         val avroData = avroEnc.get.keySerializer.serialize(row)
         out.reset()
         val encoder = EncoderFactory.get().directBinaryEncoder(out, null)
@@ -782,6 +797,7 @@ class NoPrefixKeyStateEncoder(
  * This encoder supports RocksDB StringAppendOperator merge operator. Values encoded can be
  * merged in RocksDB using merge operation, and all merged values can be read using decodeValues
  * operation.
+ * If the avroEnc is specified, we are using Avro encoding for this column family's values
  */
 class MultiValuedStateEncoder(
     valueSchema: StructType,
@@ -790,6 +806,7 @@ class MultiValuedStateEncoder(
 
   import RocksDBStateEncoder._
 
+  private val usingAvroEncoding = avroEnc.isDefined
   // Reusable objects
   private val out = new ByteArrayOutputStream
   private val valueRow = new UnsafeRow(valueSchema.size)
@@ -797,7 +814,7 @@ class MultiValuedStateEncoder(
   private val valueProj = UnsafeProjection.create(valueSchema)
 
   override def encodeValue(row: UnsafeRow): Array[Byte] = {
-    val bytes = if (avroEnc.isDefined) {
+    val bytes = if (usingAvroEncoding) {
       encodeUnsafeRow(row, avroEnc.get.valueSerializer, valueAvroType, out)
     } else {
       encodeUnsafeRow(row)
@@ -820,7 +837,7 @@ class MultiValuedStateEncoder(
       val encodedValue = new Array[Byte](numBytes)
       Platform.copyMemory(valueBytes, java.lang.Integer.BYTES + Platform.BYTE_ARRAY_OFFSET,
         encodedValue, Platform.BYTE_ARRAY_OFFSET, numBytes)
-      if (avroEnc.isDefined) {
+      if (usingAvroEncoding) {
         decodeToUnsafeRow(
           encodedValue, avroEnc.get.valueDeserializer, valueAvroType, valueProj)
       } else {
@@ -851,7 +868,7 @@ class MultiValuedStateEncoder(
 
           pos += numBytes
           pos += 1 // eat the delimiter character
-          if (avroEnc.isDefined) {
+          if (usingAvroEncoding) {
             decodeToUnsafeRow(
               encodedValue, avroEnc.get.valueDeserializer, valueAvroType, valueProj)
           } else {
@@ -876,6 +893,7 @@ class MultiValuedStateEncoder(
  *    The bytes of a UnsafeRow is written unmodified to starting from offset 1
  *    (offset 0 is the version byte of value 0). That is, if the unsafe row has N bytes,
  *    then the generated array byte will be N+1 bytes.
+ * If the avroEnc is specified, we are using Avro encoding for this column family's values
  */
 class SingleValueStateEncoder(
     valueSchema: StructType,
@@ -884,6 +902,7 @@ class SingleValueStateEncoder(
 
   import RocksDBStateEncoder._
 
+  private val usingAvroEncoding = avroEnc.isDefined
   // Reusable objects
   private val out = new ByteArrayOutputStream
   private val valueRow = new UnsafeRow(valueSchema.size)
@@ -891,7 +910,7 @@ class SingleValueStateEncoder(
   private val valueProj = UnsafeProjection.create(valueSchema)
 
   override def encodeValue(row: UnsafeRow): Array[Byte] = {
-    if (avroEnc.isDefined) {
+    if (usingAvroEncoding) {
       encodeUnsafeRow(row, avroEnc.get.valueSerializer, valueAvroType, out)
     } else {
       encodeUnsafeRow(row)
@@ -908,7 +927,7 @@ class SingleValueStateEncoder(
     if (valueBytes == null) {
       return null
     }
-    if (avroEnc.isDefined) {
+    if (usingAvroEncoding) {
       decodeToUnsafeRow(
         valueBytes, avroEnc.get.valueDeserializer, valueAvroType, valueProj)
     } else {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
index d8094f78f587f..ff07666ce0fe9 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
@@ -38,6 +38,8 @@ case class StateSchemaValidationResult(
     schemaPath: String
 )
 
+// Avro encoder that is used by the RocksDBStateStoreProvider and RocksDBStateEncoder
+// in order to serialize from UnsafeRow to a byte array of Avro encoding.
 case class AvroEncoderSpec(
   keySerializer: AvroSerializer,
   keyDeserializer: AvroDeserializer,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
index d5f79f27c7ff7..255577dfccaa6 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
@@ -133,7 +133,8 @@ trait StateStore extends ReadStateStore {
 
   /**
    * Create column family with given name, if absent.
-   *
+   * If Avro encoding is enabled for this query, we expect the avroEncoderSpec to
+   * be defined so that the Key and Value StateEncoders will use this.
    * @return column family ID
    */
   def createColFamilyIfAbsent(
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateSuite.scala
index 8e5a2fd183a8d..9442c95d1b016 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateSuite.scala
@@ -296,13 +296,12 @@ class TransformWithListStateSuite extends StreamTest
         AddData(inputData, InputRow("k5", "append", "v4")),
         AddData(inputData, InputRow("k5", "put", "v5,v6")),
         AddData(inputData, InputRow("k5", "emitAllInState", "")),
-        CheckNewAnswer(("k5", "v5"), ("k5", "v6"))
-        // TODO: Uncomment once we have implemented ListStateMetrics for Avro encoding
-//        Execute { q =>
-//          assert(q.lastProgress.stateOperators(0).customMetrics.get("numListStateVars") > 0)
-//          assert(q.lastProgress.stateOperators(0).numRowsUpdated === 2)
-//          assert(q.lastProgress.stateOperators(0).numRowsRemoved === 2)
-//        }
+        CheckNewAnswer(("k5", "v5"), ("k5", "v6")),
+        Execute { q =>
+          assert(q.lastProgress.stateOperators(0).customMetrics.get("numListStateVars") > 0)
+          assert(q.lastProgress.stateOperators(0).numRowsUpdated === 2)
+          assert(q.lastProgress.stateOperators(0).numRowsRemoved === 2)
+        }
       )
     }
   }

From 0d37ffd337f0463196f0dc7a6879ff72e5c49617 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Mon, 4 Nov 2024 14:47:33 -0800
Subject: [PATCH 14/30] calling encodeUnsafeRow

---
 .../execution/streaming/state/RocksDBStateEncoder.scala   | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
index 9f30e97faa54a..7fe7da2ff290a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
@@ -736,13 +736,7 @@ class NoPrefixKeyStateEncoder(
       // If avroEnc is defined, we know that we need to use Avro to
       // encode this UnsafeRow to Avro bytes
       val bytesToEncode = if (usingAvroEncoding) {
-        val avroData = avroEnc.get.keySerializer.serialize(row)
-        out.reset()
-        val encoder = EncoderFactory.get().directBinaryEncoder(out, null)
-        val writer = new GenericDatumWriter[Any](keyAvroType)
-        writer.write(avroData, encoder)
-        encoder.flush()
-        out.toByteArray
+        encodeUnsafeRow(row, avroEnc.get.keySerializer, keyAvroType, out)
       } else row.getBytes
       val (encodedBytes, startingOffset) = encodeColumnFamilyPrefix(
         bytesToEncode.length +

From 9b8dd5d9c0820295bc2fad111314a7ab848a2bd5 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <132308037+ericm-db@users.noreply.github.com>
Date: Wed, 6 Nov 2024 16:04:34 -0800
Subject: [PATCH 15/30] [SPARK-50127] Implement Avro encoding for MapState and
 PrefixKeyScanStateEncoder (#22)

* [WIP] Supporting Map State with Avro encoding

* adding comments

* [WIP] Avro Range Scan (#23)

* valuestatettl

* mapstate, valuestate ttl works

* timers

* renaming to suffix key

* cleaning up
---
 .../streaming/ListStateImplWithTTL.scala      |  15 +-
 .../execution/streaming/MapStateImpl.scala    |   9 +-
 .../streaming/MapStateImplWithTTL.scala       |  15 +-
 .../StateStoreColumnFamilySchemaUtils.scala   | 130 +++++++++++++-
 .../StatefulProcessorHandleImpl.scala         |  32 +++-
 .../sql/execution/streaming/TTLState.scala    |  17 +-
 .../execution/streaming/TimerStateImpl.scala  |  17 +-
 .../streaming/ValueStateImplWithTTL.scala     |  14 +-
 .../streaming/state/RocksDBStateEncoder.scala | 158 +++++++++++++++---
 .../StateSchemaCompatibilityChecker.scala     |   4 +-
 .../TransformWithListStateTTLSuite.scala      |   6 +-
 .../TransformWithMapStateSuite.scala          |   6 +-
 .../TransformWithMapStateTTLSuite.scala       |   4 +-
 .../streaming/TransformWithStateTTLTest.scala |  12 +-
 .../TransformWithValueStateTTLSuite.scala     |   2 +-
 15 files changed, 369 insertions(+), 72 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala
index 4c8dd6a193c25..445ca743b3855 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala
@@ -20,7 +20,7 @@ import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, StateStore, StateStoreErrors}
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, NoPrefixKeyStateEncoderSpec, StateStore, StateStoreErrors}
 import org.apache.spark.sql.streaming.{ListState, TTLConfig}
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.util.NextIterator
@@ -36,6 +36,10 @@ import org.apache.spark.util.NextIterator
  * @param ttlConfig  - TTL configuration for values  stored in this state
  * @param batchTimestampMs - current batch processing timestamp.
  * @param metrics - metrics to be updated as part of stateful processing
+ * @param avroEnc - optional Avro serializer and deserializer for this state variable that
+ *                is used by the StateStore to encode state in Avro format
+ * @param ttlAvroEnc - optional Avro serializer and deserializer for TTL state that
+ *                is used by the StateStore to encode state in Avro format
  * @tparam S - data type of object that will be stored
  */
 class ListStateImplWithTTL[S](
@@ -45,8 +49,10 @@ class ListStateImplWithTTL[S](
     valEncoder: ExpressionEncoder[Any],
     ttlConfig: TTLConfig,
     batchTimestampMs: Long,
-    metrics: Map[String, SQLMetric] = Map.empty)
-  extends SingleKeyTTLStateImpl(stateName, store, keyExprEnc, batchTimestampMs)
+    metrics: Map[String, SQLMetric] = Map.empty,
+    avroEnc: Option[AvroEncoderSpec] = None,
+    ttlAvroEnc: Option[AvroEncoderSpec] = None)
+  extends SingleKeyTTLStateImpl(stateName, store, keyExprEnc, batchTimestampMs, ttlAvroEnc)
   with ListStateMetricsImpl
   with ListState[S] {
 
@@ -65,7 +71,8 @@ class ListStateImplWithTTL[S](
   private def initialize(): Unit = {
     store.createColFamilyIfAbsent(stateName, keyExprEnc.schema,
       getValueSchemaWithTTL(valEncoder.schema, true),
-      NoPrefixKeyStateEncoderSpec(keyExprEnc.schema), useMultipleValuesPerKey = true)
+      NoPrefixKeyStateEncoderSpec(keyExprEnc.schema), useMultipleValuesPerKey = true,
+      avroEncoderSpec = avroEnc)
   }
 
   /** Whether state exists or not. */
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala
index 4e608a5d5dbbe..eb96b32722aaf 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala
@@ -20,7 +20,7 @@ import org.apache.spark.internal.Logging
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{PrefixKeyScanStateEncoderSpec, StateStore, StateStoreErrors, UnsafeRowPair}
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, PrefixKeyScanStateEncoderSpec, StateStore, StateStoreErrors, UnsafeRowPair}
 import org.apache.spark.sql.streaming.MapState
 import org.apache.spark.sql.types.StructType
 
@@ -32,6 +32,8 @@ import org.apache.spark.sql.types.StructType
  * @param keyExprEnc - Spark SQL encoder for key
  * @param valEncoder - Spark SQL encoder for value
  * @param metrics - metrics to be updated as part of stateful processing
+ * @param avroEnc - optional Avro serializer and deserializer for this state variable that
+ *                is used by the StateStore to encode state in Avro format
  * @tparam K - type of key for map state variable
  * @tparam V - type of value for map state variable
  */
@@ -41,7 +43,8 @@ class MapStateImpl[K, V](
     keyExprEnc: ExpressionEncoder[Any],
     userKeyEnc: ExpressionEncoder[Any],
     valEncoder: ExpressionEncoder[Any],
-    metrics: Map[String, SQLMetric] = Map.empty) extends MapState[K, V] with Logging {
+    metrics: Map[String, SQLMetric] = Map.empty,
+    avroEnc: Option[AvroEncoderSpec] = None) extends MapState[K, V] with Logging {
 
   // Pack grouping key and user key together as a prefixed composite key
   private val schemaForCompositeKeyRow: StructType = {
@@ -52,7 +55,7 @@ class MapStateImpl[K, V](
     keyExprEnc, userKeyEnc, valEncoder, stateName)
 
   store.createColFamilyIfAbsent(stateName, schemaForCompositeKeyRow, schemaForValueRow,
-    PrefixKeyScanStateEncoderSpec(schemaForCompositeKeyRow, 1))
+    PrefixKeyScanStateEncoderSpec(schemaForCompositeKeyRow, 1), avroEncoderSpec = avroEnc)
 
   /** Whether state exists or not. */
   override def exists(): Boolean = {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala
index 19704b6d1bd59..11554d8532396 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala
@@ -21,7 +21,7 @@ import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{PrefixKeyScanStateEncoderSpec, StateStore, StateStoreErrors}
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, PrefixKeyScanStateEncoderSpec, StateStore, StateStoreErrors}
 import org.apache.spark.sql.streaming.{MapState, TTLConfig}
 import org.apache.spark.util.NextIterator
 
@@ -36,6 +36,10 @@ import org.apache.spark.util.NextIterator
  * @param ttlConfig  - the ttl configuration (time to live duration etc.)
  * @param batchTimestampMs - current batch processing timestamp.
  * @param metrics - metrics to be updated as part of stateful processing
+ * @param avroEnc - optional Avro serializer and deserializer for this state variable that
+ *                is used by the StateStore to encode state in Avro format
+ * @param ttlAvroEnc - optional Avro serializer and deserializer for TTL state that
+ *                is used by the StateStore to encode state in Avro format
  * @tparam K - type of key for map state variable
  * @tparam V - type of value for map state variable
  * @return - instance of MapState of type [K,V] that can be used to store state persistently
@@ -48,9 +52,11 @@ class MapStateImplWithTTL[K, V](
     valEncoder: ExpressionEncoder[Any],
     ttlConfig: TTLConfig,
     batchTimestampMs: Long,
-    metrics: Map[String, SQLMetric] = Map.empty)
+    metrics: Map[String, SQLMetric] = Map.empty,
+    avroEnc: Option[AvroEncoderSpec] = None,
+    ttlAvroEnc: Option[AvroEncoderSpec] = None)
   extends CompositeKeyTTLStateImpl[K](stateName, store,
-    keyExprEnc, userKeyEnc, batchTimestampMs)
+    keyExprEnc, userKeyEnc, batchTimestampMs, ttlAvroEnc)
   with MapState[K, V] with Logging {
 
   private val stateTypesEncoder = new CompositeKeyStateEncoder(
@@ -66,7 +72,8 @@ class MapStateImplWithTTL[K, V](
       getCompositeKeySchema(keyExprEnc.schema, userKeyEnc.schema)
     store.createColFamilyIfAbsent(stateName, schemaForCompositeKeyRow,
       getValueSchemaWithTTL(valEncoder.schema, true),
-      PrefixKeyScanStateEncoderSpec(schemaForCompositeKeyRow, 1))
+      PrefixKeyScanStateEncoderSpec(schemaForCompositeKeyRow, 1),
+      avroEncoderSpec = avroEnc)
   }
 
   /** Whether state exists or not. */
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
index 78b9040dabf8b..734436f7b2c27 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
@@ -20,14 +20,45 @@ import org.apache.spark.sql.Encoder
 import org.apache.spark.sql.avro.{AvroDeserializer, AvroOptions, AvroSerializer, SchemaConverters}
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, PrefixKeyScanStateEncoderSpec, StateStoreColFamilySchema}
-import org.apache.spark.sql.execution.streaming.state.AvroEncoderSpec
-import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, NoPrefixKeyStateEncoderSpec, PrefixKeyScanStateEncoderSpec, RangeKeyScanStateEncoderSpec, StateStoreColFamilySchema}
+import org.apache.spark.sql.types.{BinaryType, BooleanType, ByteType, DataType, DoubleType, FloatType, IntegerType, LongType, NullType, ShortType, StructField, StructType}
 
 object StateStoreColumnFamilySchemaUtils {
 
   def apply(initializeAvroSerde: Boolean): StateStoreColumnFamilySchemaUtils =
     new StateStoreColumnFamilySchemaUtils(initializeAvroSerde)
+
+
+  /**
+   * Avro uses zig-zag encoding for some fixed-length types, like Longs and Ints. For range scans
+   * we want to use big-endian encoding, so we need to convert the source schema to replace these
+   * types with BinaryType.
+   *
+   * @param schema The schema to convert
+   * @param ordinals If non-empty, only convert fields at these ordinals.
+   *                 If empty, convert all fields.
+   */
+  def convertForRangeScan(schema: StructType, ordinals: Seq[Int] = Seq.empty): StructType = {
+    val ordinalSet = ordinals.toSet
+    StructType(schema.fields.zipWithIndex.map { case (field, idx) =>
+      if ((ordinals.isEmpty || ordinalSet.contains(idx)) && isFixedSize(field.dataType)) {
+        // Convert numeric types to BinaryType while preserving nullability
+        field.copy(dataType = BinaryType)
+      } else {
+        field
+      }
+    })
+  }
+
+  private def isFixedSize(dataType: DataType): Boolean = dataType match {
+    case _: ByteType | _: BooleanType | _: ShortType | _: IntegerType | _: LongType |
+         _: FloatType | _: DoubleType => true
+    case _ => false
+  }
+
+  def getTtlColFamilyName(stateName: String): String = {
+    "$ttl_" + stateName
+  }
 }
 
 /**
@@ -43,7 +74,10 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) {
    * for a particular key and value schema.
    */
   private def getAvroSerde(
-      keySchema: StructType, valSchema: StructType): Option[AvroEncoderSpec] = {
+      keySchema: StructType,
+      valSchema: StructType,
+      suffixKeySchema: Option[StructType] = None
+  ): Option[AvroEncoderSpec] = {
     if (initializeAvroSerde) {
       val avroType = SchemaConverters.toAvroType(valSchema)
       val avroOptions = AvroOptions(Map.empty)
@@ -56,7 +90,18 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) {
       val valueDeserializer = new AvroDeserializer(avroType, valSchema,
         avroOptions.datetimeRebaseModeInRead, avroOptions.useStableIdForUnionType,
         avroOptions.stableIdPrefixForUnionType, avroOptions.recursiveFieldMaxDepth)
-      Some(AvroEncoderSpec(keySer, keyDe, valueSerializer, valueDeserializer))
+      val (suffixKeySer, suffixKeyDe) = if (suffixKeySchema.isDefined) {
+        val userKeyAvroType = SchemaConverters.toAvroType(suffixKeySchema.get)
+        val skSer = new AvroSerializer(suffixKeySchema.get, userKeyAvroType, nullable = false)
+        val skDe = new AvroDeserializer(userKeyAvroType, suffixKeySchema.get,
+          avroOptions.datetimeRebaseModeInRead, avroOptions.useStableIdForUnionType,
+          avroOptions.stableIdPrefixForUnionType, avroOptions.recursiveFieldMaxDepth)
+        (Some(skSer), Some(skDe))
+      } else {
+        (None, None)
+      }
+      Some(AvroEncoderSpec(
+        keySer, keyDe, valueSerializer, valueDeserializer, suffixKeySer, suffixKeyDe))
     } else {
       None
     }
@@ -97,12 +142,60 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) {
       valEncoder: Encoder[V],
       hasTtl: Boolean): StateStoreColFamilySchema = {
     val compositeKeySchema = getCompositeKeySchema(keyEncoder.schema, userKeyEnc.schema)
+    val valSchema = getValueSchemaWithTTL(valEncoder.schema, hasTtl)
     StateStoreColFamilySchema(
       stateName,
       compositeKeySchema,
       getValueSchemaWithTTL(valEncoder.schema, hasTtl),
       Some(PrefixKeyScanStateEncoderSpec(compositeKeySchema, 1)),
-      Some(userKeyEnc.schema))
+      Some(userKeyEnc.schema),
+      avroEnc = getAvroSerde(
+        StructType(compositeKeySchema.take(1)),
+        valSchema,
+        Some(StructType(compositeKeySchema.drop(1)))
+      )
+    )
+  }
+
+  def getTtlStateSchema(
+      stateName: String,
+      keyEncoder: ExpressionEncoder[Any]): StateStoreColFamilySchema = {
+    val ttlKeySchema = StateStoreColumnFamilySchemaUtils.convertForRangeScan(
+      getSingleKeyTTLRowSchema(keyEncoder.schema), Seq(0))
+    val ttlValSchema = StructType(
+      Array(StructField("__dummy__", NullType)))
+    StateStoreColFamilySchema(
+      stateName,
+      ttlKeySchema,
+      ttlValSchema,
+      Some(RangeKeyScanStateEncoderSpec(ttlKeySchema, Seq(0))),
+      avroEnc = getAvroSerde(
+        StructType(ttlKeySchema.take(1)),
+        ttlValSchema,
+        Some(StructType(ttlKeySchema.drop(1)))
+      )
+    )
+  }
+
+  def getTtlStateSchema(
+      stateName: String,
+      keyEncoder: ExpressionEncoder[Any],
+      userKeySchema: StructType): StateStoreColFamilySchema = {
+    val ttlKeySchema = StateStoreColumnFamilySchemaUtils.convertForRangeScan(
+      getCompositeKeyTTLRowSchema(keyEncoder.schema, userKeySchema), Seq(0))
+    val ttlValSchema = StructType(
+      Array(StructField("__dummy__", NullType)))
+    StateStoreColFamilySchema(
+      stateName,
+      ttlKeySchema,
+      ttlValSchema,
+      Some(RangeKeyScanStateEncoderSpec(ttlKeySchema, Seq(0))),
+      avroEnc = getAvroSerde(
+        StructType(ttlKeySchema.take(1)),
+        ttlValSchema,
+        Some(StructType(ttlKeySchema.drop(1)))
+      )
+    )
   }
 
   def getTimerStateSchema(
@@ -113,6 +206,29 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) {
       stateName,
       keySchema,
       valSchema,
-      Some(PrefixKeyScanStateEncoderSpec(keySchema, 1)))
+      Some(PrefixKeyScanStateEncoderSpec(keySchema, 1)),
+      avroEnc = getAvroSerde(
+        StructType(keySchema.take(1)),
+        valSchema,
+        Some(StructType(keySchema.drop(1)))
+      ))
+  }
+
+  def getTimerStateSchemaForSecIndex(
+      stateName: String,
+      keySchema: StructType,
+      valSchema: StructType): StateStoreColFamilySchema = {
+    val avroKeySchema = StateStoreColumnFamilySchemaUtils.
+      convertForRangeScan(keySchema, Seq(0))
+    StateStoreColFamilySchema(
+      stateName,
+      keySchema,
+      valSchema,
+      Some(RangeKeyScanStateEncoderSpec(keySchema, Seq(0))),
+      avroEnc = getAvroSerde(
+        StructType(avroKeySchema.take(1)),
+        valSchema,
+        Some(StructType(avroKeySchema.drop(1)))
+      ))
   }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
index ab5b5d42041c2..0258691e51d66 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
@@ -27,6 +27,7 @@ import org.apache.spark.sql.Encoder
 import org.apache.spark.sql.catalyst.encoders.{encoderFor, ExpressionEncoder}
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.streaming.StatefulProcessorHandleState.PRE_INIT
+import org.apache.spark.sql.execution.streaming.StateStoreColumnFamilySchemaUtils.getTtlColFamilyName
 import org.apache.spark.sql.execution.streaming.state._
 import org.apache.spark.sql.streaming.{ListState, MapState, QueryInfo, TimeMode, TTLConfig, ValueState}
 import org.apache.spark.util.Utils
@@ -140,7 +141,13 @@ class StatefulProcessorHandleImpl(
 
   override def getQueryInfo(): QueryInfo = currQueryInfo
 
-  private lazy val timerState = new TimerStateImpl(store, timeMode, keyEncoder)
+  private lazy val timerStateName = TimerStateUtils.getTimerStateVarName(
+    timeMode.toString)
+  private lazy val timerSecIndexColFamily = TimerStateUtils.getSecIndexColFamilyName(
+    timeMode.toString)
+  private lazy val timerState = new TimerStateImpl(
+    store, timeMode, keyEncoder, schemas(timerStateName).avroEnc,
+    schemas(timerSecIndexColFamily).avroEnc)
 
   /**
    * Function to register a timer for the given expiryTimestampMs
@@ -323,7 +330,7 @@ class StatefulProcessorHandleImpl(
       mapStateWithTTL
     } else {
       val mapStateWithoutTTL = new MapStateImpl[K, V](store, stateName, keyEncoder,
-        userKeyEnc, valEncoder, metrics)
+        userKeyEnc, valEncoder, metrics, schemas(stateName).avroEnc)
       TWSMetricsUtils.incrementMetric(metrics, "numMapStateVars")
       mapStateWithoutTTL
     }
@@ -382,10 +389,16 @@ class DriverStatefulProcessorHandleImpl(
 
   private def addTimerColFamily(): Unit = {
     val stateName = TimerStateUtils.getTimerStateVarName(timeMode.toString)
+    val secIndexColFamilyName = TimerStateUtils.getSecIndexColFamilyName(timeMode.toString)
     val timerEncoder = new TimerKeyEncoder(keyExprEnc)
     val colFamilySchema = schemaUtils.
       getTimerStateSchema(stateName, timerEncoder.schemaForKeyRow, timerEncoder.schemaForValueRow)
+    val secIndexColFamilySchema = schemaUtils.
+      getTimerStateSchemaForSecIndex(secIndexColFamilyName,
+        timerEncoder.keySchemaForSecIndex,
+        timerEncoder.schemaForValueRow)
     columnFamilySchemas.put(stateName, colFamilySchema)
+    columnFamilySchemas.put(secIndexColFamilyName, secIndexColFamilySchema)
     val stateVariableInfo = TransformWithStateVariableUtils.getTimerState(stateName)
     stateVariableInfos.put(stateName, stateVariableInfo)
   }
@@ -404,6 +417,9 @@ class DriverStatefulProcessorHandleImpl(
     val ttlEnabled = if (ttlConfig.ttlDuration != null && ttlConfig.ttlDuration.isZero) {
       false
     } else {
+      val ttlColFamilyName = getTtlColFamilyName(stateName)
+      val ttlColFamilySchema = schemaUtils.getTtlStateSchema(ttlColFamilyName, keyExprEnc)
+      columnFamilySchemas.put(ttlColFamilyName, ttlColFamilySchema)
       true
     }
 
@@ -432,6 +448,9 @@ class DriverStatefulProcessorHandleImpl(
     val ttlEnabled = if (ttlConfig.ttlDuration != null && ttlConfig.ttlDuration.isZero) {
       false
     } else {
+      val ttlColFamilyName = getTtlColFamilyName(stateName)
+      val ttlColFamilySchema = schemaUtils.getTtlStateSchema(ttlColFamilyName, keyExprEnc)
+      columnFamilySchemas.put(ttlColFamilyName, ttlColFamilySchema)
       true
     }
 
@@ -459,14 +478,19 @@ class DriverStatefulProcessorHandleImpl(
       ttlConfig: TTLConfig): MapState[K, V] = {
     verifyStateVarOperations("get_map_state", PRE_INIT)
 
+    val userKeyEnc = encoderFor[K]
+    val valEncoder = encoderFor[V]
     val ttlEnabled = if (ttlConfig.ttlDuration != null && ttlConfig.ttlDuration.isZero) {
       false
     } else {
+      val ttlColFamilyName = getTtlColFamilyName(stateName)
+      val ttlColFamilySchema = schemaUtils.getTtlStateSchema(
+        ttlColFamilyName, keyExprEnc, userKeyEnc.schema)
+      columnFamilySchemas.put(ttlColFamilyName, ttlColFamilySchema)
       true
     }
 
-    val userKeyEnc = encoderFor[K]
-    val valEncoder = encoderFor[V]
+
     val colFamilySchema = schemaUtils.
       getMapStateSchema(stateName, keyExprEnc, userKeyEnc, valEncoder, ttlEnabled)
     columnFamilySchemas.put(stateName, colFamilySchema)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TTLState.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TTLState.scala
index 87d1a15dff1a9..e4e45fbd74bbc 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TTLState.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TTLState.scala
@@ -21,8 +21,9 @@ import java.time.Duration
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions.{UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.execution.streaming.StateStoreColumnFamilySchemaUtils.getTtlColFamilyName
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{RangeKeyScanStateEncoderSpec, StateStore}
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, RangeKeyScanStateEncoderSpec, StateStore}
 import org.apache.spark.sql.types._
 
 object StateTTLSchema {
@@ -79,12 +80,13 @@ abstract class SingleKeyTTLStateImpl(
     stateName: String,
     store: StateStore,
     keyExprEnc: ExpressionEncoder[Any],
-    ttlExpirationMs: Long)
+    ttlExpirationMs: Long,
+    avroEnc: Option[AvroEncoderSpec] = None)
   extends TTLState {
 
   import org.apache.spark.sql.execution.streaming.StateTTLSchema._
 
-  private val ttlColumnFamilyName = "$ttl_" + stateName
+  private val ttlColumnFamilyName = getTtlColFamilyName(stateName)
   private val keySchema = getSingleKeyTTLRowSchema(keyExprEnc.schema)
   private val keyTTLRowEncoder = new SingleKeyTTLEncoder(keyExprEnc)
 
@@ -93,7 +95,7 @@ abstract class SingleKeyTTLStateImpl(
     UnsafeProjection.create(Array[DataType](NullType)).apply(InternalRow.apply(null))
 
   store.createColFamilyIfAbsent(ttlColumnFamilyName, keySchema, TTL_VALUE_ROW_SCHEMA,
-    RangeKeyScanStateEncoderSpec(keySchema, Seq(0)), isInternal = true)
+    RangeKeyScanStateEncoderSpec(keySchema, Seq(0)), isInternal = true, avroEncoderSpec = avroEnc)
 
   /**
    * This function will be called when clear() on State Variables
@@ -199,12 +201,13 @@ abstract class CompositeKeyTTLStateImpl[K](
     store: StateStore,
     keyExprEnc: ExpressionEncoder[Any],
     userKeyEncoder: ExpressionEncoder[Any],
-    ttlExpirationMs: Long)
+    ttlExpirationMs: Long,
+    avroEnc: Option[AvroEncoderSpec] = None)
   extends TTLState {
 
   import org.apache.spark.sql.execution.streaming.StateTTLSchema._
 
-  private val ttlColumnFamilyName = "$ttl_" + stateName
+  private val ttlColumnFamilyName = getTtlColFamilyName(stateName)
   private val keySchema = getCompositeKeyTTLRowSchema(
     keyExprEnc.schema, userKeyEncoder.schema
   )
@@ -218,7 +221,7 @@ abstract class CompositeKeyTTLStateImpl[K](
 
   store.createColFamilyIfAbsent(ttlColumnFamilyName, keySchema,
     TTL_VALUE_ROW_SCHEMA, RangeKeyScanStateEncoderSpec(keySchema,
-      Seq(0)), isInternal = true)
+      Seq(0)), isInternal = true, avroEncoderSpec = avroEnc)
 
   def clearTTLState(): Unit = {
     val iterator = store.iterator(ttlColumnFamilyName)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TimerStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TimerStateImpl.scala
index d0fbaf6600609..5459e65526f5f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TimerStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TimerStateImpl.scala
@@ -43,6 +43,15 @@ object TimerStateUtils {
       TimerStateUtils.PROC_TIMERS_STATE_NAME + TimerStateUtils.KEY_TO_TIMESTAMP_CF
     }
   }
+
+  def getSecIndexColFamilyName(timeMode: String): String = {
+    assert(timeMode == TimeMode.EventTime.toString || timeMode == TimeMode.ProcessingTime.toString)
+    if (timeMode == TimeMode.EventTime.toString) {
+      TimerStateUtils.EVENT_TIMERS_STATE_NAME + TimerStateUtils.TIMESTAMP_TO_KEY_CF
+    } else {
+      TimerStateUtils.PROC_TIMERS_STATE_NAME + TimerStateUtils.TIMESTAMP_TO_KEY_CF
+    }
+  }
 }
 
 /**
@@ -55,7 +64,9 @@ object TimerStateUtils {
 class TimerStateImpl(
     store: StateStore,
     timeMode: TimeMode,
-    keyExprEnc: ExpressionEncoder[Any]) extends Logging {
+    keyExprEnc: ExpressionEncoder[Any],
+    avroEnc: Option[AvroEncoderSpec] = None,
+    secIndexAvroEnc: Option[AvroEncoderSpec] = None) extends Logging {
 
   private val EMPTY_ROW =
     UnsafeProjection.create(Array[DataType](NullType)).apply(InternalRow.apply(null))
@@ -75,7 +86,7 @@ class TimerStateImpl(
   private val keyToTsCFName = timerCFName + TimerStateUtils.KEY_TO_TIMESTAMP_CF
   store.createColFamilyIfAbsent(keyToTsCFName, schemaForKeyRow,
     schemaForValueRow, PrefixKeyScanStateEncoderSpec(schemaForKeyRow, 1),
-    useMultipleValuesPerKey = false, isInternal = true)
+    useMultipleValuesPerKey = false, isInternal = true, avroEncoderSpec = avroEnc)
 
   // We maintain a secondary index that inverts the ordering of the timestamp
   // and grouping key
@@ -83,7 +94,7 @@ class TimerStateImpl(
   private val tsToKeyCFName = timerCFName + TimerStateUtils.TIMESTAMP_TO_KEY_CF
   store.createColFamilyIfAbsent(tsToKeyCFName, keySchemaForSecIndex,
     schemaForValueRow, RangeKeyScanStateEncoderSpec(keySchemaForSecIndex, Seq(0)),
-    useMultipleValuesPerKey = false, isInternal = true)
+    useMultipleValuesPerKey = false, isInternal = true, avroEncoderSpec = secIndexAvroEnc)
 
   private def getGroupingKey(cfName: String): Any = {
     val keyOption = ImplicitGroupingKeyTracker.getImplicitKeyOption
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala
index 60eea5842645e..3ab9e5d226f23 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala
@@ -20,7 +20,7 @@ import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, StateStore}
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, NoPrefixKeyStateEncoderSpec, StateStore}
 import org.apache.spark.sql.streaming.{TTLConfig, ValueState}
 
 /**
@@ -34,6 +34,10 @@ import org.apache.spark.sql.streaming.{TTLConfig, ValueState}
  * @param ttlConfig  - TTL configuration for values  stored in this state
  * @param batchTimestampMs - current batch processing timestamp.
  * @param metrics - metrics to be updated as part of stateful processing
+ * @param avroEnc - optional Avro serializer and deserializer for this state variable that
+ *                is used by the StateStore to encode state in Avro format
+ * @param ttlAvroEnc - optional Avro serializer and deserializer for TTL state that
+ *                is used by the StateStore to encode state in Avro format
  * @tparam S - data type of object that will be stored
  */
 class ValueStateImplWithTTL[S](
@@ -43,9 +47,11 @@ class ValueStateImplWithTTL[S](
     valEncoder: ExpressionEncoder[Any],
     ttlConfig: TTLConfig,
     batchTimestampMs: Long,
-    metrics: Map[String, SQLMetric] = Map.empty)
+    metrics: Map[String, SQLMetric] = Map.empty,
+    avroEnc: Option[AvroEncoderSpec] = None,
+    ttlAvroEnc: Option[AvroEncoderSpec] = None)
   extends SingleKeyTTLStateImpl(
-    stateName, store, keyExprEnc, batchTimestampMs) with ValueState[S] {
+    stateName, store, keyExprEnc, batchTimestampMs, ttlAvroEnc) with ValueState[S] {
 
   private val stateTypesEncoder = StateTypesEncoder(keyExprEnc, valEncoder,
     stateName, hasTtl = true)
@@ -57,7 +63,7 @@ class ValueStateImplWithTTL[S](
   private def initialize(): Unit = {
     store.createColFamilyIfAbsent(stateName,
       keyExprEnc.schema, getValueSchemaWithTTL(valEncoder.schema, true),
-      NoPrefixKeyStateEncoderSpec(keyExprEnc.schema))
+      NoPrefixKeyStateEncoderSpec(keyExprEnc.schema), avroEncoderSpec = avroEnc)
   }
 
   /** Function to check if state exists. Returns true if present and false otherwise */
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
index 7fe7da2ff290a..ee5490daf315c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
@@ -31,6 +31,7 @@ import org.apache.spark.sql.avro.{AvroDeserializer, AvroSerializer, SchemaConver
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{BoundReference, JoinedRow, UnsafeProjection, UnsafeRow}
 import org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter
+import org.apache.spark.sql.execution.streaming.StateStoreColumnFamilySchemaUtils
 import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider.{STATE_ENCODING_NUM_VERSION_BYTES, STATE_ENCODING_VERSION, VIRTUAL_COL_FAMILY_PREFIX_BYTES}
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.Platform
@@ -237,6 +238,7 @@ class PrefixKeyScanStateEncoder(
 
   import RocksDBStateEncoder._
 
+  private val usingAvroEncoding = avroEnc.isDefined
   private val prefixKeyFieldsWithIdx: Seq[(StructField, Int)] = {
     keySchema.zipWithIndex.take(numColsPrefixKey)
   }
@@ -256,6 +258,18 @@ class PrefixKeyScanStateEncoder(
     UnsafeProjection.create(refs)
   }
 
+  // Prefix Key schema and projection definitions used by the Avro Serializers
+  // and Deserializers
+  private val prefixKeySchema = StructType(keySchema.take(numColsPrefixKey))
+  private val prefixKeyAvroType = SchemaConverters.toAvroType(prefixKeySchema)
+  private val prefixKeyProj = UnsafeProjection.create(prefixKeySchema)
+
+  // Remaining Key schema and projection definitions used by the Avro Serializers
+  // and Deserializers
+  private val remainingKeySchema = StructType(keySchema.drop(numColsPrefixKey))
+  private val remainingKeyAvroType = SchemaConverters.toAvroType(remainingKeySchema)
+  private val remainingKeyProj = UnsafeProjection.create(remainingKeySchema)
+
   // This is quite simple to do - just bind sequentially, as we don't change the order.
   private val restoreKeyProjection: UnsafeProjection = UnsafeProjection.create(keySchema)
 
@@ -263,9 +277,24 @@ class PrefixKeyScanStateEncoder(
   private val joinedRowOnKey = new JoinedRow()
 
   override def encodeKey(row: UnsafeRow): Array[Byte] = {
-    val prefixKeyEncoded = encodeUnsafeRow(extractPrefixKey(row))
-    val remainingEncoded = encodeUnsafeRow(remainingKeyProjection(row))
-
+    val (prefixKeyEncoded, remainingEncoded) = if (usingAvroEncoding) {
+      (
+        encodeUnsafeRow(
+          extractPrefixKey(row),
+          avroEnc.get.keySerializer,
+          prefixKeyAvroType,
+          out
+        ),
+        encodeUnsafeRow(
+          remainingKeyProjection(row),
+          avroEnc.get.suffixKeySerializer.get,
+          remainingKeyAvroType,
+          out
+        )
+      )
+    } else {
+      (encodeUnsafeRow(extractPrefixKey(row)), encodeUnsafeRow(remainingKeyProjection(row)))
+    }
     val (encodedBytes, startingOffset) = encodeColumnFamilyPrefix(
       prefixKeyEncoded.length + remainingEncoded.length + 4
     )
@@ -296,9 +325,25 @@ class PrefixKeyScanStateEncoder(
     Platform.copyMemory(keyBytes, decodeKeyStartOffset + 4 + prefixKeyEncodedLen,
       remainingKeyEncoded, Platform.BYTE_ARRAY_OFFSET, remainingKeyEncodedLen)
 
-    val prefixKeyDecoded = decodeToUnsafeRow(prefixKeyEncoded, numFields = numColsPrefixKey)
-    val remainingKeyDecoded = decodeToUnsafeRow(remainingKeyEncoded,
-      numFields = keySchema.length - numColsPrefixKey)
+    val (prefixKeyDecoded, remainingKeyDecoded) = if (usingAvroEncoding) {
+      (
+        decodeToUnsafeRow(
+          prefixKeyEncoded,
+          avroEnc.get.keyDeserializer,
+          prefixKeyAvroType,
+          prefixKeyProj
+        ),
+        decodeToUnsafeRow(
+          remainingKeyEncoded,
+          avroEnc.get.suffixKeyDeserializer.get,
+          remainingKeyAvroType,
+          remainingKeyProj
+        )
+      )
+    } else {
+      (decodeToUnsafeRow(prefixKeyEncoded, numFields = numColsPrefixKey),
+        decodeToUnsafeRow(remainingKeyEncoded, numFields = keySchema.length - numColsPrefixKey))
+    }
 
     restoreKeyProjection(joinedRowOnKey.withLeft(prefixKeyDecoded).withRight(remainingKeyDecoded))
   }
@@ -308,7 +353,11 @@ class PrefixKeyScanStateEncoder(
   }
 
   override def encodePrefixKey(prefixKey: UnsafeRow): Array[Byte] = {
-    val prefixKeyEncoded = encodeUnsafeRow(prefixKey)
+    val prefixKeyEncoded = if (usingAvroEncoding) {
+      encodeUnsafeRow(prefixKey, avroEnc.get.keySerializer, prefixKeyAvroType, out)
+    } else {
+      encodeUnsafeRow(prefixKey)
+    }
     val (prefix, startingOffset) = encodeColumnFamilyPrefix(
       prefixKeyEncoded.length + 4
     )
@@ -361,7 +410,7 @@ class RangeKeyScanStateEncoder(
     useColumnFamilies: Boolean = false,
     virtualColFamilyId: Option[Short] = None,
     avroEnc: Option[AvroEncoderSpec] = None)
-  extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) {
+  extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) with Logging {
 
   import RocksDBStateEncoder._
 
@@ -430,6 +479,22 @@ class RangeKeyScanStateEncoder(
     UnsafeProjection.create(refs)
   }
 
+  private val rangeScanAvroSchema = StateStoreColumnFamilySchemaUtils.convertForRangeScan(
+    StructType(rangeScanKeyFieldsWithOrdinal.map(_._1).toArray))
+
+  private val rangeScanAvroType = SchemaConverters.toAvroType(rangeScanAvroSchema)
+
+  private val rangeScanAvroProjection = UnsafeProjection.create(rangeScanAvroSchema)
+
+  // Existing remainder key schema stuff
+  private val remainingKeySchema = StructType(
+    0.to(keySchema.length - 1).diff(orderingOrdinals).map(keySchema(_))
+  )
+
+  private val remainingKeyAvroType = SchemaConverters.toAvroType(remainingKeySchema)
+
+  private val remainingKeyAvroProjection = UnsafeProjection.create(remainingKeySchema)
+
   // Reusable objects
   private val joinedRowOnKey = new JoinedRow()
 
@@ -622,10 +687,28 @@ class RangeKeyScanStateEncoder(
   override def encodeKey(row: UnsafeRow): Array[Byte] = {
     // This prefix key has the columns specified by orderingOrdinals
     val prefixKey = extractPrefixKey(row)
-    val rangeScanKeyEncoded = encodeUnsafeRow(encodePrefixKeyForRangeScan(prefixKey))
+    val rangeScanKeyEncoded = if (avroEnc.isDefined) {
+      encodeUnsafeRow(
+        encodePrefixKeyForRangeScan(prefixKey),
+        avroEnc.get.keySerializer,
+        rangeScanAvroType,
+        out
+      )
+    } else {
+      encodeUnsafeRow(encodePrefixKeyForRangeScan(prefixKey))
+    }
 
     val result = if (orderingOrdinals.length < keySchema.length) {
-      val remainingEncoded = encodeUnsafeRow(remainingKeyProjection(row))
+      val remainingEncoded = if (avroEnc.isDefined) {
+        encodeUnsafeRow(
+          remainingKeyProjection(row),
+          avroEnc.get.suffixKeySerializer.get,
+          remainingKeyAvroType,
+          out
+        )
+      } else {
+        encodeUnsafeRow(remainingKeyProjection(row))
+      }
       val (encodedBytes, startingOffset) = encodeColumnFamilyPrefix(
         rangeScanKeyEncoded.length + remainingEncoded.length + 4
       )
@@ -662,8 +745,13 @@ class RangeKeyScanStateEncoder(
     Platform.copyMemory(keyBytes, decodeKeyStartOffset + 4,
       prefixKeyEncoded, Platform.BYTE_ARRAY_OFFSET, prefixKeyEncodedLen)
 
-    val prefixKeyDecodedForRangeScan = decodeToUnsafeRow(prefixKeyEncoded,
-      numFields = orderingOrdinals.length)
+    val prefixKeyDecodedForRangeScan = if (avroEnc.isDefined) {
+      decodeToUnsafeRow(prefixKeyEncoded, avroEnc.get.keyDeserializer,
+        rangeScanAvroType, rangeScanAvroProjection)
+    } else {
+      decodeToUnsafeRow(prefixKeyEncoded,
+        numFields = orderingOrdinals.length)
+    }
     val prefixKeyDecoded = decodePrefixKeyForRangeScan(prefixKeyDecodedForRangeScan)
 
     if (orderingOrdinals.length < keySchema.length) {
@@ -676,8 +764,14 @@ class RangeKeyScanStateEncoder(
         remainingKeyEncoded, Platform.BYTE_ARRAY_OFFSET,
         remainingKeyEncodedLen)
 
-      val remainingKeyDecoded = decodeToUnsafeRow(remainingKeyEncoded,
-        numFields = keySchema.length - orderingOrdinals.length)
+      val remainingKeyDecoded = if (avroEnc.isDefined) {
+        decodeToUnsafeRow(remainingKeyEncoded,
+          avroEnc.get.suffixKeyDeserializer.get,
+          remainingKeyAvroType, remainingKeyAvroProjection)
+      } else {
+        decodeToUnsafeRow(remainingKeyEncoded,
+          numFields = keySchema.length - orderingOrdinals.length)
+      }
 
       val joined = joinedRowOnKey.withLeft(prefixKeyDecoded).withRight(remainingKeyDecoded)
       val restored = restoreKeyProjection(joined)
@@ -690,7 +784,16 @@ class RangeKeyScanStateEncoder(
   }
 
   override def encodePrefixKey(prefixKey: UnsafeRow): Array[Byte] = {
-    val rangeScanKeyEncoded = encodeUnsafeRow(encodePrefixKeyForRangeScan(prefixKey))
+    val rangeScanKeyEncoded = if (avroEnc.isDefined) {
+      encodeUnsafeRow(
+        encodePrefixKeyForRangeScan(prefixKey),
+        avroEnc.get.keySerializer,
+        rangeScanAvroType,
+        out
+      )
+    } else {
+      encodeUnsafeRow(encodePrefixKeyForRangeScan(prefixKey))
+    }
     val (prefix, startingOffset) = encodeColumnFamilyPrefix(rangeScanKeyEncoded.length + 4)
 
     Platform.putInt(prefix, startingOffset, rangeScanKeyEncoded.length)
@@ -728,6 +831,7 @@ class NoPrefixKeyStateEncoder(
   private val usingAvroEncoding = avroEnc.isDefined
   private val keyRow = new UnsafeRow(keySchema.size)
   private val keyAvroType = SchemaConverters.toAvroType(keySchema)
+  private val keyProj = UnsafeProjection.create(keySchema)
 
   override def encodeKey(row: UnsafeRow): Array[Byte] = {
     if (!useColumnFamilies) {
@@ -761,11 +865,25 @@ class NoPrefixKeyStateEncoder(
     if (useColumnFamilies) {
       if (keyBytes != null) {
         // Platform.BYTE_ARRAY_OFFSET is the recommended way refer to the 1st offset. See Platform.
-        keyRow.pointTo(
-          keyBytes,
-          decodeKeyStartOffset + STATE_ENCODING_NUM_VERSION_BYTES,
-          keyBytes.length - STATE_ENCODING_NUM_VERSION_BYTES - VIRTUAL_COL_FAMILY_PREFIX_BYTES)
-        keyRow
+        if (usingAvroEncoding) {
+          val avroBytes = new Array[Byte](
+            keyBytes.length - STATE_ENCODING_NUM_VERSION_BYTES - VIRTUAL_COL_FAMILY_PREFIX_BYTES)
+          System.arraycopy(
+            keyBytes,
+            decodeKeyStartOffset + STATE_ENCODING_NUM_VERSION_BYTES,
+            avroBytes,
+            0,
+            avroBytes.length
+          )
+          decodeToUnsafeRow(
+            keyBytes, avroEnc.get.keyDeserializer, keyAvroType, keyProj)
+        } else {
+          keyRow.pointTo(
+            keyBytes,
+            decodeKeyStartOffset + STATE_ENCODING_NUM_VERSION_BYTES,
+            keyBytes.length - STATE_ENCODING_NUM_VERSION_BYTES - VIRTUAL_COL_FAMILY_PREFIX_BYTES)
+          keyRow
+        }
       } else {
         null
       }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
index ff07666ce0fe9..a05b452bd0184 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
@@ -44,7 +44,9 @@ case class AvroEncoderSpec(
   keySerializer: AvroSerializer,
   keyDeserializer: AvroDeserializer,
   valueSerializer: AvroSerializer,
-  valueDeserializer: AvroDeserializer
+  valueDeserializer: AvroDeserializer,
+  suffixKeySerializer: Option[AvroSerializer] = None,
+  suffixKeyDeserializer: Option[AvroDeserializer] = None
 ) extends Serializable
 
 // Used to represent the schema of a column family in the state store
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala
index 409a255ae3e64..ebd29bff5d354 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala
@@ -105,7 +105,7 @@ class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
 
   override def getStateTTLMetricName: String = "numListStateWithTTLVars"
 
-  test("verify iterator works with expired values in beginning of list") {
+  testWithEncodingTypes("verify iterator works with expired values in beginning of list") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
@@ -195,7 +195,7 @@ class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
   // ascending order of TTL by stopping the query, setting the new TTL, and restarting
   // the query to check that the expired elements in the middle or end of the list
   // are not returned.
-  test("verify iterator works with expired values in middle of list") {
+  testWithEncodingTypes("verify iterator works with expired values in middle of list") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
@@ -343,7 +343,7 @@ class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
     }
   }
 
-  test("verify iterator works with expired values in end of list") {
+  testWithEncodingTypes("verify iterator works with expired values in end of list") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithMapStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithMapStateSuite.scala
index 76c5cbeee424b..63609ba96625c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithMapStateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithMapStateSuite.scala
@@ -110,7 +110,7 @@ class TransformWithMapStateSuite extends StreamTest
     }
   }
 
-  test("Test retrieving value with non-existing user key") {
+  testWithEncodingTypes("Test retrieving value with non-existing user key") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
 
@@ -134,7 +134,7 @@ class TransformWithMapStateSuite extends StreamTest
     }
   }
 
-  test("Test put value with null value") {
+  testWithEncodingTypes("Test put value with null value") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
 
@@ -158,7 +158,7 @@ class TransformWithMapStateSuite extends StreamTest
     }
   }
 
-  test("Test map state correctness") {
+  testWithEncodingTypes("Test map state correctness") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
       val inputData = MemoryStream[InputMapRow]
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithMapStateTTLSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithMapStateTTLSuite.scala
index 022280eb3bcef..a68632534c001 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithMapStateTTLSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithMapStateTTLSuite.scala
@@ -182,7 +182,7 @@ class TransformWithMapStateTTLSuite extends TransformWithStateTTLTest {
 
   override def getStateTTLMetricName: String = "numMapStateWithTTLVars"
 
-  test("validate state is evicted with multiple user keys") {
+  testWithEncodingTypes("validate state is evicted with multiple user keys") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
@@ -224,7 +224,7 @@ class TransformWithMapStateTTLSuite extends TransformWithStateTTLTest {
     }
   }
 
-  test("verify iterator doesn't return expired keys") {
+  testWithEncodingTypes("verify iterator doesn't return expired keys") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateTTLTest.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateTTLTest.scala
index 2ddf69aa49e04..e3b0a6b811742 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateTTLTest.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateTTLTest.scala
@@ -21,7 +21,7 @@ import java.sql.Timestamp
 import java.time.Duration
 
 import org.apache.spark.sql.execution.streaming.MemoryStream
-import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.execution.streaming.state.{AlsoTestWithChangelogCheckpointingEnabled, RocksDBStateStoreProvider}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.streaming.util.StreamManualClock
 
@@ -41,14 +41,14 @@ case class OutputEvent(
  * Test suite base for TransformWithState with TTL support.
  */
 abstract class TransformWithStateTTLTest
-  extends StreamTest {
+  extends StreamTest with AlsoTestWithChangelogCheckpointingEnabled {
   import testImplicits._
 
   def getProcessor(ttlConfig: TTLConfig): StatefulProcessor[String, InputEvent, OutputEvent]
 
   def getStateTTLMetricName: String
 
-  test("validate state is evicted at ttl expiry") {
+  testWithEncodingTypes("validate state is evicted at ttl expiry") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
       withTempDir { dir =>
@@ -125,7 +125,7 @@ abstract class TransformWithStateTTLTest
     }
   }
 
-  test("validate state update updates the expiration timestamp") {
+  testWithEncodingTypes("validate state update updates the expiration timestamp") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
       val inputStream = MemoryStream[InputEvent]
@@ -187,7 +187,7 @@ abstract class TransformWithStateTTLTest
     }
   }
 
-  test("validate state is evicted at ttl expiry for no data batch") {
+  testWithEncodingTypes("validate state is evicted at ttl expiry for no data batch") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
     classOf[RocksDBStateStoreProvider].getName) {
       val inputStream = MemoryStream[InputEvent]
@@ -238,7 +238,7 @@ abstract class TransformWithStateTTLTest
     }
   }
 
-  test("validate only expired keys are removed from the state") {
+  testWithEncodingTypes("validate only expired keys are removed from the state") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala
index 21c3beb79314c..403eb8c48ed67 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala
@@ -195,7 +195,7 @@ class TransformWithValueStateTTLSuite extends TransformWithStateTTLTest {
 
   override def getStateTTLMetricName: String = "numValueStateWithTTLVars"
 
-  test("validate multiple value states") {
+  testWithEncodingTypes("validate multiple value states") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
       val ttlKey = "k1"

From 448ea760696ab87687447864516ee30fe8ab2ebb Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Thu, 7 Nov 2024 12:31:55 -0800
Subject: [PATCH 16/30] making schema conversion lazy

---
 .../streaming/state/RocksDBStateEncoder.scala      | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
index ee5490daf315c..8f177745273a2 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
@@ -261,13 +261,13 @@ class PrefixKeyScanStateEncoder(
   // Prefix Key schema and projection definitions used by the Avro Serializers
   // and Deserializers
   private val prefixKeySchema = StructType(keySchema.take(numColsPrefixKey))
-  private val prefixKeyAvroType = SchemaConverters.toAvroType(prefixKeySchema)
+  private lazy val prefixKeyAvroType = SchemaConverters.toAvroType(prefixKeySchema)
   private val prefixKeyProj = UnsafeProjection.create(prefixKeySchema)
 
   // Remaining Key schema and projection definitions used by the Avro Serializers
   // and Deserializers
   private val remainingKeySchema = StructType(keySchema.drop(numColsPrefixKey))
-  private val remainingKeyAvroType = SchemaConverters.toAvroType(remainingKeySchema)
+  private lazy val remainingKeyAvroType = SchemaConverters.toAvroType(remainingKeySchema)
   private val remainingKeyProj = UnsafeProjection.create(remainingKeySchema)
 
   // This is quite simple to do - just bind sequentially, as we don't change the order.
@@ -482,7 +482,7 @@ class RangeKeyScanStateEncoder(
   private val rangeScanAvroSchema = StateStoreColumnFamilySchemaUtils.convertForRangeScan(
     StructType(rangeScanKeyFieldsWithOrdinal.map(_._1).toArray))
 
-  private val rangeScanAvroType = SchemaConverters.toAvroType(rangeScanAvroSchema)
+  private lazy val rangeScanAvroType = SchemaConverters.toAvroType(rangeScanAvroSchema)
 
   private val rangeScanAvroProjection = UnsafeProjection.create(rangeScanAvroSchema)
 
@@ -491,7 +491,7 @@ class RangeKeyScanStateEncoder(
     0.to(keySchema.length - 1).diff(orderingOrdinals).map(keySchema(_))
   )
 
-  private val remainingKeyAvroType = SchemaConverters.toAvroType(remainingKeySchema)
+  private lazy val remainingKeyAvroType = SchemaConverters.toAvroType(remainingKeySchema)
 
   private val remainingKeyAvroProjection = UnsafeProjection.create(remainingKeySchema)
 
@@ -830,7 +830,7 @@ class NoPrefixKeyStateEncoder(
   // Reusable objects
   private val usingAvroEncoding = avroEnc.isDefined
   private val keyRow = new UnsafeRow(keySchema.size)
-  private val keyAvroType = SchemaConverters.toAvroType(keySchema)
+  private lazy val keyAvroType = SchemaConverters.toAvroType(keySchema)
   private val keyProj = UnsafeProjection.create(keySchema)
 
   override def encodeKey(row: UnsafeRow): Array[Byte] = {
@@ -922,7 +922,7 @@ class MultiValuedStateEncoder(
   // Reusable objects
   private val out = new ByteArrayOutputStream
   private val valueRow = new UnsafeRow(valueSchema.size)
-  private val valueAvroType = SchemaConverters.toAvroType(valueSchema)
+  private lazy val valueAvroType = SchemaConverters.toAvroType(valueSchema)
   private val valueProj = UnsafeProjection.create(valueSchema)
 
   override def encodeValue(row: UnsafeRow): Array[Byte] = {
@@ -1018,7 +1018,7 @@ class SingleValueStateEncoder(
   // Reusable objects
   private val out = new ByteArrayOutputStream
   private val valueRow = new UnsafeRow(valueSchema.size)
-  private val valueAvroType = SchemaConverters.toAvroType(valueSchema)
+  private lazy val valueAvroType = SchemaConverters.toAvroType(valueSchema)
   private val valueProj = UnsafeProjection.create(valueSchema)
 
   override def encodeValue(row: UnsafeRow): Array[Byte] = {

From 386fbf144f4f539445a0c4203bd6e4a65a8483da Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Thu, 7 Nov 2024 12:42:27 -0800
Subject: [PATCH 17/30] batch succeeds

---
 .../spark/sql/execution/streaming/statefulOperators.scala       | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
index 0cf641c703d6a..02048ee7ce682 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
@@ -112,7 +112,7 @@ trait StatefulOperator extends SparkPlan {
   }
 
   lazy val stateStoreEncoding: String =
-    session.sessionState.conf.getConf(
+    conf.getConf(
       SQLConf.STREAMING_STATE_STORE_ENCODING_FORMAT)
 
   def metadataFilePath(): Path = {

From 896e24f27201cbbacc1e06ce5616bb060aead7d4 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Thu, 7 Nov 2024 12:52:18 -0800
Subject: [PATCH 18/30] actually enabling ttl

---
 .../streaming/StatefulProcessorHandleImpl.scala          | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
index 0258691e51d66..32e5d6d4946b0 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
@@ -237,7 +237,8 @@ class StatefulProcessorHandleImpl(
       validateTTLConfig(ttlConfig, stateName)
       assert(batchTimestampMs.isDefined)
       val valueStateWithTTL = new ValueStateImplWithTTL[T](store, stateName,
-        keyEncoder, stateEncoder, ttlConfig, batchTimestampMs.get, metrics)
+        keyEncoder, stateEncoder, ttlConfig, batchTimestampMs.get, metrics,
+        schemas(stateName).avroEnc, schemas(getTtlColFamilyName(stateName)).avroEnc)
       ttlStates.add(valueStateWithTTL)
       TWSMetricsUtils.incrementMetric(metrics, "numValueStateWithTTLVars")
       valueStateWithTTL
@@ -286,7 +287,8 @@ class StatefulProcessorHandleImpl(
       validateTTLConfig(ttlConfig, stateName)
       assert(batchTimestampMs.isDefined)
       val listStateWithTTL = new ListStateImplWithTTL[T](store, stateName,
-        keyEncoder, stateEncoder, ttlConfig, batchTimestampMs.get, metrics)
+        keyEncoder, stateEncoder, ttlConfig, batchTimestampMs.get, metrics,
+        schemas(stateName).avroEnc, schemas(getTtlColFamilyName(stateName)).avroEnc)
       TWSMetricsUtils.incrementMetric(metrics, "numListStateWithTTLVars")
       ttlStates.add(listStateWithTTL)
       listStateWithTTL
@@ -324,7 +326,8 @@ class StatefulProcessorHandleImpl(
       validateTTLConfig(ttlConfig, stateName)
       assert(batchTimestampMs.isDefined)
       val mapStateWithTTL = new MapStateImplWithTTL[K, V](store, stateName, keyEncoder, userKeyEnc,
-        valEncoder, ttlConfig, batchTimestampMs.get, metrics)
+        valEncoder, ttlConfig, batchTimestampMs.get, metrics,
+        schemas(stateName).avroEnc, schemas(getTtlColFamilyName(stateName)).avroEnc)
       TWSMetricsUtils.incrementMetric(metrics, "numMapStateWithTTLVars")
       ttlStates.add(mapStateWithTTL)
       mapStateWithTTL

From 15c5f71b3db387dd3b06f4958a2c95dfde2124a1 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Thu, 7 Nov 2024 13:00:57 -0800
Subject: [PATCH 19/30] including hidden files

---
 .../TransformWithValueStateTTLSuite.scala     | 116 ++++++++++++------
 1 file changed, 78 insertions(+), 38 deletions(-)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala
index 403eb8c48ed67..e4e03b5da0515 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala
@@ -23,7 +23,7 @@ import org.apache.hadoop.fs.Path
 
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.Encoders
-import org.apache.spark.sql.execution.streaming.{CheckpointFileManager, ListStateImplWithTTL, MapStateImplWithTTL, MemoryStream, TimerStateUtils, ValueStateImpl, ValueStateImplWithTTL}
+import org.apache.spark.sql.execution.streaming.{CheckpointFileManager, ListStateImplWithTTL, MapStateImplWithTTL, MemoryStream, ValueStateImpl, ValueStateImplWithTTL}
 import org.apache.spark.sql.execution.streaming.state._
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.streaming.util.StreamManualClock
@@ -275,60 +275,98 @@ class TransformWithValueStateTTLSuite extends TransformWithStateTTLTest {
         val fm = CheckpointFileManager.create(stateSchemaPath, hadoopConf)
 
         val keySchema = new StructType().add("value", StringType)
-        val schemaForKeyRow: StructType = new StructType()
-          .add("key", new StructType(keySchema.fields))
+        val timerKeyStruct = new StructType(keySchema.fields)
+        val schemaForTimerKeyRow: StructType = new StructType()
+          .add("key", timerKeyStruct)
           .add("expiryTimestampMs", LongType, nullable = false)
-        val schemaForValueRow: StructType = StructType(Array(StructField("__dummy__", NullType)))
+        val schemaForTimerValueRow: StructType =
+          StructType(Array(StructField("__dummy__", NullType)))
+
+        // Timer schemas
         val schema0 = StateStoreColFamilySchema(
-          TimerStateUtils.getTimerStateVarName(TimeMode.ProcessingTime().toString),
-          schemaForKeyRow,
-          schemaForValueRow,
-          Some(PrefixKeyScanStateEncoderSpec(schemaForKeyRow, 1)))
+          "$procTimers_keyToTimestamp",
+          schemaForTimerKeyRow,
+          schemaForTimerValueRow,
+          Some(PrefixKeyScanStateEncoderSpec(schemaForTimerKeyRow, 1)))
+
+        val schemaForTimerReverseKeyRow: StructType = new StructType()
+          .add("expiryTimestampMs", LongType, nullable = false)
+          .add("key", timerKeyStruct)
         val schema1 = StateStoreColFamilySchema(
-          "valueStateTTL",
-          keySchema,
-          new StructType().add("value",
-            new StructType()
-              .add("value", IntegerType, false))
-          .add("ttlExpirationMs", LongType),
-          Some(NoPrefixKeyStateEncoderSpec(keySchema)),
-          None
-        )
+          "$procTimers_timestampToKey",
+          schemaForTimerReverseKeyRow,
+          schemaForTimerValueRow,
+          Some(RangeKeyScanStateEncoderSpec(schemaForTimerReverseKeyRow, List(0))))
+
+        // TTL tracking schemas
+        val ttlKeySchema = new StructType()
+          .add("expirationMs", BinaryType)
+          .add("groupingKey", keySchema)
+
         val schema2 = StateStoreColFamilySchema(
-          "valueState",
-          keySchema,
-          new StructType().add("value", IntegerType, false),
-          Some(NoPrefixKeyStateEncoderSpec(keySchema)),
-          None
-        )
+          "$ttl_listState",
+          ttlKeySchema,
+          schemaForTimerValueRow,
+          Some(RangeKeyScanStateEncoderSpec(ttlKeySchema, List(0))))
+
+        val userKeySchema = new StructType()
+          .add("id", IntegerType, false)
+          .add("name", StringType)
+        val ttlMapKeySchema = new StructType()
+          .add("expirationMs", BinaryType)
+          .add("groupingKey", keySchema)
+          .add("userKey", userKeySchema)
+
         val schema3 = StateStoreColFamilySchema(
+          "$ttl_mapState",
+          ttlMapKeySchema,
+          schemaForTimerValueRow,
+          Some(RangeKeyScanStateEncoderSpec(ttlMapKeySchema, List(0))))
+
+        val schema4 = StateStoreColFamilySchema(
+          "$ttl_valueStateTTL",
+          ttlKeySchema,
+          schemaForTimerValueRow,
+          Some(RangeKeyScanStateEncoderSpec(ttlKeySchema, List(0))))
+
+        // Main state schemas
+        val schema5 = StateStoreColFamilySchema(
           "listState",
           keySchema,
-          new StructType().add("value",
-            new StructType()
+          new StructType()
+            .add("value", new StructType()
               .add("id", LongType, false)
               .add("name", StringType))
             .add("ttlExpirationMs", LongType),
-          Some(NoPrefixKeyStateEncoderSpec(keySchema)),
-          None
-        )
+          Some(NoPrefixKeyStateEncoderSpec(keySchema)))
 
-        val userKeySchema = new StructType()
-          .add("id", IntegerType, false)
-          .add("name", StringType)
         val compositeKeySchema = new StructType()
           .add("key", new StructType().add("value", StringType))
           .add("userKey", userKeySchema)
-        val schema4 = StateStoreColFamilySchema(
+        val schema6 = StateStoreColFamilySchema(
           "mapState",
           compositeKeySchema,
-          new StructType().add("value",
-            new StructType()
+          new StructType()
+            .add("value", new StructType()
               .add("value", StringType))
             .add("ttlExpirationMs", LongType),
           Some(PrefixKeyScanStateEncoderSpec(compositeKeySchema, 1)),
-          Option(userKeySchema)
-        )
+          Option(userKeySchema))
+
+        val schema7 = StateStoreColFamilySchema(
+          "valueState",
+          keySchema,
+          new StructType().add("value", IntegerType, false),
+          Some(NoPrefixKeyStateEncoderSpec(keySchema)))
+
+        val schema8 = StateStoreColFamilySchema(
+          "valueStateTTL",
+          keySchema,
+          new StructType()
+            .add("value", new StructType()
+              .add("value", IntegerType, false))
+            .add("ttlExpirationMs", LongType),
+          Some(NoPrefixKeyStateEncoderSpec(keySchema)))
 
         val ttlKey = "k1"
         val noTtlKey = "k2"
@@ -370,9 +408,11 @@ class TransformWithValueStateTTLSuite extends TransformWithStateTTLTest {
               q.lastProgress.stateOperators.head.customMetrics
                 .get("numMapStateWithTTLVars").toInt)
 
-            assert(colFamilySeq.length == 5)
+            // Now expect 9 column families
+            assert(colFamilySeq.length == 9)
             assert(colFamilySeq.map(_.toString).toSet == Set(
-              schema0, schema1, schema2, schema3, schema4
+              schema0, schema1, schema2, schema3, schema4,
+              schema5, schema6, schema7, schema8
             ).map(_.toString))
           },
           StopStream

From 1f5e5f748ace644920b0954fd78e3979c2d2ebbf Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Thu, 7 Nov 2024 13:03:23 -0800
Subject: [PATCH 20/30] testWithEncodingTypes

---
 .../spark/sql/streaming/TransformWithValueStateTTLSuite.scala  | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala
index e4e03b5da0515..91e88eec06cd1 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala
@@ -262,7 +262,8 @@ class TransformWithValueStateTTLSuite extends TransformWithStateTTLTest {
     }
   }
 
-  test("verify StateSchemaV3 writes correct SQL schema of key/value and with TTL") {
+  testWithEncodingTypes(
+    "verify StateSchemaV3 writes correct SQL schema of key/value and with TTL") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key ->

From 1826d5a75e1da240fbbc4fcd956b0c499a3b6995 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Fri, 8 Nov 2024 14:58:52 -0800
Subject: [PATCH 21/30] no longer relying on unsaferow

---
 .../StateStoreColumnFamilySchemaUtils.scala   |  31 +++--
 .../streaming/state/RocksDBStateEncoder.scala | 113 +++++++++++++++---
 2 files changed, 113 insertions(+), 31 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
index 734436f7b2c27..a2c98127b3b91 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
@@ -16,6 +16,7 @@
  */
 package org.apache.spark.sql.execution.streaming
 
+import org.apache.spark.internal.Logging
 import org.apache.spark.sql.Encoder
 import org.apache.spark.sql.avro.{AvroDeserializer, AvroOptions, AvroSerializer, SchemaConverters}
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
@@ -40,12 +41,18 @@ object StateStoreColumnFamilySchemaUtils {
    */
   def convertForRangeScan(schema: StructType, ordinals: Seq[Int] = Seq.empty): StructType = {
     val ordinalSet = ordinals.toSet
-    StructType(schema.fields.zipWithIndex.map { case (field, idx) =>
+
+    StructType(schema.fields.zipWithIndex.flatMap { case (field, idx) =>
       if ((ordinals.isEmpty || ordinalSet.contains(idx)) && isFixedSize(field.dataType)) {
-        // Convert numeric types to BinaryType while preserving nullability
-        field.copy(dataType = BinaryType)
+        // For each numeric field, create two fields:
+        // 1. A boolean for sign (positive = true, negative = false)
+        // 2. The original numeric value in big-endian format
+        Seq(
+          StructField(s"${field.name}_marker", BooleanType, nullable = false),
+          field.copy(name = s"${field.name}_value", BinaryType)
+        )
       } else {
-        field
+        Seq(field)
       }
     })
   }
@@ -67,13 +74,13 @@ object StateStoreColumnFamilySchemaUtils {
  *                            for this state type. This class is used to create the
  *                            StateStoreColumnFamilySchema for each state variable from the driver
  */
-class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) {
+class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) extends Logging {
 
   /**
    * If initializeAvroSerde is true, this method will create an Avro Serializer and Deserializer
    * for a particular key and value schema.
    */
-  private def getAvroSerde(
+  private[sql] def getAvroSerde(
       keySchema: StructType,
       valSchema: StructType,
       suffixKeySchema: Option[StructType] = None
@@ -170,9 +177,9 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) {
       ttlValSchema,
       Some(RangeKeyScanStateEncoderSpec(ttlKeySchema, Seq(0))),
       avroEnc = getAvroSerde(
-        StructType(ttlKeySchema.take(1)),
+        StructType(ttlKeySchema.take(2)),
         ttlValSchema,
-        Some(StructType(ttlKeySchema.drop(1)))
+        Some(StructType(ttlKeySchema.drop(2)))
       )
     )
   }
@@ -191,9 +198,9 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) {
       ttlValSchema,
       Some(RangeKeyScanStateEncoderSpec(ttlKeySchema, Seq(0))),
       avroEnc = getAvroSerde(
-        StructType(ttlKeySchema.take(1)),
+        StructType(ttlKeySchema.take(2)),
         ttlValSchema,
-        Some(StructType(ttlKeySchema.drop(1)))
+        Some(StructType(ttlKeySchema.drop(2)))
       )
     )
   }
@@ -226,9 +233,9 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) {
       valSchema,
       Some(RangeKeyScanStateEncoderSpec(keySchema, Seq(0))),
       avroEnc = getAvroSerde(
-        StructType(avroKeySchema.take(1)),
+        StructType(avroKeySchema.take(2)),
         valSchema,
-        Some(StructType(avroKeySchema.drop(1)))
+        Some(StructType(avroKeySchema.drop(2)))
       ))
   }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
index 8f177745273a2..3a77db1cf7beb 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
@@ -23,7 +23,7 @@ import java.lang.Float.{floatToRawIntBits, intBitsToFloat}
 import java.nio.{ByteBuffer, ByteOrder}
 
 import org.apache.avro.Schema
-import org.apache.avro.generic.{GenericDatumReader, GenericDatumWriter}
+import org.apache.avro.generic.{GenericData, GenericDatumReader, GenericDatumWriter, GenericRecord}
 import org.apache.avro.io.{DecoderFactory, EncoderFactory}
 
 import org.apache.spark.internal.Logging
@@ -684,16 +684,98 @@ class RangeKeyScanStateEncoder(
     writer.getRow()
   }
 
+  def encodePrefixKeyForRangeScan(
+      row: UnsafeRow, avroType: Schema): Array[Byte] = {
+    val record = new GenericData.Record(avroType)
+    var fieldIdx = 0
+    rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
+      val field = fieldWithOrdinal._1
+      val value = row.get(idx, field.dataType)
+      if (value == null) {
+        record.put(fieldIdx, false) // isNull marker
+        record.put(fieldIdx + 1, new Array[Byte](field.dataType.defaultSize))
+      } else {
+        field.dataType match {
+          case LongType =>
+            val longVal = value.asInstanceOf[Long]
+            val marker = longVal >= 0
+            record.put(fieldIdx, marker)
+
+            // Convert long to byte array in big endian format
+            val bbuf = ByteBuffer.allocate(8)
+            bbuf.order(ByteOrder.BIG_ENDIAN)
+            bbuf.putLong(longVal)
+            // Create a new byte array to avoid Avro's issue with direct ByteBuffer arrays
+            val bytes = new Array[Byte](8)
+            bbuf.position(0)
+            bbuf.get(bytes)
+
+            // Wrap bytes in Avro's ByteBuffer to ensure proper handling
+            record.put(fieldIdx + 1, ByteBuffer.wrap(bytes))
+
+          case _ => throw new UnsupportedOperationException(
+            s"Range scan encoding not supported for data type: ${field.dataType}")
+        }
+      }
+      fieldIdx += 2
+    }
+
+    out.reset()
+    val writer = new GenericDatumWriter[GenericRecord](rangeScanAvroType)
+    val encoder = EncoderFactory.get().binaryEncoder(out, null)
+    writer.write(record, encoder)
+    encoder.flush()
+    out.toByteArray
+  }
+
+  def decodePrefixKeyForRangeScan(
+      bytes: Array[Byte],
+      avroType: Schema): UnsafeRow = {
+
+    val reader = new GenericDatumReader[GenericRecord](avroType)
+    val decoder = DecoderFactory.get().binaryDecoder(bytes, 0, bytes.length, null)
+    val record = reader.read(null, decoder)
+
+    val rowWriter = new UnsafeRowWriter(rangeScanKeyFieldsWithOrdinal.length)
+    rowWriter.resetRowWriter()
+
+    var fieldIdx = 0
+    rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
+      val field = fieldWithOrdinal._1
+      val isMarkerNull = record.get(fieldIdx) == null
+
+      if (isMarkerNull) {
+        rowWriter.setNullAt(idx)
+      } else {
+        field.dataType match {
+          case LongType =>
+            // Get bytes from Avro ByteBuffer
+            val byteBuf = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer]
+            val bbuf = ByteBuffer.allocate(8)
+            bbuf.order(ByteOrder.BIG_ENDIAN)
+
+            // Copy bytes to our ByteBuffer
+            bbuf.put(byteBuf.array(), byteBuf.position(), byteBuf.remaining())
+            bbuf.flip()
+
+            val longVal = bbuf.getLong(fieldIdx)
+            rowWriter.write(idx, longVal)
+
+          case _ => throw new UnsupportedOperationException(
+            s"Range scan decoding not supported for data type: ${field.dataType}")
+        }
+      }
+      fieldIdx += 2
+    }
+
+    rowWriter.getRow()
+  }
+
   override def encodeKey(row: UnsafeRow): Array[Byte] = {
     // This prefix key has the columns specified by orderingOrdinals
     val prefixKey = extractPrefixKey(row)
     val rangeScanKeyEncoded = if (avroEnc.isDefined) {
-      encodeUnsafeRow(
-        encodePrefixKeyForRangeScan(prefixKey),
-        avroEnc.get.keySerializer,
-        rangeScanAvroType,
-        out
-      )
+      encodePrefixKeyForRangeScan(prefixKey, rangeScanAvroType)
     } else {
       encodeUnsafeRow(encodePrefixKeyForRangeScan(prefixKey))
     }
@@ -745,14 +827,12 @@ class RangeKeyScanStateEncoder(
     Platform.copyMemory(keyBytes, decodeKeyStartOffset + 4,
       prefixKeyEncoded, Platform.BYTE_ARRAY_OFFSET, prefixKeyEncodedLen)
 
-    val prefixKeyDecodedForRangeScan = if (avroEnc.isDefined) {
-      decodeToUnsafeRow(prefixKeyEncoded, avroEnc.get.keyDeserializer,
-        rangeScanAvroType, rangeScanAvroProjection)
+    val prefixKeyDecoded = if (avroEnc.isDefined) {
+      decodePrefixKeyForRangeScan(prefixKeyEncoded, rangeScanAvroType)
     } else {
-      decodeToUnsafeRow(prefixKeyEncoded,
-        numFields = orderingOrdinals.length)
+      decodePrefixKeyForRangeScan(decodeToUnsafeRow(prefixKeyEncoded,
+        numFields = orderingOrdinals.length))
     }
-    val prefixKeyDecoded = decodePrefixKeyForRangeScan(prefixKeyDecodedForRangeScan)
 
     if (orderingOrdinals.length < keySchema.length) {
       // Here we calculate the remainingKeyEncodedLen leveraging the length of keyBytes
@@ -785,12 +865,7 @@ class RangeKeyScanStateEncoder(
 
   override def encodePrefixKey(prefixKey: UnsafeRow): Array[Byte] = {
     val rangeScanKeyEncoded = if (avroEnc.isDefined) {
-      encodeUnsafeRow(
-        encodePrefixKeyForRangeScan(prefixKey),
-        avroEnc.get.keySerializer,
-        rangeScanAvroType,
-        out
-      )
+      encodePrefixKeyForRangeScan(prefixKey, rangeScanAvroType)
     } else {
       encodeUnsafeRow(encodePrefixKeyForRangeScan(prefixKey))
     }

From c5ef895875cd8d677ec70f7cf7612116d06e0c8b Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Fri, 8 Nov 2024 15:41:47 -0800
Subject: [PATCH 22/30] everything but batch works

---
 .../streaming/IncrementalExecution.scala      | 19 ++++++-
 .../execution/streaming/ListStateImpl.scala   |  4 +-
 .../streaming/ListStateImplWithTTL.scala      | 11 ++--
 .../execution/streaming/MapStateImpl.scala    |  4 +-
 .../streaming/MapStateImplWithTTL.scala       | 26 ++++-----
 .../StateStoreColumnFamilySchemaUtils.scala   | 55 ++++++++++++-------
 .../StatefulProcessorHandleImpl.scala         | 16 +++---
 .../sql/execution/streaming/TTLState.scala    |  6 +-
 .../execution/streaming/TimerStateImpl.scala  |  4 +-
 .../streaming/TransformWithStateExec.scala    | 43 ++++++++++-----
 .../execution/streaming/ValueStateImpl.scala  |  4 +-
 .../streaming/ValueStateImplWithTTL.scala     | 10 ++--
 .../state/HDFSBackedStateStoreProvider.scala  |  2 +-
 .../streaming/state/RocksDBStateEncoder.scala | 16 +++---
 .../state/RocksDBStateStoreProvider.scala     |  2 +-
 .../StateSchemaCompatibilityChecker.scala     |  4 +-
 .../streaming/state/StateStore.scala          |  2 +-
 .../streaming/state/MemoryStateStore.scala    |  2 +-
 ...sDBStateStoreCheckpointFormatV2Suite.scala |  2 +-
 19 files changed, 137 insertions(+), 95 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala
index 4b8bc72b2ed7f..e5604ffc2eb6b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala
@@ -259,6 +259,19 @@ class IncrementalExecution(
     }
   }
 
+  object StateStoreColumnFamilySchemas extends SparkPlanPartialRule {
+    override val rule: PartialFunction[SparkPlan, SparkPlan] = {
+      case statefulOp: StatefulOperator =>
+        statefulOp match {
+          case transformWithStateExec: TransformWithStateExec =>
+            transformWithStateExec.copy(
+              columnFamilySchemas = transformWithStateExec.getColFamilySchemas()
+            )
+          case _ => statefulOp
+        }
+    }
+  }
+
   object StateOpIdRule extends SparkPlanPartialRule {
     override val rule: PartialFunction[SparkPlan, SparkPlan] = {
       case StateStoreSaveExec(keys, None, None, None, None, stateFormatVersion,
@@ -552,9 +565,9 @@ class IncrementalExecution(
       // The rule below doesn't change the plan but can cause the side effect that
       // metadata/schema is written in the checkpoint directory of stateful operator.
       planWithStateOpId transform StateSchemaAndOperatorMetadataRule.rule
-
-      simulateWatermarkPropagation(planWithStateOpId)
-      planWithStateOpId transform WatermarkPropagationRule.rule
+      val planWithStateSchemas = planWithStateOpId transform StateStoreColumnFamilySchemas.rule
+      simulateWatermarkPropagation(planWithStateSchemas)
+      planWithStateSchemas transform WatermarkPropagationRule.rule
     }
   }
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
index 08eaa3677d46f..e0a201834e24d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
@@ -19,7 +19,7 @@ package org.apache.spark.sql.execution.streaming
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.execution.metric.SQLMetric
-import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, NoPrefixKeyStateEncoderSpec, StateStore, StateStoreErrors}
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoder, NoPrefixKeyStateEncoderSpec, StateStore, StateStoreErrors}
 import org.apache.spark.sql.streaming.ListState
 import org.apache.spark.sql.types.StructType
 
@@ -42,7 +42,7 @@ class ListStateImpl[S](
      keyExprEnc: ExpressionEncoder[Any],
      valEncoder: ExpressionEncoder[Any],
      metrics: Map[String, SQLMetric] = Map.empty,
-     avroEnc: Option[AvroEncoderSpec] = None)
+     avroEnc: Option[AvroEncoder] = None)
   extends ListStateMetricsImpl
   with ListState[S]
   with Logging {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala
index 445ca743b3855..313e974fd5a7c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala
@@ -20,7 +20,7 @@ import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, NoPrefixKeyStateEncoderSpec, StateStore, StateStoreErrors}
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoder, NoPrefixKeyStateEncoderSpec, StateStore, StateStoreErrors}
 import org.apache.spark.sql.streaming.{ListState, TTLConfig}
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.util.NextIterator
@@ -38,7 +38,7 @@ import org.apache.spark.util.NextIterator
  * @param metrics - metrics to be updated as part of stateful processing
  * @param avroEnc - optional Avro serializer and deserializer for this state variable that
  *                is used by the StateStore to encode state in Avro format
- * @param ttlAvroEnc - optional Avro serializer and deserializer for TTL state that
+ * @param secondaryIndexAvroEnc - optional Avro serializer and deserializer for TTL state that
  *                is used by the StateStore to encode state in Avro format
  * @tparam S - data type of object that will be stored
  */
@@ -50,9 +50,10 @@ class ListStateImplWithTTL[S](
     ttlConfig: TTLConfig,
     batchTimestampMs: Long,
     metrics: Map[String, SQLMetric] = Map.empty,
-    avroEnc: Option[AvroEncoderSpec] = None,
-    ttlAvroEnc: Option[AvroEncoderSpec] = None)
-  extends SingleKeyTTLStateImpl(stateName, store, keyExprEnc, batchTimestampMs, ttlAvroEnc)
+    avroEnc: Option[AvroEncoder] = None,
+    secondaryIndexAvroEnc: Option[AvroEncoder] = None)
+  extends SingleKeyTTLStateImpl(
+    stateName, store, keyExprEnc, batchTimestampMs, secondaryIndexAvroEnc)
   with ListStateMetricsImpl
   with ListState[S] {
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala
index eb96b32722aaf..b57eaec8d1e3c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala
@@ -20,7 +20,7 @@ import org.apache.spark.internal.Logging
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, PrefixKeyScanStateEncoderSpec, StateStore, StateStoreErrors, UnsafeRowPair}
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoder, PrefixKeyScanStateEncoderSpec, StateStore, StateStoreErrors, UnsafeRowPair}
 import org.apache.spark.sql.streaming.MapState
 import org.apache.spark.sql.types.StructType
 
@@ -44,7 +44,7 @@ class MapStateImpl[K, V](
     userKeyEnc: ExpressionEncoder[Any],
     valEncoder: ExpressionEncoder[Any],
     metrics: Map[String, SQLMetric] = Map.empty,
-    avroEnc: Option[AvroEncoderSpec] = None) extends MapState[K, V] with Logging {
+    avroEnc: Option[AvroEncoder] = None) extends MapState[K, V] with Logging {
 
   // Pack grouping key and user key together as a prefixed composite key
   private val schemaForCompositeKeyRow: StructType = {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala
index 11554d8532396..f267304b1fe4a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImplWithTTL.scala
@@ -21,7 +21,7 @@ import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, PrefixKeyScanStateEncoderSpec, StateStore, StateStoreErrors}
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoder, PrefixKeyScanStateEncoderSpec, StateStore, StateStoreErrors}
 import org.apache.spark.sql.streaming.{MapState, TTLConfig}
 import org.apache.spark.util.NextIterator
 
@@ -38,25 +38,25 @@ import org.apache.spark.util.NextIterator
  * @param metrics - metrics to be updated as part of stateful processing
  * @param avroEnc - optional Avro serializer and deserializer for this state variable that
  *                is used by the StateStore to encode state in Avro format
- * @param ttlAvroEnc - optional Avro serializer and deserializer for TTL state that
+ * @param secondaryIndexAvroEnc - optional Avro serializer and deserializer for TTL state that
  *                is used by the StateStore to encode state in Avro format
  * @tparam K - type of key for map state variable
  * @tparam V - type of value for map state variable
  * @return - instance of MapState of type [K,V] that can be used to store state persistently
  */
 class MapStateImplWithTTL[K, V](
-    store: StateStore,
-    stateName: String,
-    keyExprEnc: ExpressionEncoder[Any],
-    userKeyEnc: ExpressionEncoder[Any],
-    valEncoder: ExpressionEncoder[Any],
-    ttlConfig: TTLConfig,
-    batchTimestampMs: Long,
-    metrics: Map[String, SQLMetric] = Map.empty,
-    avroEnc: Option[AvroEncoderSpec] = None,
-    ttlAvroEnc: Option[AvroEncoderSpec] = None)
+      store: StateStore,
+      stateName: String,
+      keyExprEnc: ExpressionEncoder[Any],
+      userKeyEnc: ExpressionEncoder[Any],
+      valEncoder: ExpressionEncoder[Any],
+      ttlConfig: TTLConfig,
+      batchTimestampMs: Long,
+      metrics: Map[String, SQLMetric] = Map.empty,
+      avroEnc: Option[AvroEncoder] = None,
+      secondaryIndexAvroEnc: Option[AvroEncoder] = None)
   extends CompositeKeyTTLStateImpl[K](stateName, store,
-    keyExprEnc, userKeyEnc, batchTimestampMs, ttlAvroEnc)
+    keyExprEnc, userKeyEnc, batchTimestampMs, secondaryIndexAvroEnc)
   with MapState[K, V] with Logging {
 
   private val stateTypesEncoder = new CompositeKeyStateEncoder(
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
index a2c98127b3b91..2791277fbfe9f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
@@ -21,7 +21,7 @@ import org.apache.spark.sql.Encoder
 import org.apache.spark.sql.avro.{AvroDeserializer, AvroOptions, AvroSerializer, SchemaConverters}
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, NoPrefixKeyStateEncoderSpec, PrefixKeyScanStateEncoderSpec, RangeKeyScanStateEncoderSpec, StateStoreColFamilySchema}
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoder, NoPrefixKeyStateEncoderSpec, PrefixKeyScanStateEncoderSpec, RangeKeyScanStateEncoderSpec, StateStoreColFamilySchema}
 import org.apache.spark.sql.types.{BinaryType, BooleanType, ByteType, DataType, DoubleType, FloatType, IntegerType, LongType, NullType, ShortType, StructField, StructType}
 
 object StateStoreColumnFamilySchemaUtils {
@@ -29,7 +29,6 @@ object StateStoreColumnFamilySchemaUtils {
   def apply(initializeAvroSerde: Boolean): StateStoreColumnFamilySchemaUtils =
     new StateStoreColumnFamilySchemaUtils(initializeAvroSerde)
 
-
   /**
    * Avro uses zig-zag encoding for some fixed-length types, like Longs and Ints. For range scans
    * we want to use big-endian encoding, so we need to convert the source schema to replace these
@@ -76,6 +75,16 @@ object StateStoreColumnFamilySchemaUtils {
  */
 class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) extends Logging {
 
+  private def getAvroSerdeForSchema(schema: StructType): (AvroSerializer, AvroDeserializer) = {
+    val avroType = SchemaConverters.toAvroType(schema)
+    val avroOptions = AvroOptions(Map.empty)
+    val serializer = new AvroSerializer(schema, avroType, nullable = false)
+    val deserializer = new AvroDeserializer(avroType, schema,
+      avroOptions.datetimeRebaseModeInRead, avroOptions.useStableIdForUnionType,
+      avroOptions.stableIdPrefixForUnionType, avroOptions.recursiveFieldMaxDepth)
+    (serializer, deserializer)
+  }
+
   /**
    * If initializeAvroSerde is true, this method will create an Avro Serializer and Deserializer
    * for a particular key and value schema.
@@ -84,30 +93,19 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) extends Lo
       keySchema: StructType,
       valSchema: StructType,
       suffixKeySchema: Option[StructType] = None
-  ): Option[AvroEncoderSpec] = {
+  ): Option[AvroEncoder] = {
     if (initializeAvroSerde) {
-      val avroType = SchemaConverters.toAvroType(valSchema)
-      val avroOptions = AvroOptions(Map.empty)
-      val keyAvroType = SchemaConverters.toAvroType(keySchema)
-      val keySer = new AvroSerializer(keySchema, keyAvroType, nullable = false)
-      val keyDe = new AvroDeserializer(keyAvroType, keySchema,
-        avroOptions.datetimeRebaseModeInRead, avroOptions.useStableIdForUnionType,
-        avroOptions.stableIdPrefixForUnionType, avroOptions.recursiveFieldMaxDepth)
-      val valueSerializer = new AvroSerializer(valSchema, avroType, nullable = false)
-      val valueDeserializer = new AvroDeserializer(avroType, valSchema,
-        avroOptions.datetimeRebaseModeInRead, avroOptions.useStableIdForUnionType,
-        avroOptions.stableIdPrefixForUnionType, avroOptions.recursiveFieldMaxDepth)
+      val (keySer, keyDe) =
+        getAvroSerdeForSchema(keySchema)
+      val (valueSerializer, valueDeserializer) =
+        getAvroSerdeForSchema(valSchema)
       val (suffixKeySer, suffixKeyDe) = if (suffixKeySchema.isDefined) {
-        val userKeyAvroType = SchemaConverters.toAvroType(suffixKeySchema.get)
-        val skSer = new AvroSerializer(suffixKeySchema.get, userKeyAvroType, nullable = false)
-        val skDe = new AvroDeserializer(userKeyAvroType, suffixKeySchema.get,
-          avroOptions.datetimeRebaseModeInRead, avroOptions.useStableIdForUnionType,
-          avroOptions.stableIdPrefixForUnionType, avroOptions.recursiveFieldMaxDepth)
-        (Some(skSer), Some(skDe))
+        val serde = getAvroSerdeForSchema(suffixKeySchema.get)
+        (Some(serde._1), Some(serde._2))
       } else {
         (None, None)
       }
-      Some(AvroEncoderSpec(
+      Some(AvroEncoder(
         keySer, keyDe, valueSerializer, valueDeserializer, suffixKeySer, suffixKeyDe))
     } else {
       None
@@ -164,6 +162,11 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) extends Lo
     )
   }
 
+  // This function creates the StateStoreColFamilySchema for
+  // the TTL secondary index.
+  // Because we want to encode fixed-length types as binary types
+  // if we are using Avro, we need to do some schema conversion to ensure
+  // we can use range scan
   def getTtlStateSchema(
       stateName: String,
       keyEncoder: ExpressionEncoder[Any]): StateStoreColFamilySchema = {
@@ -184,6 +187,11 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) extends Lo
     )
   }
 
+  // This function creates the StateStoreColFamilySchema for
+  // the TTL secondary index.
+  // Because we want to encode fixed-length types as binary types
+  // if we are using Avro, we need to do some schema conversion to ensure
+  // we can use range scan
   def getTtlStateSchema(
       stateName: String,
       keyEncoder: ExpressionEncoder[Any],
@@ -221,6 +229,11 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) extends Lo
       ))
   }
 
+  // This function creates the StateStoreColFamilySchema for
+  // Timers' secondary index.
+  // Because we want to encode fixed-length types as binary types
+  // if we are using Avro, we need to do some schema conversion to ensure
+  // we can use range scan
   def getTimerStateSchemaForSecIndex(
       stateName: String,
       keySchema: StructType,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
index 32e5d6d4946b0..1fd244ea167dd 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
@@ -394,10 +394,10 @@ class DriverStatefulProcessorHandleImpl(
     val stateName = TimerStateUtils.getTimerStateVarName(timeMode.toString)
     val secIndexColFamilyName = TimerStateUtils.getSecIndexColFamilyName(timeMode.toString)
     val timerEncoder = new TimerKeyEncoder(keyExprEnc)
-    val colFamilySchema = schemaUtils.
-      getTimerStateSchema(stateName, timerEncoder.schemaForKeyRow, timerEncoder.schemaForValueRow)
-    val secIndexColFamilySchema = schemaUtils.
-      getTimerStateSchemaForSecIndex(secIndexColFamilyName,
+    val colFamilySchema = schemaUtils
+      .getTimerStateSchema(stateName, timerEncoder.schemaForKeyRow, timerEncoder.schemaForValueRow)
+    val secIndexColFamilySchema = schemaUtils
+      .getTimerStateSchemaForSecIndex(secIndexColFamilyName,
         timerEncoder.keySchemaForSecIndex,
         timerEncoder.schemaForValueRow)
     columnFamilySchemas.put(stateName, colFamilySchema)
@@ -458,8 +458,8 @@ class DriverStatefulProcessorHandleImpl(
     }
 
     val stateEncoder = encoderFor[T]
-    val colFamilySchema = schemaUtils.
-      getListStateSchema(stateName, keyExprEnc, stateEncoder, ttlEnabled)
+    val colFamilySchema = schemaUtils
+      .getListStateSchema(stateName, keyExprEnc, stateEncoder, ttlEnabled)
     checkIfDuplicateVariableDefined(stateName)
     columnFamilySchemas.put(stateName, colFamilySchema)
     val stateVariableInfo = TransformWithStateVariableUtils.
@@ -494,8 +494,8 @@ class DriverStatefulProcessorHandleImpl(
     }
 
 
-    val colFamilySchema = schemaUtils.
-      getMapStateSchema(stateName, keyExprEnc, userKeyEnc, valEncoder, ttlEnabled)
+    val colFamilySchema = schemaUtils
+      .getMapStateSchema(stateName, keyExprEnc, userKeyEnc, valEncoder, ttlEnabled)
     columnFamilySchemas.put(stateName, colFamilySchema)
     val stateVariableInfo = TransformWithStateVariableUtils.
       getMapState(stateName, ttlEnabled = ttlEnabled)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TTLState.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TTLState.scala
index e4e45fbd74bbc..02008a1ba4fd0 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TTLState.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TTLState.scala
@@ -23,7 +23,7 @@ import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions.{UnsafeProjection, UnsafeRow}
 import org.apache.spark.sql.execution.streaming.StateStoreColumnFamilySchemaUtils.getTtlColFamilyName
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, RangeKeyScanStateEncoderSpec, StateStore}
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoder, RangeKeyScanStateEncoderSpec, StateStore}
 import org.apache.spark.sql.types._
 
 object StateTTLSchema {
@@ -81,7 +81,7 @@ abstract class SingleKeyTTLStateImpl(
     store: StateStore,
     keyExprEnc: ExpressionEncoder[Any],
     ttlExpirationMs: Long,
-    avroEnc: Option[AvroEncoderSpec] = None)
+    avroEnc: Option[AvroEncoder] = None)
   extends TTLState {
 
   import org.apache.spark.sql.execution.streaming.StateTTLSchema._
@@ -202,7 +202,7 @@ abstract class CompositeKeyTTLStateImpl[K](
     keyExprEnc: ExpressionEncoder[Any],
     userKeyEncoder: ExpressionEncoder[Any],
     ttlExpirationMs: Long,
-    avroEnc: Option[AvroEncoderSpec] = None)
+    avroEnc: Option[AvroEncoder] = None)
   extends TTLState {
 
   import org.apache.spark.sql.execution.streaming.StateTTLSchema._
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TimerStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TimerStateImpl.scala
index 5459e65526f5f..74eaf062ec547 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TimerStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TimerStateImpl.scala
@@ -65,8 +65,8 @@ class TimerStateImpl(
     store: StateStore,
     timeMode: TimeMode,
     keyExprEnc: ExpressionEncoder[Any],
-    avroEnc: Option[AvroEncoderSpec] = None,
-    secIndexAvroEnc: Option[AvroEncoderSpec] = None) extends Logging {
+    avroEnc: Option[AvroEncoder] = None,
+    secIndexAvroEnc: Option[AvroEncoder] = None) extends Logging {
 
   private val EMPTY_ROW =
     UnsafeProjection.create(Array[DataType](NullType)).apply(InternalRow.apply(null))
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
index 87b452a6a38cd..c1900e6c5e25f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
@@ -75,7 +75,8 @@ case class TransformWithStateExec(
     initialStateGroupingAttrs: Seq[Attribute],
     initialStateDataAttrs: Seq[Attribute],
     initialStateDeserializer: Expression,
-    initialState: SparkPlan)
+    initialState: SparkPlan,
+    columnFamilySchemas: Map[String, StateStoreColFamilySchema] = Map.empty)
   extends BinaryExecNode with StateStoreWriter with WatermarkSupport with ObjectProducerExec {
 
   override def shortName: String = "transformWithStateExec"
@@ -118,7 +119,7 @@ case class TransformWithStateExec(
    * Fetching the columnFamilySchemas from the StatefulProcessorHandle
    * after init is called.
    */
-  private def getColFamilySchemas(): Map[String, StateStoreColFamilySchema] = {
+  def getColFamilySchemas(): Map[String, StateStoreColFamilySchema] = {
     val columnFamilySchemas = getDriverProcessorHandle().getColumnFamilySchemas
     closeProcessorHandle()
     columnFamilySchemas
@@ -524,7 +525,6 @@ case class TransformWithStateExec(
 
   override protected def doExecute(): RDD[InternalRow] = {
     metrics // force lazy init at driver
-
     validateTimeMode()
 
     if (hasInitialState) {
@@ -535,11 +535,10 @@ case class TransformWithStateExec(
         initialState.execute(),
         getStateInfo,
         storeNames = Seq(),
-        session.streams.stateStoreCoordinator,
-        getColFamilySchemas()) {
+        session.streams.stateStoreCoordinator) {
         // The state store aware zip partitions will provide us with two iterators,
         // child data iterator and the initial state iterator per partition.
-        case (partitionId, childDataIterator, initStateIterator, colFamilySchemas) =>
+        case (partitionId, childDataIterator, initStateIterator) =>
           if (isStreaming) {
             val stateStoreId = StateStoreId(stateInfo.get.checkpointLocation,
               stateInfo.get.operatorId, partitionId)
@@ -557,27 +556,26 @@ case class TransformWithStateExec(
             )
 
             processDataWithInitialState(
-              store, childDataIterator, initStateIterator, colFamilySchemas)
+              store, childDataIterator, initStateIterator, columnFamilySchemas)
           } else {
             initNewStateStoreAndProcessData(
-              partitionId, hadoopConfBroadcast, getColFamilySchemas()) { (store, schemas) =>
+              partitionId, hadoopConfBroadcast, columnFamilySchemas) { (store, schemas) =>
               processDataWithInitialState(store, childDataIterator, initStateIterator, schemas)
             }
           }
       }
     } else {
       if (isStreaming) {
-        child.execute().mapPartitionsWithStateStoreWithSchemas[InternalRow](
+        child.execute().mapPartitionsWithStateStore[InternalRow](
           getStateInfo,
           keyEncoder.schema,
           DUMMY_VALUE_ROW_SCHEMA,
           NoPrefixKeyStateEncoderSpec(keyEncoder.schema),
           session.sessionState,
           Some(session.streams.stateStoreCoordinator),
-          useColumnFamilies = true,
-          columnFamilySchemas = getColFamilySchemas()
+          useColumnFamilies = true
         ) {
-          case (store: StateStore, singleIterator: Iterator[InternalRow], columnFamilySchemas) =>
+          case (store: StateStore, singleIterator: Iterator[InternalRow]) =>
             processData(store, singleIterator, columnFamilySchemas)
         }
       } else {
@@ -588,7 +586,7 @@ case class TransformWithStateExec(
         child.execute().mapPartitionsWithIndex[InternalRow](
           (i: Int, iter: Iterator[InternalRow]) => {
             initNewStateStoreAndProcessData(
-              i, hadoopConfBroadcast, getColFamilySchemas()) { (store, schemas) =>
+              i, hadoopConfBroadcast, columnFamilySchemas) { (store, schemas) =>
               processData(store, iter, schemas)
             }
           }
@@ -733,6 +731,22 @@ object TransformWithStateExec {
       stateStoreCkptIds = None
     )
 
+    val stateStoreEncoding = child.session.sessionState.conf.getConf(
+      SQLConf.STREAMING_STATE_STORE_ENCODING_FORMAT
+    )
+
+    def getDriverProcessorHandle(): DriverStatefulProcessorHandleImpl = {
+      val driverProcessorHandle = new DriverStatefulProcessorHandleImpl(
+        timeMode, keyEncoder, initializeAvroEnc =
+          stateStoreEncoding == StateStoreEncoding.Avro.toString)
+      driverProcessorHandle.setHandleState(StatefulProcessorHandleState.PRE_INIT)
+      statefulProcessor.setHandle(driverProcessorHandle)
+      statefulProcessor.init(outputMode, timeMode)
+      driverProcessorHandle
+    }
+
+    val columnFamilySchemas = getDriverProcessorHandle().getColumnFamilySchemas
+
     new TransformWithStateExec(
       keyDeserializer,
       valueDeserializer,
@@ -753,7 +767,8 @@ object TransformWithStateExec {
       initialStateGroupingAttrs,
       initialStateDataAttrs,
       initialStateDeserializer,
-      initialState)
+      initialState,
+      columnFamilySchemas)
   }
 }
 // scalastyle:on argcount
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala
index e12c02f4ea2a3..9eb51abaee6e3 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala
@@ -19,7 +19,7 @@ package org.apache.spark.sql.execution.streaming
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.execution.metric.SQLMetric
-import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, NoPrefixKeyStateEncoderSpec, StateStore}
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoder, NoPrefixKeyStateEncoderSpec, StateStore}
 import org.apache.spark.sql.streaming.ValueState
 
 /**
@@ -40,7 +40,7 @@ class ValueStateImpl[S](
     keyExprEnc: ExpressionEncoder[Any],
     valEncoder: ExpressionEncoder[Any],
     metrics: Map[String, SQLMetric] = Map.empty,
-    avroEnc: Option[AvroEncoderSpec] = None)
+    avroEnc: Option[AvroEncoder] = None)
   extends ValueState[S] with Logging {
 
   private val stateTypesEncoder = StateTypesEncoder(keyExprEnc, valEncoder, stateName)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala
index 3ab9e5d226f23..7c2401dffb2f3 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala
@@ -20,7 +20,7 @@ import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.state.{AvroEncoderSpec, NoPrefixKeyStateEncoderSpec, StateStore}
+import org.apache.spark.sql.execution.streaming.state.{AvroEncoder, NoPrefixKeyStateEncoderSpec, StateStore}
 import org.apache.spark.sql.streaming.{TTLConfig, ValueState}
 
 /**
@@ -36,7 +36,7 @@ import org.apache.spark.sql.streaming.{TTLConfig, ValueState}
  * @param metrics - metrics to be updated as part of stateful processing
  * @param avroEnc - optional Avro serializer and deserializer for this state variable that
  *                is used by the StateStore to encode state in Avro format
- * @param ttlAvroEnc - optional Avro serializer and deserializer for TTL state that
+ * @param secondaryIndexAvroEnc - optional Avro serializer and deserializer for TTL state that
  *                is used by the StateStore to encode state in Avro format
  * @tparam S - data type of object that will be stored
  */
@@ -48,10 +48,10 @@ class ValueStateImplWithTTL[S](
     ttlConfig: TTLConfig,
     batchTimestampMs: Long,
     metrics: Map[String, SQLMetric] = Map.empty,
-    avroEnc: Option[AvroEncoderSpec] = None,
-    ttlAvroEnc: Option[AvroEncoderSpec] = None)
+    avroEnc: Option[AvroEncoder] = None,
+    secondaryIndexAvroEnc: Option[AvroEncoder] = None)
   extends SingleKeyTTLStateImpl(
-    stateName, store, keyExprEnc, batchTimestampMs, ttlAvroEnc) with ValueState[S] {
+    stateName, store, keyExprEnc, batchTimestampMs, secondaryIndexAvroEnc) with ValueState[S] {
 
   private val stateTypesEncoder = StateTypesEncoder(keyExprEnc, valEncoder,
     stateName, hasTtl = true)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala
index 11bf8ce53b560..423ce50776fa4 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala
@@ -128,7 +128,7 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
         keyStateEncoderSpec: KeyStateEncoderSpec,
         useMultipleValuesPerKey: Boolean = false,
         isInternal: Boolean = false,
-        avroEnc: Option[AvroEncoderSpec]): Unit = {
+        avroEnc: Option[AvroEncoder]): Unit = {
       throw StateStoreErrors.multipleColumnFamiliesNotSupported(providerName)
     }
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
index 3a77db1cf7beb..3ffe792c3e618 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
@@ -98,12 +98,12 @@ abstract class RocksDBKeyStateEncoderBase(
   }
 }
 
-object RocksDBStateEncoder {
+object RocksDBStateEncoder extends Logging {
   def getKeyEncoder(
       keyStateEncoderSpec: KeyStateEncoderSpec,
       useColumnFamilies: Boolean,
       virtualColFamilyId: Option[Short] = None,
-      avroEnc: Option[AvroEncoderSpec] = None): RocksDBKeyStateEncoder = {
+      avroEnc: Option[AvroEncoder] = None): RocksDBKeyStateEncoder = {
     // Return the key state encoder based on the requested type
     keyStateEncoderSpec match {
       case NoPrefixKeyStateEncoderSpec(keySchema) =>
@@ -126,7 +126,7 @@ object RocksDBStateEncoder {
   def getValueEncoder(
       valueSchema: StructType,
       useMultipleValuesPerKey: Boolean,
-      avroEnc: Option[AvroEncoderSpec] = None): RocksDBValueStateEncoder = {
+      avroEnc: Option[AvroEncoder] = None): RocksDBValueStateEncoder = {
     if (useMultipleValuesPerKey) {
       new MultiValuedStateEncoder(valueSchema, avroEnc)
     } else {
@@ -233,7 +233,7 @@ class PrefixKeyScanStateEncoder(
     numColsPrefixKey: Int,
     useColumnFamilies: Boolean = false,
     virtualColFamilyId: Option[Short] = None,
-    avroEnc: Option[AvroEncoderSpec] = None)
+    avroEnc: Option[AvroEncoder] = None)
   extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) {
 
   import RocksDBStateEncoder._
@@ -409,7 +409,7 @@ class RangeKeyScanStateEncoder(
     orderingOrdinals: Seq[Int],
     useColumnFamilies: Boolean = false,
     virtualColFamilyId: Option[Short] = None,
-    avroEnc: Option[AvroEncoderSpec] = None)
+    avroEnc: Option[AvroEncoder] = None)
   extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) with Logging {
 
   import RocksDBStateEncoder._
@@ -897,7 +897,7 @@ class NoPrefixKeyStateEncoder(
     keySchema: StructType,
     useColumnFamilies: Boolean = false,
     virtualColFamilyId: Option[Short] = None,
-    avroEnc: Option[AvroEncoderSpec] = None)
+    avroEnc: Option[AvroEncoder] = None)
   extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) with Logging {
 
   import RocksDBStateEncoder._
@@ -988,7 +988,7 @@ class NoPrefixKeyStateEncoder(
  */
 class MultiValuedStateEncoder(
     valueSchema: StructType,
-    avroEnc: Option[AvroEncoderSpec] = None)
+    avroEnc: Option[AvroEncoder] = None)
   extends RocksDBValueStateEncoder with Logging {
 
   import RocksDBStateEncoder._
@@ -1084,7 +1084,7 @@ class MultiValuedStateEncoder(
  */
 class SingleValueStateEncoder(
     valueSchema: StructType,
-    avroEnc: Option[AvroEncoderSpec] = None)
+    avroEnc: Option[AvroEncoder] = None)
   extends RocksDBValueStateEncoder with Logging {
 
   import RocksDBStateEncoder._
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala
index 0ab10a6fbdb98..146c983be3170 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala
@@ -72,7 +72,7 @@ private[sql] class RocksDBStateStoreProvider
         keyStateEncoderSpec: KeyStateEncoderSpec,
         useMultipleValuesPerKey: Boolean = false,
         isInternal: Boolean = false,
-        avroEnc: Option[AvroEncoderSpec]): Unit = {
+        avroEnc: Option[AvroEncoder]): Unit = {
       verifyColFamilyCreationOrDeletion("create_col_family", colFamilyName, isInternal)
       val newColFamilyId = rocksDB.createColFamilyIfAbsent(colFamilyName)
       keyValueEncoderMap.putIfAbsent(colFamilyName,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
index a05b452bd0184..b5f2f318de418 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
@@ -40,7 +40,7 @@ case class StateSchemaValidationResult(
 
 // Avro encoder that is used by the RocksDBStateStoreProvider and RocksDBStateEncoder
 // in order to serialize from UnsafeRow to a byte array of Avro encoding.
-case class AvroEncoderSpec(
+case class AvroEncoder(
   keySerializer: AvroSerializer,
   keyDeserializer: AvroDeserializer,
   valueSerializer: AvroSerializer,
@@ -56,7 +56,7 @@ case class StateStoreColFamilySchema(
     valueSchema: StructType,
     keyStateEncoderSpec: Option[KeyStateEncoderSpec] = None,
     userKeyEncoderSchema: Option[StructType] = None,
-    avroEnc: Option[AvroEncoderSpec] = None
+    avroEnc: Option[AvroEncoder] = None
 ) extends Serializable
 
 class StateSchemaCompatibilityChecker(
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
index 255577dfccaa6..50843b1aeb438 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
@@ -144,7 +144,7 @@ trait StateStore extends ReadStateStore {
       keyStateEncoderSpec: KeyStateEncoderSpec,
       useMultipleValuesPerKey: Boolean = false,
       isInternal: Boolean = false,
-      avroEncoderSpec: Option[AvroEncoderSpec] = None): Unit
+      avroEncoderSpec: Option[AvroEncoder] = None): Unit
 
   /**
    * Put a new non-null value for a non-null key. Implementations must be aware that the UnsafeRows
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MemoryStateStore.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MemoryStateStore.scala
index bfcd828d01cc3..9a982a2264701 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MemoryStateStore.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MemoryStateStore.scala
@@ -37,7 +37,7 @@ class MemoryStateStore extends StateStore() {
       keyStateEncoderSpec: KeyStateEncoderSpec,
       useMultipleValuesPerKey: Boolean = false,
       isInternal: Boolean = false,
-      avroEnc: Option[AvroEncoderSpec]): Unit = {
+      avroEnc: Option[AvroEncoder]): Unit = {
     throw StateStoreErrors.multipleColumnFamiliesNotSupported("MemoryStateStoreProvider")
   }
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala
index e6454d3c77a2f..346bfd37798f6 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala
@@ -92,7 +92,7 @@ case class CkptIdCollectingStateStoreWrapper(innerStore: StateStore) extends Sta
       keyStateEncoderSpec: KeyStateEncoderSpec,
       useMultipleValuesPerKey: Boolean = false,
       isInternal: Boolean = false,
-      avroEnc: Option[AvroEncoderSpec]): Unit = {
+      avroEnc: Option[AvroEncoder]): Unit = {
     innerStore.createColFamilyIfAbsent(
       colFamilyName,
       keySchema,

From e22e1a298fbb6d0eaecd4ac4fdd97a4d71cfac08 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Fri, 8 Nov 2024 15:48:34 -0800
Subject: [PATCH 23/30] splitting it up

---
 .../StateStoreColumnFamilySchemaUtils.scala   | 25 +++++++++++--------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
index 2791277fbfe9f..bb296e2f4aae0 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
@@ -74,15 +74,17 @@ object StateStoreColumnFamilySchemaUtils {
  *                            StateStoreColumnFamilySchema for each state variable from the driver
  */
 class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) extends Logging {
+  private def getAvroSerializer(schema: StructType): AvroSerializer = {
+    val avroType = SchemaConverters.toAvroType(schema)
+    new AvroSerializer(schema, avroType, nullable = false)
+  }
 
-  private def getAvroSerdeForSchema(schema: StructType): (AvroSerializer, AvroDeserializer) = {
+  private def getAvroDeserializer(schema: StructType): AvroDeserializer = {
     val avroType = SchemaConverters.toAvroType(schema)
     val avroOptions = AvroOptions(Map.empty)
-    val serializer = new AvroSerializer(schema, avroType, nullable = false)
-    val deserializer = new AvroDeserializer(avroType, schema,
+    new AvroDeserializer(avroType, schema,
       avroOptions.datetimeRebaseModeInRead, avroOptions.useStableIdForUnionType,
       avroOptions.stableIdPrefixForUnionType, avroOptions.recursiveFieldMaxDepth)
-    (serializer, deserializer)
   }
 
   /**
@@ -95,18 +97,19 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) extends Lo
       suffixKeySchema: Option[StructType] = None
   ): Option[AvroEncoder] = {
     if (initializeAvroSerde) {
-      val (keySer, keyDe) =
-        getAvroSerdeForSchema(keySchema)
-      val (valueSerializer, valueDeserializer) =
-        getAvroSerdeForSchema(valSchema)
+
       val (suffixKeySer, suffixKeyDe) = if (suffixKeySchema.isDefined) {
-        val serde = getAvroSerdeForSchema(suffixKeySchema.get)
-        (Some(serde._1), Some(serde._2))
+        (Some(getAvroSerializer(suffixKeySchema.get)),
+          Some(getAvroDeserializer(suffixKeySchema.get)))
       } else {
         (None, None)
       }
       Some(AvroEncoder(
-        keySer, keyDe, valueSerializer, valueDeserializer, suffixKeySer, suffixKeyDe))
+        getAvroSerializer(keySchema),
+        getAvroDeserializer(keySchema),
+        getAvroSerializer(valSchema),
+        getAvroDeserializer(valSchema),
+        suffixKeySer, suffixKeyDe))
     } else {
       None
     }

From 730cae08f5d6e3308a0f56ee6a06cec385b814a8 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Fri, 8 Nov 2024 22:15:10 -0800
Subject: [PATCH 24/30] easy feedback to address

---
 .../streaming/IncrementalExecution.scala      | 10 ++--
 .../StreamingSymmetricHashJoinHelper.scala    | 50 -------------------
 .../streaming/state/RocksDBStateEncoder.scala |  6 ++-
 .../execution/streaming/state/package.scala   | 43 ----------------
 4 files changed, 9 insertions(+), 100 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala
index e5604ffc2eb6b..634222e785a44 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala
@@ -259,13 +259,13 @@ class IncrementalExecution(
     }
   }
 
-  object StateStoreColumnFamilySchemas extends SparkPlanPartialRule {
+  object StateStoreColumnFamilySchemasRule extends SparkPlanPartialRule {
     override val rule: PartialFunction[SparkPlan, SparkPlan] = {
       case statefulOp: StatefulOperator =>
         statefulOp match {
-          case transformWithStateExec: TransformWithStateExec =>
-            transformWithStateExec.copy(
-              columnFamilySchemas = transformWithStateExec.getColFamilySchemas()
+          case op: TransformWithStateExec =>
+            op.copy(
+              columnFamilySchemas = op.getColFamilySchemas()
             )
           case _ => statefulOp
         }
@@ -565,7 +565,7 @@ class IncrementalExecution(
       // The rule below doesn't change the plan but can cause the side effect that
       // metadata/schema is written in the checkpoint directory of stateful operator.
       planWithStateOpId transform StateSchemaAndOperatorMetadataRule.rule
-      val planWithStateSchemas = planWithStateOpId transform StateStoreColumnFamilySchemas.rule
+      val planWithStateSchemas = planWithStateOpId transform StateStoreColumnFamilySchemasRule.rule
       simulateWatermarkPropagation(planWithStateSchemas)
       planWithStateSchemas transform WatermarkPropagationRule.rule
     }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamingSymmetricHashJoinHelper.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamingSymmetricHashJoinHelper.scala
index 468d0df75fee4..fe3b1683fce0c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamingSymmetricHashJoinHelper.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamingSymmetricHashJoinHelper.scala
@@ -303,56 +303,6 @@ object StreamingSymmetricHashJoinHelper extends Logging {
     }
   }
 
-  /**
-   * A custom RDD that allows partitions to be "zipped" together, while ensuring the tasks'
-   * preferred location is based on which executors have the required join state stores already
-   * loaded. This class is a variant of [[org.apache.spark.rdd.ZippedPartitionsRDD2]] which only
-   * changes signature of `f` by taking in a map of column family schemas. This is used for
-   * passing the column family schemas when there is initial state for the TransformWithStateExec
-   * operator
-   */
-  class StateStoreAwareZipPartitionsRDDWithSchemas[A: ClassTag, B: ClassTag, V: ClassTag](
-      sc: SparkContext,
-      var f: (Int, Iterator[A], Iterator[B], Map[String, StateStoreColFamilySchema]) => Iterator[V],
-      var rdd1: RDD[A],
-      var rdd2: RDD[B],
-      stateInfo: StatefulOperatorStateInfo,
-      stateStoreNames: Seq[String],
-      @transient private val storeCoordinator: Option[StateStoreCoordinatorRef],
-      schemas: Map[String, StateStoreColFamilySchema])
-    extends ZippedPartitionsBaseRDD[V](sc, List(rdd1, rdd2)) {
-
-    /**
-     * Set the preferred location of each partition using the executor that has the related
-     * [[StateStoreProvider]] already loaded.
-     */
-    override def getPreferredLocations(partition: Partition): Seq[String] = {
-      stateStoreNames.flatMap { storeName =>
-        val stateStoreProviderId = StateStoreProviderId(stateInfo, partition.index, storeName)
-        storeCoordinator.flatMap(_.getLocation(stateStoreProviderId))
-      }.distinct
-    }
-
-    override def compute(s: Partition, context: TaskContext): Iterator[V] = {
-      val partitions = s.asInstanceOf[ZippedPartitionsPartition].partitions
-      if (partitions(0).index != partitions(1).index) {
-        throw new IllegalStateException(s"Partition ID should be same in both side: " +
-          s"left ${partitions(0).index} , right ${partitions(1).index}")
-      }
-
-      val partitionId = partitions(0).index
-      f(partitionId, rdd1.iterator(partitions(0), context),
-        rdd2.iterator(partitions(1), context), schemas)
-    }
-
-    override def clearDependencies(): Unit = {
-      super.clearDependencies()
-      rdd1 = null
-      rdd2 = null
-      f = null
-    }
-  }
-
   implicit class StateStoreAwareZipPartitionsHelper[T: ClassTag](dataRDD: RDD[T]) {
     /**
      * Function used by `StreamingSymmetricHashJoinExec` to zip together the partitions of two
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
index 3ffe792c3e618..afeec3988538b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
@@ -685,7 +685,9 @@ class RangeKeyScanStateEncoder(
   }
 
   def encodePrefixKeyForRangeScan(
-      row: UnsafeRow, avroType: Schema): Array[Byte] = {
+      row: UnsafeRow,
+      avroType: Schema
+  ): Array[Byte] = {
     val record = new GenericData.Record(avroType)
     var fieldIdx = 0
     rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
@@ -887,11 +889,11 @@ class RangeKeyScanStateEncoder(
  * It uses the first byte of the generated byte array to store the version the describes how the
  * row is encoded in the rest of the byte array. Currently, the default version is 0,
  *
+ * If the avroEnc is specified, we are using Avro encoding for this column family's keys
  * VERSION 0:  [ VERSION (1 byte) | ROW (N bytes) ]
  *    The bytes of a UnsafeRow is written unmodified to starting from offset 1
  *    (offset 0 is the version byte of value 0). That is, if the unsafe row has N bytes,
  *    then the generated array byte will be N+1 bytes.
- * If the avroEnc is specified, we are using Avro encoding for this column family's keys
  */
 class NoPrefixKeyStateEncoder(
     keySchema: StructType,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/package.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/package.scala
index 19a90c6978df0..e1a95dd10be74 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/package.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/package.scala
@@ -89,49 +89,6 @@ package object state {
         extraOptions,
         useMultipleValuesPerKey)
     }
-
-    /** Map each partition of an RDD along with data in a [[StateStore]] that passes the
-     * column family schemas to the storeUpdateFunction. Used to pass Avro encoders/decoders
-     * to executors */
-    def mapPartitionsWithStateStoreWithSchemas[U: ClassTag](
-        stateInfo: StatefulOperatorStateInfo,
-        keySchema: StructType,
-        valueSchema: StructType,
-        keyStateEncoderSpec: KeyStateEncoderSpec,
-        sessionState: SessionState,
-        storeCoordinator: Option[StateStoreCoordinatorRef],
-        useColumnFamilies: Boolean = false,
-        extraOptions: Map[String, String] = Map.empty,
-        useMultipleValuesPerKey: Boolean = false,
-        columnFamilySchemas: Map[String, StateStoreColFamilySchema] = Map.empty)(
-        storeUpdateFunction: (StateStore, Iterator[T], Map[String, StateStoreColFamilySchema]) => Iterator[U]): StateStoreRDD[T, U] = {
-
-      val cleanedF = dataRDD.sparkContext.clean(storeUpdateFunction)
-      val wrappedF = (store: StateStore, iter: Iterator[T]) => {
-        // Abort the state store in case of error
-        TaskContext.get().addTaskCompletionListener[Unit](_ => {
-          if (!store.hasCommitted) store.abort()
-        })
-        cleanedF(store, iter, columnFamilySchemas)
-      }
-
-      new StateStoreRDD(
-        dataRDD,
-        wrappedF,
-        stateInfo.checkpointLocation,
-        stateInfo.queryRunId,
-        stateInfo.operatorId,
-        stateInfo.storeVersion,
-        stateInfo.stateStoreCkptIds,
-        keySchema,
-        valueSchema,
-        keyStateEncoderSpec,
-        sessionState,
-        storeCoordinator,
-        useColumnFamilies,
-        extraOptions,
-        useMultipleValuesPerKey)
-    }
     // scalastyle:on
 
     /** Map each partition of an RDD along with data in a [[ReadStateStore]]. */

From 754ce6ca0fc5f7268bfc0ae604327611a9d35823 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Sat, 9 Nov 2024 12:57:31 -0800
Subject: [PATCH 25/30] batch works

---
 .../StateStoreColumnFamilySchemaUtils.scala   |  5 +++--
 .../StreamingSymmetricHashJoinHelper.scala    | 15 +--------------
 .../streaming/TransformWithStateExec.scala    | 19 ++++++++++---------
 3 files changed, 14 insertions(+), 25 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
index bb296e2f4aae0..ab15bf78a03ca 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
@@ -24,7 +24,7 @@ import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSch
 import org.apache.spark.sql.execution.streaming.state.{AvroEncoder, NoPrefixKeyStateEncoderSpec, PrefixKeyScanStateEncoderSpec, RangeKeyScanStateEncoderSpec, StateStoreColFamilySchema}
 import org.apache.spark.sql.types.{BinaryType, BooleanType, ByteType, DataType, DoubleType, FloatType, IntegerType, LongType, NullType, ShortType, StructField, StructType}
 
-object StateStoreColumnFamilySchemaUtils {
+object StateStoreColumnFamilySchemaUtils extends Serializable {
 
   def apply(initializeAvroSerde: Boolean): StateStoreColumnFamilySchemaUtils =
     new StateStoreColumnFamilySchemaUtils(initializeAvroSerde)
@@ -73,7 +73,8 @@ object StateStoreColumnFamilySchemaUtils {
  *                            for this state type. This class is used to create the
  *                            StateStoreColumnFamilySchema for each state variable from the driver
  */
-class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) extends Logging {
+class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean)
+    extends Logging with Serializable {
   private def getAvroSerializer(schema: StructType): AvroSerializer = {
     val avroType = SchemaConverters.toAvroType(schema)
     new AvroSerializer(schema, avroType, nullable = false)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamingSymmetricHashJoinHelper.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamingSymmetricHashJoinHelper.scala
index fe3b1683fce0c..497e71070a09a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamingSymmetricHashJoinHelper.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamingSymmetricHashJoinHelper.scala
@@ -27,7 +27,7 @@ import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeSet,
 import org.apache.spark.sql.catalyst.plans.logical.EventTimeWatermark._
 import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.execution.streaming.WatermarkSupport.watermarkExpression
-import org.apache.spark.sql.execution.streaming.state.{StateStoreCheckpointInfo, StateStoreColFamilySchema, StateStoreCoordinatorRef, StateStoreProviderId}
+import org.apache.spark.sql.execution.streaming.state.{StateStoreCheckpointInfo, StateStoreCoordinatorRef, StateStoreProviderId}
 
 
 /**
@@ -319,19 +319,6 @@ object StreamingSymmetricHashJoinHelper extends Logging {
       new StateStoreAwareZipPartitionsRDD(
         dataRDD.sparkContext, f, dataRDD, dataRDD2, stateInfo, storeNames, Some(storeCoordinator))
     }
-
-    def stateStoreAwareZipPartitions[U: ClassTag, V: ClassTag](
-        dataRDD2: RDD[U],
-        stateInfo: StatefulOperatorStateInfo,
-        storeNames: Seq[String],
-        storeCoordinator: StateStoreCoordinatorRef,
-        schemas: Map[String, StateStoreColFamilySchema]
-      )(f: (Int, Iterator[T], Iterator[U], Map[String, StateStoreColFamilySchema]) =>
-      Iterator[V]): RDD[V] = {
-      new StateStoreAwareZipPartitionsRDDWithSchemas(
-        dataRDD.sparkContext, f, dataRDD, dataRDD2, stateInfo,
-        storeNames, Some(storeCoordinator), schemas)
-    }
   }
 
   case class JoinerStateStoreCkptInfo(
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
index c1900e6c5e25f..277820ce6206b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
@@ -23,6 +23,7 @@ import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 
 import org.apache.spark.broadcast.Broadcast
+import org.apache.spark.internal.Logging
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
@@ -559,8 +560,9 @@ case class TransformWithStateExec(
               store, childDataIterator, initStateIterator, columnFamilySchemas)
           } else {
             initNewStateStoreAndProcessData(
-              partitionId, hadoopConfBroadcast, columnFamilySchemas) { (store, schemas) =>
-              processDataWithInitialState(store, childDataIterator, initStateIterator, schemas)
+              partitionId, hadoopConfBroadcast) { store =>
+              processDataWithInitialState(
+                store, childDataIterator, initStateIterator, columnFamilySchemas)
             }
           }
       }
@@ -586,8 +588,8 @@ case class TransformWithStateExec(
         child.execute().mapPartitionsWithIndex[InternalRow](
           (i: Int, iter: Iterator[InternalRow]) => {
             initNewStateStoreAndProcessData(
-              i, hadoopConfBroadcast, columnFamilySchemas) { (store, schemas) =>
-              processData(store, iter, schemas)
+              i, hadoopConfBroadcast) { store =>
+              processData(store, iter, columnFamilySchemas)
             }
           }
         )
@@ -601,9 +603,8 @@ case class TransformWithStateExec(
    */
   private def initNewStateStoreAndProcessData(
       partitionId: Int,
-      hadoopConfBroadcast: Broadcast[SerializableConfiguration],
-      schemas: Map[String, StateStoreColFamilySchema])
-    (f: (StateStore, Map[String, StateStoreColFamilySchema]) =>
+      hadoopConfBroadcast: Broadcast[SerializableConfiguration])
+    (f: StateStore =>
       CompletionIterator[InternalRow, Iterator[InternalRow]]):
     CompletionIterator[InternalRow, Iterator[InternalRow]] = {
 
@@ -630,7 +631,7 @@ case class TransformWithStateExec(
       useMultipleValuesPerKey = true)
 
     val store = stateStoreProvider.getStore(0)
-    val outputIterator = f(store, schemas)
+    val outputIterator = f(store)
     CompletionIterator[InternalRow, Iterator[InternalRow]](outputIterator.iterator, {
       stateStoreProvider.close()
       statefulProcessor.close()
@@ -702,7 +703,7 @@ case class TransformWithStateExec(
 }
 
 // scalastyle:off argcount
-object TransformWithStateExec {
+object TransformWithStateExec extends Logging {
 
   // Plan logical transformWithState for batch queries
   def generateSparkPlanForBatchQueries(

From b6dbfdb7ae2665f96c595d7bb079c8a2504306a2 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Sun, 10 Nov 2024 20:41:43 -0800
Subject: [PATCH 26/30] added test suite for non-contiguous ordinals

---
 .../streaming/state/RocksDBStateEncoder.scala | 175 +++++++++++++++---
 .../state/RocksDBStateStoreSuite.scala        |  79 +++++++-
 .../streaming/state/StateStoreSuite.scala     |   6 +
 3 files changed, 236 insertions(+), 24 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
index afeec3988538b..de9f2ab4d2f85 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
@@ -159,7 +159,7 @@ object RocksDBStateEncoder extends Logging {
   /**
    * This method takes an UnsafeRow, and serializes to a byte array using Avro encoding.
    */
-  def encodeUnsafeRow(
+  def encodeUnsafeRowToAvro(
      row: UnsafeRow,
      avroSerializer: AvroSerializer,
      valueAvroType: Schema,
@@ -189,7 +189,7 @@ object RocksDBStateEncoder extends Logging {
    * This method takes a byte array written using Avro encoding, and
    * deserializes to an UnsafeRow using the Avro deserializer
    */
-  def decodeToUnsafeRow(
+  def decodeFromAvroToUnsafeRow(
       valueBytes: Array[Byte],
       avroDeserializer: AvroDeserializer,
       valueAvroType: Schema,
@@ -279,13 +279,13 @@ class PrefixKeyScanStateEncoder(
   override def encodeKey(row: UnsafeRow): Array[Byte] = {
     val (prefixKeyEncoded, remainingEncoded) = if (usingAvroEncoding) {
       (
-        encodeUnsafeRow(
+        encodeUnsafeRowToAvro(
           extractPrefixKey(row),
           avroEnc.get.keySerializer,
           prefixKeyAvroType,
           out
         ),
-        encodeUnsafeRow(
+        encodeUnsafeRowToAvro(
           remainingKeyProjection(row),
           avroEnc.get.suffixKeySerializer.get,
           remainingKeyAvroType,
@@ -327,13 +327,13 @@ class PrefixKeyScanStateEncoder(
 
     val (prefixKeyDecoded, remainingKeyDecoded) = if (usingAvroEncoding) {
       (
-        decodeToUnsafeRow(
+        decodeFromAvroToUnsafeRow(
           prefixKeyEncoded,
           avroEnc.get.keyDeserializer,
           prefixKeyAvroType,
           prefixKeyProj
         ),
-        decodeToUnsafeRow(
+        decodeFromAvroToUnsafeRow(
           remainingKeyEncoded,
           avroEnc.get.suffixKeyDeserializer.get,
           remainingKeyAvroType,
@@ -354,7 +354,7 @@ class PrefixKeyScanStateEncoder(
 
   override def encodePrefixKey(prefixKey: UnsafeRow): Array[Byte] = {
     val prefixKeyEncoded = if (usingAvroEncoding) {
-      encodeUnsafeRow(prefixKey, avroEnc.get.keySerializer, prefixKeyAvroType, out)
+      encodeUnsafeRowToAvro(prefixKey, avroEnc.get.keySerializer, prefixKeyAvroType, out)
     } else {
       encodeUnsafeRow(prefixKey)
     }
@@ -483,6 +483,8 @@ class RangeKeyScanStateEncoder(
     StructType(rangeScanKeyFieldsWithOrdinal.map(_._1).toArray))
 
   private lazy val rangeScanAvroType = SchemaConverters.toAvroType(rangeScanAvroSchema)
+  logError(s"### rangeScanAvroSchema: $rangeScanAvroSchema")
+  logError(s"### rangeScanAvroType: $rangeScanAvroType")
 
   private val rangeScanAvroProjection = UnsafeProjection.create(rangeScanAvroSchema)
 
@@ -698,21 +700,97 @@ class RangeKeyScanStateEncoder(
         record.put(fieldIdx + 1, new Array[Byte](field.dataType.defaultSize))
       } else {
         field.dataType match {
+          case BooleanType =>
+            val boolVal = value.asInstanceOf[Boolean]
+            record.put(fieldIdx, true) // not null marker
+            record.put(fieldIdx + 1, ByteBuffer.wrap(Array[Byte](if (boolVal) 1 else 0)))
+
+          case ByteType =>
+            val byteVal = value.asInstanceOf[Byte]
+            val marker = byteVal >= 0
+            record.put(fieldIdx, marker)
+
+            val bytes = new Array[Byte](1)
+            bytes(0) = byteVal
+            record.put(fieldIdx + 1, ByteBuffer.wrap(bytes))
+
+          case ShortType =>
+            val shortVal = value.asInstanceOf[Short]
+            val marker = shortVal >= 0
+            record.put(fieldIdx, marker)
+
+            val bbuf = ByteBuffer.allocate(2)
+            bbuf.order(ByteOrder.BIG_ENDIAN)
+            bbuf.putShort(shortVal)
+            val bytes = new Array[Byte](2)
+            bbuf.position(0)
+            bbuf.get(bytes)
+            record.put(fieldIdx + 1, ByteBuffer.wrap(bytes))
+
+          case IntegerType =>
+            val intVal = value.asInstanceOf[Int]
+            val marker = intVal >= 0
+            record.put(fieldIdx, marker)
+
+            val bbuf = ByteBuffer.allocate(4)
+            bbuf.order(ByteOrder.BIG_ENDIAN)
+            bbuf.putInt(intVal)
+            val bytes = new Array[Byte](4)
+            bbuf.position(0)
+            bbuf.get(bytes)
+            record.put(fieldIdx + 1, ByteBuffer.wrap(bytes))
+
           case LongType =>
             val longVal = value.asInstanceOf[Long]
             val marker = longVal >= 0
             record.put(fieldIdx, marker)
 
-            // Convert long to byte array in big endian format
             val bbuf = ByteBuffer.allocate(8)
             bbuf.order(ByteOrder.BIG_ENDIAN)
             bbuf.putLong(longVal)
-            // Create a new byte array to avoid Avro's issue with direct ByteBuffer arrays
             val bytes = new Array[Byte](8)
             bbuf.position(0)
             bbuf.get(bytes)
+            record.put(fieldIdx + 1, ByteBuffer.wrap(bytes))
+
+          case FloatType =>
+            val floatVal = value.asInstanceOf[Float]
+            val rawBits = floatToRawIntBits(floatVal)
+            val marker = (rawBits & floatSignBitMask) == 0
+            record.put(fieldIdx, marker)
 
-            // Wrap bytes in Avro's ByteBuffer to ensure proper handling
+            val bbuf = ByteBuffer.allocate(4)
+            bbuf.order(ByteOrder.BIG_ENDIAN)
+            if (!marker) {
+              // For negative values, flip the bits to maintain proper ordering
+              val updatedVal = rawBits ^ floatFlipBitMask
+              bbuf.putFloat(intBitsToFloat(updatedVal))
+            } else {
+              bbuf.putFloat(floatVal)
+            }
+            val bytes = new Array[Byte](4)
+            bbuf.position(0)
+            bbuf.get(bytes)
+            record.put(fieldIdx + 1, ByteBuffer.wrap(bytes))
+
+          case DoubleType =>
+            val doubleVal = value.asInstanceOf[Double]
+            val rawBits = doubleToRawLongBits(doubleVal)
+            val marker = (rawBits & doubleSignBitMask) == 0
+            record.put(fieldIdx, marker)
+
+            val bbuf = ByteBuffer.allocate(8)
+            bbuf.order(ByteOrder.BIG_ENDIAN)
+            if (!marker) {
+              // For negative values, flip the bits to maintain proper ordering
+              val updatedVal = rawBits ^ doubleFlipBitMask
+              bbuf.putDouble(longBitsToDouble(updatedVal))
+            } else {
+              bbuf.putDouble(doubleVal)
+            }
+            val bytes = new Array[Byte](8)
+            bbuf.position(0)
+            bbuf.get(bytes)
             record.put(fieldIdx + 1, ByteBuffer.wrap(bytes))
 
           case _ => throw new UnsupportedOperationException(
@@ -750,18 +828,69 @@ class RangeKeyScanStateEncoder(
         rowWriter.setNullAt(idx)
       } else {
         field.dataType match {
+          case BooleanType =>
+            val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+            rowWriter.write(idx, bytes(0) == 1)
+
+          case ByteType =>
+            val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+            rowWriter.write(idx, bytes(0))
+
+          case ShortType =>
+            val byteBuf = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer]
+            val bbuf = ByteBuffer.allocate(2)
+            bbuf.order(ByteOrder.BIG_ENDIAN)
+            bbuf.put(byteBuf.array(), byteBuf.position(), byteBuf.remaining())
+            bbuf.flip()
+            rowWriter.write(idx, bbuf.getShort)
+
+          case IntegerType =>
+            val byteBuf = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer]
+            val bbuf = ByteBuffer.allocate(4)
+            bbuf.order(ByteOrder.BIG_ENDIAN)
+            bbuf.put(byteBuf.array(), byteBuf.position(), byteBuf.remaining())
+            bbuf.flip()
+            rowWriter.write(idx, bbuf.getInt)
+
           case LongType =>
-            // Get bytes from Avro ByteBuffer
             val byteBuf = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer]
             val bbuf = ByteBuffer.allocate(8)
             bbuf.order(ByteOrder.BIG_ENDIAN)
+            bbuf.put(byteBuf.array(), byteBuf.position(), byteBuf.remaining())
+            bbuf.flip()
+            rowWriter.write(idx, bbuf.getLong)
 
-            // Copy bytes to our ByteBuffer
+          case FloatType =>
+            val byteBuf = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer]
+            val bbuf = ByteBuffer.allocate(4)
+            bbuf.order(ByteOrder.BIG_ENDIAN)
             bbuf.put(byteBuf.array(), byteBuf.position(), byteBuf.remaining())
             bbuf.flip()
 
-            val longVal = bbuf.getLong(fieldIdx)
-            rowWriter.write(idx, longVal)
+            val isNegative = !record.get(fieldIdx).asInstanceOf[Boolean]
+            if (isNegative) {
+              val floatVal = bbuf.getFloat
+              val updatedVal = floatToRawIntBits(floatVal) ^ floatFlipBitMask
+              rowWriter.write(idx, intBitsToFloat(updatedVal))
+            } else {
+              rowWriter.write(idx, bbuf.getFloat)
+            }
+
+          case DoubleType =>
+            val byteBuf = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer]
+            val bbuf = ByteBuffer.allocate(8)
+            bbuf.order(ByteOrder.BIG_ENDIAN)
+            bbuf.put(byteBuf.array(), byteBuf.position(), byteBuf.remaining())
+            bbuf.flip()
+
+            val isNegative = !record.get(fieldIdx).asInstanceOf[Boolean]
+            if (isNegative) {
+              val doubleVal = bbuf.getDouble
+              val updatedVal = doubleToRawLongBits(doubleVal) ^ doubleFlipBitMask
+              rowWriter.write(idx, longBitsToDouble(updatedVal))
+            } else {
+              rowWriter.write(idx, bbuf.getDouble)
+            }
 
           case _ => throw new UnsupportedOperationException(
             s"Range scan decoding not supported for data type: ${field.dataType}")
@@ -784,7 +913,7 @@ class RangeKeyScanStateEncoder(
 
     val result = if (orderingOrdinals.length < keySchema.length) {
       val remainingEncoded = if (avroEnc.isDefined) {
-        encodeUnsafeRow(
+        encodeUnsafeRowToAvro(
           remainingKeyProjection(row),
           avroEnc.get.suffixKeySerializer.get,
           remainingKeyAvroType,
@@ -847,7 +976,7 @@ class RangeKeyScanStateEncoder(
         remainingKeyEncodedLen)
 
       val remainingKeyDecoded = if (avroEnc.isDefined) {
-        decodeToUnsafeRow(remainingKeyEncoded,
+        decodeFromAvroToUnsafeRow(remainingKeyEncoded,
           avroEnc.get.suffixKeyDeserializer.get,
           remainingKeyAvroType, remainingKeyAvroProjection)
       } else {
@@ -917,7 +1046,7 @@ class NoPrefixKeyStateEncoder(
       // If avroEnc is defined, we know that we need to use Avro to
       // encode this UnsafeRow to Avro bytes
       val bytesToEncode = if (usingAvroEncoding) {
-        encodeUnsafeRow(row, avroEnc.get.keySerializer, keyAvroType, out)
+        encodeUnsafeRowToAvro(row, avroEnc.get.keySerializer, keyAvroType, out)
       } else row.getBytes
       val (encodedBytes, startingOffset) = encodeColumnFamilyPrefix(
         bytesToEncode.length +
@@ -952,7 +1081,7 @@ class NoPrefixKeyStateEncoder(
             0,
             avroBytes.length
           )
-          decodeToUnsafeRow(
+          decodeFromAvroToUnsafeRow(
             keyBytes, avroEnc.get.keyDeserializer, keyAvroType, keyProj)
         } else {
           keyRow.pointTo(
@@ -1004,7 +1133,7 @@ class MultiValuedStateEncoder(
 
   override def encodeValue(row: UnsafeRow): Array[Byte] = {
     val bytes = if (usingAvroEncoding) {
-      encodeUnsafeRow(row, avroEnc.get.valueSerializer, valueAvroType, out)
+      encodeUnsafeRowToAvro(row, avroEnc.get.valueSerializer, valueAvroType, out)
     } else {
       encodeUnsafeRow(row)
     }
@@ -1027,7 +1156,7 @@ class MultiValuedStateEncoder(
       Platform.copyMemory(valueBytes, java.lang.Integer.BYTES + Platform.BYTE_ARRAY_OFFSET,
         encodedValue, Platform.BYTE_ARRAY_OFFSET, numBytes)
       if (usingAvroEncoding) {
-        decodeToUnsafeRow(
+        decodeFromAvroToUnsafeRow(
           encodedValue, avroEnc.get.valueDeserializer, valueAvroType, valueProj)
       } else {
         decodeToUnsafeRow(encodedValue, valueRow)
@@ -1058,7 +1187,7 @@ class MultiValuedStateEncoder(
           pos += numBytes
           pos += 1 // eat the delimiter character
           if (usingAvroEncoding) {
-            decodeToUnsafeRow(
+            decodeFromAvroToUnsafeRow(
               encodedValue, avroEnc.get.valueDeserializer, valueAvroType, valueProj)
           } else {
             decodeToUnsafeRow(encodedValue, valueRow)
@@ -1100,7 +1229,7 @@ class SingleValueStateEncoder(
 
   override def encodeValue(row: UnsafeRow): Array[Byte] = {
     if (usingAvroEncoding) {
-      encodeUnsafeRow(row, avroEnc.get.valueSerializer, valueAvroType, out)
+      encodeUnsafeRowToAvro(row, avroEnc.get.valueSerializer, valueAvroType, out)
     } else {
       encodeUnsafeRow(row)
     }
@@ -1117,7 +1246,7 @@ class SingleValueStateEncoder(
       return null
     }
     if (usingAvroEncoding) {
-      decodeToUnsafeRow(
+      decodeFromAvroToUnsafeRow(
         valueBytes, avroEnc.get.valueDeserializer, valueAvroType, valueProj)
     } else {
       decodeToUnsafeRow(valueBytes, valueRow)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala
index e1bd9dd38066b..e5f69fba1a97b 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala
@@ -31,7 +31,7 @@ import org.apache.spark.sql.LocalSparkSession.withSparkSession
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeProjection, UnsafeRow}
 import org.apache.spark.sql.catalyst.util.quietly
-import org.apache.spark.sql.execution.streaming.StatefulOperatorStateInfo
+import org.apache.spark.sql.execution.streaming.{StatefulOperatorStateInfo, StateStoreColumnFamilySchemaUtils}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
 import org.apache.spark.sql.types._
@@ -339,6 +339,83 @@ class RocksDBStateStoreSuite extends StateStoreSuiteBase[RocksDBStateStoreProvid
     }
   }
 
+  test("rocksdb range scan - fixed size non-ordering columns with Avro encoding") {
+
+
+    val keySchemaWithLong: StructType = StructType(
+      Seq(StructField("key1", StringType, false), StructField("key2", LongType, false),
+        StructField("key3", StringType, false), StructField("key4", LongType, false)))
+
+    val remainingKeySchema: StructType = StructType(
+      Seq(StructField("key1", StringType, false), StructField("key3", StringType, false)))
+    tryWithProviderResource(newStoreProvider(keySchemaWithLong,
+      RangeKeyScanStateEncoderSpec(keySchemaWithLong, Seq(1, 3)),
+      useColumnFamilies = true)) { provider =>
+      val store = provider.getStore(0)
+
+      // use non-default col family if column families are enabled
+      val cfName = "testColFamily"
+      val convertedKeySchema = StateStoreColumnFamilySchemaUtils.convertForRangeScan(
+        keySchemaWithLong)
+      val avroSerde = StateStoreColumnFamilySchemaUtils(true).getAvroSerde(
+        StructType(convertedKeySchema.drop(1)),
+        valueSchema,
+        Some(remainingKeySchema)
+      )
+      store.createColFamilyIfAbsent(cfName,
+        keySchemaWithLong, valueSchema,
+        RangeKeyScanStateEncoderSpec(keySchemaWithLong, Seq(1, 3)),
+        avroEncoderSpec = avroSerde)
+
+      val timerTimestamps = Seq(931L, 8000L, 452300L, 4200L, -1L, 90L, 1L, 2L, 8L,
+        -230L, -14569L, -92L, -7434253L, 35L, 6L, 9L, -323L, 5L)
+      val otherLongs = Seq(3L, 2L, 1L)
+
+      // Create all combinations using flatMap
+      val testPairs = timerTimestamps.flatMap { ts1 =>
+        timerTimestamps.map { ts2 =>
+          (ts1, ts2)
+        }
+      }
+
+      testPairs.foreach { ts =>
+        // non-timestamp col is of fixed size
+        val keyRow = dataToKeyRowWithRangeScan("a", ts._1, ts._2)
+        val valueRow = dataToValueRow(1)
+        store.put(keyRow, valueRow, cfName)
+        assert(valueRowToData(store.get(keyRow, cfName)) === 1)
+      }
+
+      val result = store.iterator(cfName).map { kv =>
+        (kv.key.getLong(1), kv.key.getLong(3))
+      }.toSeq
+      assert(result === testPairs.sortBy(pair => (pair._1, pair._2)))
+      store.commit()
+
+      // test with a different set of power of 2 timestamps
+      val store1 = provider.getStore(1)
+      val timerTimestamps1 = Seq(-32L, -64L, -256L, 64L, 32L, 1024L, 4096L, 0L)
+      val testPairs1 = timerTimestamps1.flatMap { ts1 =>
+        otherLongs.map { ts2 =>
+          (ts1, ts2)
+        }
+      }
+      testPairs1.foreach { ts =>
+        // non-timestamp col is of fixed size
+        val keyRow = dataToKeyRowWithRangeScan("a", ts._1, ts._2)
+        val valueRow = dataToValueRow(1)
+        store1.put(keyRow, valueRow, cfName)
+        assert(valueRowToData(store1.get(keyRow, cfName)) === 1)
+      }
+
+      val result1 = store1.iterator(cfName).map { kv =>
+        (kv.key.getLong(1), kv.key.getLong(3))
+      }.toSeq
+      logError(s"### result1: ${result1}")
+      assert(result1 === (testPairs ++ testPairs1).sortBy(pair => (pair._1, pair._2)))
+    }
+  }
+
   testWithColumnFamilies("rocksdb range scan - variable size non-ordering columns with " +
     "double type values are supported",
     TestWithBothChangelogCheckpointingEnabledAndDisabled) { colFamiliesEnabled =>
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreSuite.scala
index 47dd77f1bb9fd..90d8b157d94fa 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreSuite.scala
@@ -1847,6 +1847,12 @@ object StateStoreTestsHelper {
     rangeScanProj.apply(new GenericInternalRow(Array[Any](ts, UTF8String.fromString(s)))).copy()
   }
 
+  def dataToKeyRowWithRangeScan(s: String, ts: Long, otherLong: Long): UnsafeRow = {
+    UnsafeProjection.create(Array[DataType](StringType, LongType, StringType, LongType))
+      .apply(new GenericInternalRow(Array[Any](UTF8String.fromString(s), ts,
+        UTF8String.fromString(s), otherLong))).copy()
+  }
+
   def dataToValueRow(i: Int): UnsafeRow = {
     valueProj.apply(new GenericInternalRow(Array[Any](i))).copy()
   }

From e6f0b7a606ee81ab7efe093fe408471cc17f4213 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Sun, 10 Nov 2024 20:58:54 -0800
Subject: [PATCH 27/30] using negative/null val marker

---
 .../StateStoreColumnFamilySchemaUtils.scala   |  8 ++---
 .../streaming/state/RocksDBStateEncoder.scala | 32 +++++++++----------
 .../state/RocksDBStateStoreSuite.scala        |  2 +-
 3 files changed, 20 insertions(+), 22 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
index ab15bf78a03ca..5662e91c9926a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
@@ -47,7 +47,7 @@ object StateStoreColumnFamilySchemaUtils extends Serializable {
         // 1. A boolean for sign (positive = true, negative = false)
         // 2. The original numeric value in big-endian format
         Seq(
-          StructField(s"${field.name}_marker", BooleanType, nullable = false),
+          StructField(s"${field.name}_marker", ByteType, nullable = false),
           field.copy(name = s"${field.name}_value", BinaryType)
         )
       } else {
@@ -184,7 +184,7 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean)
       ttlValSchema,
       Some(RangeKeyScanStateEncoderSpec(ttlKeySchema, Seq(0))),
       avroEnc = getAvroSerde(
-        StructType(ttlKeySchema.take(2)),
+        getSingleKeyTTLRowSchema(keyEncoder.schema),
         ttlValSchema,
         Some(StructType(ttlKeySchema.drop(2)))
       )
@@ -210,7 +210,7 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean)
       ttlValSchema,
       Some(RangeKeyScanStateEncoderSpec(ttlKeySchema, Seq(0))),
       avroEnc = getAvroSerde(
-        StructType(ttlKeySchema.take(2)),
+        getCompositeKeyTTLRowSchema(keyEncoder.schema, userKeySchema),
         ttlValSchema,
         Some(StructType(ttlKeySchema.drop(2)))
       )
@@ -250,7 +250,7 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean)
       valSchema,
       Some(RangeKeyScanStateEncoderSpec(keySchema, Seq(0))),
       avroEnc = getAvroSerde(
-        StructType(avroKeySchema.take(2)),
+        keySchema,
         valSchema,
         Some(StructType(avroKeySchema.drop(2)))
       ))
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
index de9f2ab4d2f85..c69cf6efc813b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
@@ -483,8 +483,6 @@ class RangeKeyScanStateEncoder(
     StructType(rangeScanKeyFieldsWithOrdinal.map(_._1).toArray))
 
   private lazy val rangeScanAvroType = SchemaConverters.toAvroType(rangeScanAvroSchema)
-  logError(s"### rangeScanAvroSchema: $rangeScanAvroSchema")
-  logError(s"### rangeScanAvroType: $rangeScanAvroType")
 
   private val rangeScanAvroProjection = UnsafeProjection.create(rangeScanAvroSchema)
 
@@ -696,18 +694,18 @@ class RangeKeyScanStateEncoder(
       val field = fieldWithOrdinal._1
       val value = row.get(idx, field.dataType)
       if (value == null) {
-        record.put(fieldIdx, false) // isNull marker
+        record.put(fieldIdx, nullValMarker) // isNull marker
         record.put(fieldIdx + 1, new Array[Byte](field.dataType.defaultSize))
       } else {
         field.dataType match {
           case BooleanType =>
             val boolVal = value.asInstanceOf[Boolean]
-            record.put(fieldIdx, true) // not null marker
+            record.put(fieldIdx, positiveValMarker) // not null marker
             record.put(fieldIdx + 1, ByteBuffer.wrap(Array[Byte](if (boolVal) 1 else 0)))
 
           case ByteType =>
             val byteVal = value.asInstanceOf[Byte]
-            val marker = byteVal >= 0
+            val marker = positiveValMarker
             record.put(fieldIdx, marker)
 
             val bytes = new Array[Byte](1)
@@ -716,7 +714,7 @@ class RangeKeyScanStateEncoder(
 
           case ShortType =>
             val shortVal = value.asInstanceOf[Short]
-            val marker = shortVal >= 0
+            val marker = if (shortVal >= 0) positiveValMarker else negativeValMarker
             record.put(fieldIdx, marker)
 
             val bbuf = ByteBuffer.allocate(2)
@@ -729,7 +727,7 @@ class RangeKeyScanStateEncoder(
 
           case IntegerType =>
             val intVal = value.asInstanceOf[Int]
-            val marker = intVal >= 0
+            val marker = if (intVal >= 0) positiveValMarker else negativeValMarker
             record.put(fieldIdx, marker)
 
             val bbuf = ByteBuffer.allocate(4)
@@ -742,7 +740,7 @@ class RangeKeyScanStateEncoder(
 
           case LongType =>
             val longVal = value.asInstanceOf[Long]
-            val marker = longVal >= 0
+            val marker = if (longVal >= 0) positiveValMarker else negativeValMarker
             record.put(fieldIdx, marker)
 
             val bbuf = ByteBuffer.allocate(8)
@@ -756,16 +754,16 @@ class RangeKeyScanStateEncoder(
           case FloatType =>
             val floatVal = value.asInstanceOf[Float]
             val rawBits = floatToRawIntBits(floatVal)
-            val marker = (rawBits & floatSignBitMask) == 0
-            record.put(fieldIdx, marker)
 
             val bbuf = ByteBuffer.allocate(4)
             bbuf.order(ByteOrder.BIG_ENDIAN)
-            if (!marker) {
+            if ((rawBits & floatSignBitMask) != 0) {
+              record.put(fieldIdx, negativeValMarker)
               // For negative values, flip the bits to maintain proper ordering
               val updatedVal = rawBits ^ floatFlipBitMask
               bbuf.putFloat(intBitsToFloat(updatedVal))
             } else {
+              record.put(fieldIdx, positiveValMarker)
               bbuf.putFloat(floatVal)
             }
             val bytes = new Array[Byte](4)
@@ -776,16 +774,16 @@ class RangeKeyScanStateEncoder(
           case DoubleType =>
             val doubleVal = value.asInstanceOf[Double]
             val rawBits = doubleToRawLongBits(doubleVal)
-            val marker = (rawBits & doubleSignBitMask) == 0
-            record.put(fieldIdx, marker)
 
             val bbuf = ByteBuffer.allocate(8)
             bbuf.order(ByteOrder.BIG_ENDIAN)
-            if (!marker) {
+            if ((rawBits & doubleSignBitMask) != 0) {
               // For negative values, flip the bits to maintain proper ordering
+              record.put(fieldIdx, negativeValMarker)
               val updatedVal = rawBits ^ doubleFlipBitMask
               bbuf.putDouble(longBitsToDouble(updatedVal))
             } else {
+              record.put(fieldIdx, positiveValMarker)
               bbuf.putDouble(doubleVal)
             }
             val bytes = new Array[Byte](8)
@@ -822,7 +820,7 @@ class RangeKeyScanStateEncoder(
     var fieldIdx = 0
     rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
       val field = fieldWithOrdinal._1
-      val isMarkerNull = record.get(fieldIdx) == null
+      val isMarkerNull = record.get(fieldIdx) == nullValMarker
 
       if (isMarkerNull) {
         rowWriter.setNullAt(idx)
@@ -867,7 +865,7 @@ class RangeKeyScanStateEncoder(
             bbuf.put(byteBuf.array(), byteBuf.position(), byteBuf.remaining())
             bbuf.flip()
 
-            val isNegative = !record.get(fieldIdx).asInstanceOf[Boolean]
+            val isNegative = record.get(fieldIdx).asInstanceOf[Byte] == negativeValMarker
             if (isNegative) {
               val floatVal = bbuf.getFloat
               val updatedVal = floatToRawIntBits(floatVal) ^ floatFlipBitMask
@@ -883,7 +881,7 @@ class RangeKeyScanStateEncoder(
             bbuf.put(byteBuf.array(), byteBuf.position(), byteBuf.remaining())
             bbuf.flip()
 
-            val isNegative = !record.get(fieldIdx).asInstanceOf[Boolean]
+            val isNegative = record.get(fieldIdx).asInstanceOf[Byte] == negativeValMarker
             if (isNegative) {
               val doubleVal = bbuf.getDouble
               val updatedVal = doubleToRawLongBits(doubleVal) ^ doubleFlipBitMask
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala
index e5f69fba1a97b..7755f3b6980d2 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala
@@ -358,7 +358,7 @@ class RocksDBStateStoreSuite extends StateStoreSuiteBase[RocksDBStateStoreProvid
       val convertedKeySchema = StateStoreColumnFamilySchemaUtils.convertForRangeScan(
         keySchemaWithLong)
       val avroSerde = StateStoreColumnFamilySchemaUtils(true).getAvroSerde(
-        StructType(convertedKeySchema.drop(1)),
+        keySchemaWithLong,
         valueSchema,
         Some(remainingKeySchema)
       )

From ca660c0b6c3e8633237bdeb60390fcfd55f574be Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Sun, 10 Nov 2024 21:01:09 -0800
Subject: [PATCH 28/30] removing log line

---
 .../sql/execution/streaming/state/RocksDBStateStoreSuite.scala   | 1 -
 1 file changed, 1 deletion(-)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala
index 7755f3b6980d2..c7a895d165037 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala
@@ -411,7 +411,6 @@ class RocksDBStateStoreSuite extends StateStoreSuiteBase[RocksDBStateStoreProvid
       val result1 = store1.iterator(cfName).map { kv =>
         (kv.key.getLong(1), kv.key.getLong(3))
       }.toSeq
-      logError(s"### result1: ${result1}")
       assert(result1 === (testPairs ++ testPairs1).sortBy(pair => (pair._1, pair._2)))
     }
   }

From 41de8aee1d079ca34d34c0731f3c279f4247a254 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Mon, 11 Nov 2024 14:44:22 -0800
Subject: [PATCH 29/30] getAvroEnc

---
 .../StatefulProcessorHandleImpl.scala         | 24 ++++++++++++-------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
index 1fd244ea167dd..8d5ad2ec11098 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
@@ -122,6 +122,14 @@ class StatefulProcessorHandleImpl(
 
   currState = CREATED
 
+  private def getAvroEnc(stateName: String): Option[AvroEncoder] = {
+    if (!schemas.contains(stateName)) {
+      None
+    } else {
+      schemas(stateName).avroEnc
+    }
+  }
+
   private def buildQueryInfo(): QueryInfo = {
     val taskCtxOpt = Option(TaskContext.get())
     val (queryId, batchId) = if (!isStreaming) {
@@ -146,8 +154,8 @@ class StatefulProcessorHandleImpl(
   private lazy val timerSecIndexColFamily = TimerStateUtils.getSecIndexColFamilyName(
     timeMode.toString)
   private lazy val timerState = new TimerStateImpl(
-    store, timeMode, keyEncoder, schemas(timerStateName).avroEnc,
-    schemas(timerSecIndexColFamily).avroEnc)
+    store, timeMode, keyEncoder, getAvroEnc(timerStateName),
+    getAvroEnc(timerSecIndexColFamily))
 
   /**
    * Function to register a timer for the given expiryTimestampMs
@@ -238,13 +246,13 @@ class StatefulProcessorHandleImpl(
       assert(batchTimestampMs.isDefined)
       val valueStateWithTTL = new ValueStateImplWithTTL[T](store, stateName,
         keyEncoder, stateEncoder, ttlConfig, batchTimestampMs.get, metrics,
-        schemas(stateName).avroEnc, schemas(getTtlColFamilyName(stateName)).avroEnc)
+        getAvroEnc(stateName), getAvroEnc(getTtlColFamilyName(stateName)))
       ttlStates.add(valueStateWithTTL)
       TWSMetricsUtils.incrementMetric(metrics, "numValueStateWithTTLVars")
       valueStateWithTTL
     } else {
       val valueStateWithoutTTL = new ValueStateImpl[T](store, stateName,
-        keyEncoder, stateEncoder, metrics, schemas(stateName).avroEnc)
+        keyEncoder, stateEncoder, metrics, getAvroEnc(stateName))
       TWSMetricsUtils.incrementMetric(metrics, "numValueStateVars")
       valueStateWithoutTTL
     }
@@ -288,13 +296,13 @@ class StatefulProcessorHandleImpl(
       assert(batchTimestampMs.isDefined)
       val listStateWithTTL = new ListStateImplWithTTL[T](store, stateName,
         keyEncoder, stateEncoder, ttlConfig, batchTimestampMs.get, metrics,
-        schemas(stateName).avroEnc, schemas(getTtlColFamilyName(stateName)).avroEnc)
+        getAvroEnc(stateName), getAvroEnc(getTtlColFamilyName(stateName)))
       TWSMetricsUtils.incrementMetric(metrics, "numListStateWithTTLVars")
       ttlStates.add(listStateWithTTL)
       listStateWithTTL
     } else {
       val listStateWithoutTTL = new ListStateImpl[T](store, stateName, keyEncoder,
-        stateEncoder, metrics, schemas(stateName).avroEnc)
+        stateEncoder, metrics, getAvroEnc(stateName))
       TWSMetricsUtils.incrementMetric(metrics, "numListStateVars")
       listStateWithoutTTL
     }
@@ -327,13 +335,13 @@ class StatefulProcessorHandleImpl(
       assert(batchTimestampMs.isDefined)
       val mapStateWithTTL = new MapStateImplWithTTL[K, V](store, stateName, keyEncoder, userKeyEnc,
         valEncoder, ttlConfig, batchTimestampMs.get, metrics,
-        schemas(stateName).avroEnc, schemas(getTtlColFamilyName(stateName)).avroEnc)
+        getAvroEnc(stateName), getAvroEnc(getTtlColFamilyName(stateName)))
       TWSMetricsUtils.incrementMetric(metrics, "numMapStateWithTTLVars")
       ttlStates.add(mapStateWithTTL)
       mapStateWithTTL
     } else {
       val mapStateWithoutTTL = new MapStateImpl[K, V](store, stateName, keyEncoder,
-        userKeyEnc, valEncoder, metrics, schemas(stateName).avroEnc)
+        userKeyEnc, valEncoder, metrics, getAvroEnc(stateName))
       TWSMetricsUtils.incrementMetric(metrics, "numMapStateVars")
       mapStateWithoutTTL
     }

From c49acd28e8a965ce9462819a94cfcfc51920d092 Mon Sep 17 00:00:00 2001
From: Eric Marnadi <eric.marnadi@databricks.com>
Date: Tue, 5 Nov 2024 15:25:16 -0800
Subject: [PATCH 30/30] init

---
 .../streaming/TransformWithStateExec.scala    |  3 +-
 .../StateSchemaCompatibilityChecker.scala     | 40 ++++++++++++++-----
 2 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
index 277820ce6206b..adb7c27363ae8 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
@@ -475,7 +475,8 @@ case class TransformWithStateExec(
       newSchemas.values.toList, session.sessionState, stateSchemaVersion,
       storeName = StateStoreId.DEFAULT_STORE_NAME,
       oldSchemaFilePath = oldStateSchemaFilePath,
-      newSchemaFilePath = Some(newStateSchemaFilePath)))
+      newSchemaFilePath = Some(newStateSchemaFilePath),
+      usingAvro = true))
   }
 
   /** Metadata of this stateful operator and its states stores. */
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
index b5f2f318de418..5bb511d5d5567 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
@@ -17,14 +17,16 @@
 
 package org.apache.spark.sql.execution.streaming.state
 
+import scala.jdk.CollectionConverters.IterableHasAsJava
 import scala.util.Try
 
+import org.apache.avro.SchemaValidatorBuilder
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 
 import org.apache.spark.SparkUnsupportedOperationException
 import org.apache.spark.internal.{Logging, LogKeys, MDC}
-import org.apache.spark.sql.avro.{AvroDeserializer, AvroSerializer}
+import org.apache.spark.sql.avro.{AvroDeserializer, AvroSerializer, SchemaConverters}
 import org.apache.spark.sql.catalyst.util.UnsafeRowUtils
 import org.apache.spark.sql.execution.streaming.{CheckpointFileManager, StatefulOperatorStateInfo}
 import org.apache.spark.sql.execution.streaming.state.SchemaHelper.{SchemaReader, SchemaWriter}
@@ -151,7 +153,8 @@ class StateSchemaCompatibilityChecker(
   private def check(
       oldSchema: StateStoreColFamilySchema,
       newSchema: StateStoreColFamilySchema,
-      ignoreValueSchema: Boolean) : Unit = {
+      ignoreValueSchema: Boolean,
+      usingAvro: Boolean) : Boolean = {
     val (storedKeySchema, storedValueSchema) = (oldSchema.keySchema,
       oldSchema.valueSchema)
     val (keySchema, valueSchema) = (newSchema.keySchema, newSchema.valueSchema)
@@ -159,14 +162,27 @@ class StateSchemaCompatibilityChecker(
     if (storedKeySchema.equals(keySchema) &&
       (ignoreValueSchema || storedValueSchema.equals(valueSchema))) {
       // schema is exactly same
+      false
     } else if (!schemasCompatible(storedKeySchema, keySchema)) {
       throw StateStoreErrors.stateStoreKeySchemaNotCompatible(storedKeySchema.toString,
         keySchema.toString)
+    } else if (!ignoreValueSchema && usingAvro) {
+      // By this point, we know that old value schema is not equal to new value schema
+      val oldAvroSchema = SchemaConverters.toAvroType(storedValueSchema)
+      val newAvroSchema = SchemaConverters.toAvroType(valueSchema)
+      val validator = new SchemaValidatorBuilder().canReadStrategy.validateAll()
+      // This will throw a SchemaValidation exception if the schema has evolved in an
+      // unacceptable way.
+      validator.validate(newAvroSchema, Iterable(oldAvroSchema).asJava)
+      // If no exception is thrown, then we know that the schema evolved in an
+      // acceptable way
+      true
     } else if (!ignoreValueSchema && !schemasCompatible(storedValueSchema, valueSchema)) {
       throw StateStoreErrors.stateStoreValueSchemaNotCompatible(storedValueSchema.toString,
         valueSchema.toString)
     } else {
       logInfo("Detected schema change which is compatible. Allowing to put rows.")
+      true
     }
   }
 
@@ -180,7 +196,8 @@ class StateSchemaCompatibilityChecker(
   def validateAndMaybeEvolveStateSchema(
       newStateSchema: List[StateStoreColFamilySchema],
       ignoreValueSchema: Boolean,
-      stateSchemaVersion: Int): Boolean = {
+      stateSchemaVersion: Int,
+      usingAvro: Boolean): Boolean = {
     val existingStateSchemaList = getExistingKeyAndValueSchema()
     val newStateSchemaList = newStateSchema
 
@@ -195,18 +212,18 @@ class StateSchemaCompatibilityChecker(
       }.toMap
       // For each new state variable, we want to compare it to the old state variable
       // schema with the same name
-      newStateSchemaList.foreach { newSchema =>
-        existingSchemaMap.get(newSchema.colFamilyName).foreach { existingStateSchema =>
-          check(existingStateSchema, newSchema, ignoreValueSchema)
-        }
+      val hasEvolvedSchema = newStateSchemaList.exists { newSchema =>
+        existingSchemaMap.get(newSchema.colFamilyName)
+          .exists(existingSchema => check(existingSchema, newSchema, ignoreValueSchema, usingAvro))
       }
       val colFamiliesAddedOrRemoved =
         (newStateSchemaList.map(_.colFamilyName).toSet != existingSchemaMap.keySet)
-      if (stateSchemaVersion == SCHEMA_FORMAT_V3 && colFamiliesAddedOrRemoved) {
+      val newSchemaFileWritten = hasEvolvedSchema || colFamiliesAddedOrRemoved
+      if (stateSchemaVersion == SCHEMA_FORMAT_V3 && newSchemaFileWritten) {
         createSchemaFile(newStateSchemaList.sortBy(_.colFamilyName), stateSchemaVersion)
       }
       // TODO: [SPARK-49535] Write Schema files after schema has changed for StateSchemaV3
-      colFamiliesAddedOrRemoved
+      newSchemaFileWritten
     }
   }
 
@@ -255,7 +272,8 @@ object StateSchemaCompatibilityChecker {
       extraOptions: Map[String, String] = Map.empty,
       storeName: String = StateStoreId.DEFAULT_STORE_NAME,
       oldSchemaFilePath: Option[Path] = None,
-      newSchemaFilePath: Option[Path] = None): StateSchemaValidationResult = {
+      newSchemaFilePath: Option[Path] = None,
+      usingAvro: Boolean = false): StateSchemaValidationResult = {
     // SPARK-47776: collation introduces the concept of binary (in)equality, which means
     // in some collation we no longer be able to just compare the binary format of two
     // UnsafeRows to determine equality. For example, 'aaa' and 'AAA' can be "semantically"
@@ -286,7 +304,7 @@ object StateSchemaCompatibilityChecker {
     val result = Try(
       checker.validateAndMaybeEvolveStateSchema(newStateSchema,
         ignoreValueSchema = !storeConf.formatValidationCheckValue,
-        stateSchemaVersion = stateSchemaVersion)
+        stateSchemaVersion = stateSchemaVersion, usingAvro)
     ).toEither.fold(Some(_),
       hasEvolvedSchema => {
         evolvedSchema = hasEvolvedSchema