stuff

apache · Dec 23, 2024 · 03c6297 · 03c6297
1 parent 61e99cb
commit 03c6297
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 5 deletions.
diff --git a/...in/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala b/...in/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala
@@ -19,7 +19,7 @@ package org.apache.spark.sql.execution.streaming
 import org.apache.spark.sql.Encoder
 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.{NoPrefixKeyStateEncoderSpec, PrefixKeyScanStateEncoderSpec, RangeKeyScanStateEncoderSpec, StateStoreColFamilySchema}
 import org.apache.spark.sql.types._
 
 object StateStoreColumnFamilySchemaUtils {
@@ -112,4 +112,15 @@ object StateStoreColumnFamilySchemaUtils {
       valSchema,
       Some(PrefixKeyScanStateEncoderSpec(keySchema, 1)))
   }
+
+  def getSecIndexTimerStateSchema(
+      stateName: String,
+      keySchema: StructType,
+      valSchema: StructType): StateStoreColFamilySchema = {
+    StateStoreColFamilySchema(
+      stateName,
+      keySchema,
+      valSchema,
+      Some(RangeKeyScanStateEncoderSpec(keySchema, Seq(0))))
+  }
 }
diff --git a/...src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/...src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala
@@ -376,6 +376,12 @@ class DriverStatefulProcessorHandleImpl(timeMode: TimeMode, keyExprEnc: Expressi
     val colFamilySchema = StateStoreColumnFamilySchemaUtils.
       getTimerStateSchema(stateName, timerEncoder.schemaForKeyRow, timerEncoder.schemaForValueRow)
     columnFamilySchemas.put(stateName, colFamilySchema)
+    val tsToKeyCFName = TimerStateUtils.getTimerStateSecIndexName(timeMode.toString)
+    val secondaryColFamilySchema = StateStoreColumnFamilySchemaUtils.
+      getSecIndexTimerStateSchema(
+        tsToKeyCFName, timerEncoder.keySchemaForSecIndex, timerEncoder.schemaForValueRow)
+    columnFamilySchemas.put(stateName, colFamilySchema)
+    columnFamilySchemas.put(tsToKeyCFName, secondaryColFamilySchema)
     val stateVariableInfo = TransformWithStateVariableUtils.getTimerState(stateName)
     stateVariableInfos.put(stateName, stateVariableInfo)
   }

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
@@ -43,6 +43,15 @@ object TimerStateUtils {
       TimerStateUtils.PROC_TIMERS_STATE_NAME + TimerStateUtils.KEY_TO_TIMESTAMP_CF
     }
   }
+
+  def getTimerStateSecIndexName(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
+    }
+  }
 }
 
 /**

diff --git a/...e/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala b/...e/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
@@ -1766,7 +1766,6 @@ class SingleValueStateEncoder(
     columnFamilyInfo = None,
     supportSchemaEvolution = dataEncoder.supportsSchemaEvolution
   ) with RocksDBValueStateEncoder with Logging {
-  assert(currentSchemaId.isDefined)
 
   override def getCurrentSchemaId: Option[Short] = currentSchemaId
 
@@ -1783,8 +1782,7 @@ class SingleValueStateEncoder(
     }
     // First decode the metadata prefixes
     val prefix = decodeStateRowPrefix(valueBytes)
-    logError(s"### prefix: $prefix")
-    logError(s"### currentSchemaId: ${currentSchemaId}")
+
     val data = decodeStateRowData(valueBytes)
     // Decode the actual value using either Avro or UnsafeRow
     dataEncoder.decodeValue(data, prefix)

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
@@ -668,7 +668,6 @@ class TransformWithStateSuite extends StateStoreMetricsTest
           CheckNewAnswer(("a", "1"), ("c", "1"))
         )
 
-        logError(s"### starting query 2")
         val result2 = inputData.toDS()
           .groupByKey(x => x)
           .transformWithState(new RunningCountStatefulProcessor(),