metric changes

ericm-db · Apr 12, 2024 · 624e6a5 · 624e6a5
1 parent be2d7a7
commit 624e6a5
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 1 deletion.
diff --git a/...core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala b/...core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala
@@ -311,6 +311,8 @@ case class TransformWithStateExec(
       // metrics around TTL
       StatefulOperatorCustomSumMetric("numValueStateWithTTLVars",
         "Number of value state variables with TTL"),
+      StatefulOperatorCustomSumMetric("numListStateWithTTLVars",
+        "Number of list state variables with TTL"),
       StatefulOperatorCustomSumMetric("numValuesRemovedDueToTTLExpiry",
         "Number of values removed due to TTL expiry")
     )

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
@@ -103,6 +103,8 @@ class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
     new ListStateTTLProcessor(ttlConfig)
   }
 
+  override def getStateTTLMetricName: String = "numListStateWithTTLVars"
+
   test("verify iterator works with expired values in middle of list") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,

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
@@ -46,6 +46,8 @@ abstract class TransformWithStateTTLTest
 
   def getProcessor(ttlConfig: TTLConfig): StatefulProcessor[String, InputEvent, OutputEvent]
 
+  def getStateTTLMetricName: String
+
   test("validate state is evicted at ttl expiry") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName) {
@@ -113,7 +115,7 @@ abstract class TransformWithStateTTLTest
             // for stateful operator
             val progData = q.recentProgress.filter(prog => prog.stateOperators.size > 0)
             assert(progData.filter(prog =>
-              prog.stateOperators(0).customMetrics.get("numValueStateWithTTLVars") > 0).size > 0)
+              prog.stateOperators(0).customMetrics.get(getStateTTLMetricName) > 0).size > 0)
             assert(progData.filter(prog =>
               prog.stateOperators(0).customMetrics
                 .get("numValuesRemovedDueToTTLExpiry") > 0).size > 0)

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
@@ -168,6 +168,7 @@ class TransformWithValueStateTTLSuite extends TransformWithStateTTLTest {
     new ValueStateTTLProcessor(ttlConfig)
   }
 
+  override def getStateTTLMetricName: String = "numValueStateWithTTLVars"
 
   test("validate multiple value states") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->