timers

Author: ericm-db

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

@@ -22,12 +22,30 @@ import org.apache.spark.sql.avro.{AvroDeserializer, AvroOptions, AvroSerializer,
 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.types.{NullType, StructField, StructType}
+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)
+
+
+  def convertForRangeScan(schema: StructType): StructType = {
+    StructType(schema.fields.map { field =>
+      if (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
+  }
 }
 
 class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) extends Logging {
@@ -119,7 +137,8 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) extends Lo
   def getTtlStateSchema(
       stateName: String,
       keyEncoder: ExpressionEncoder[Any]): StateStoreColFamilySchema = {
-    val ttlKeySchema = getSingleKeyTTLAvroRowSchema(keyEncoder.schema)
+    val ttlKeySchema = StateStoreColumnFamilySchemaUtils.convertForRangeScan(
+      getSingleKeyTTLRowSchema(keyEncoder.schema))
     val ttlValSchema = StructType(
       Array(StructField("__dummy__", NullType)))
     StateStoreColFamilySchema(
@@ -170,4 +189,21 @@ class StateStoreColumnFamilySchemaUtils(initializeAvroSerde: Boolean) extends Lo
         Some(StructType(keySchema.drop(1)))
       ))
   }
+
+  def getTimerStateSchemaForSecIndex(
+      stateName: String,
+      keySchema: StructType,
+      valSchema: StructType): StateStoreColFamilySchema = {
+    val avroKeySchema = StateStoreColumnFamilySchemaUtils.convertForRangeScan(keySchema)
+    StateStoreColFamilySchema(
+      stateName,
+      keySchema,
+      valSchema,
+      Some(RangeKeyScanStateEncoderSpec(keySchema, Seq(0))),
+      avroEnc = getAvroSerde(
+        StructType(avroKeySchema.take(1)),
+        valSchema,
+        Some(StructType(avroKeySchema.drop(1)))
+      ))
+  }
 }

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

@@ -43,11 +43,6 @@ object TransformWithStateKeyValueRowSchemaUtils {
       .add("expirationMs", LongType)
       .add("groupingKey", keySchema)
 
-  def getSingleKeyTTLAvroRowSchema(keySchema: StructType): StructType =
-    new StructType()
-      .add("expirationMs", BinaryType)
-      .add("groupingKey", keySchema)
-
   def getCompositeKeyTTLRowSchema(
       groupingKeySchema: StructType,
       userKeySchema: StructType): StructType =

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

@@ -165,7 +165,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
@@ -355,10 +361,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)
   }

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,15 +86,15 @@ 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
   private val keySchemaForSecIndex = rowEncoder.keySchemaForSecIndex
   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

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
@@ -394,6 +395,7 @@ class RangeKeyScanStateEncoder(
   extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) with Logging {
 
   import RocksDBStateEncoder._
+  logError(s"### avroEnc.isDefined: ${avroEnc.isDefined}")
 
   private val rangeScanKeyFieldsWithOrdinal: Seq[(StructField, Int)] = {
     orderingOrdinals.map { ordinal =>
@@ -460,18 +462,7 @@ class RangeKeyScanStateEncoder(
     UnsafeProjection.create(refs)
   }
 
-  def convertForRangeScan(schema: StructType): StructType = {
-    StructType(schema.fields.map { field =>
-      if (isFixedSize(field.dataType)) {
-        // Convert numeric types to BinaryType while preserving nullability
-        field.copy(dataType = BinaryType)
-      } else {
-        field
-      }
-    })
-  }
-
-  private val rangeScanAvroSchema = convertForRangeScan(
+  private val rangeScanAvroSchema = StateStoreColumnFamilySchemaUtils.convertForRangeScan(
     StructType(rangeScanKeyFieldsWithOrdinal.map(_._1).toArray))
 
   private val rangeScanAvroType = SchemaConverters.toAvroType(rangeScanAvroSchema)