comments

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](

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(

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,

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](

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,17 +184,23 @@ 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,
       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
+    // 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,14 +806,15 @@ class MultiValuedStateEncoder(
 
   import RocksDBStateEncoder._
 
+  private val usingAvroEncoding = avroEnc.isDefined
   // 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 = 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,14 +902,15 @@ class SingleValueStateEncoder(
 
   import RocksDBStateEncoder._
 
+  private val usingAvroEncoding = avroEnc.isDefined
   // 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] = {
-    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 {

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,

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(

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)
+        }
       )
     }
   }