HeartSaVioR commented on code in PR #48317:
URL: https://github.com/apache/spark/pull/48317#discussion_r1802453142
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala:
##########
@@ -76,6 +86,7 @@ class ListStateImpl[S](
val encodedKey = stateTypesEncoder.encodeGroupingKey()
var isFirst = true
+ var entryCount = getEntryCount(encodedKey)
Review Comment:
Shouldn't this start with 0 as we are resetting the list?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -99,6 +108,7 @@ class ListStateImplWithTTL[S](
val encodedKey = stateTypesEncoder.encodeGroupingKey()
var isFirst = true
+ var entryCount = getEntryCount(encodedKey)
Review Comment:
ditto
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateMetricsImpl.scala:
##########
@@ -0,0 +1,85 @@
+/*
+ * 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.execution.streaming
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow,
UnsafeProjection, UnsafeRow}
+import
org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec,
StateStore}
+import org.apache.spark.sql.types._
+
+/**
+ * Trait that provides helper methods to maintain metrics for a list state.
+ * For list state, we keep track of the count of entries in the list in a
separate column family
+ * to get an accurate view of the number of entries that are updated/removed
from the list and
+ * reported as part of the query progress metrics.
+ */
+trait ListStateMetricsImpl {
+ def stateStore: StateStore
+
+ def baseStateName: String
+
+ def exprEncSchema: StructType
+
+ // We keep track of the count of entries in the list in a separate column
family
+ // to avoid scanning the entire list to get the count.
+ private val counterCFValueSchema: StructType =
+ StructType(Seq(StructField("count", LongType, nullable = false)))
+
+ private val counterCFProjection =
UnsafeProjection.create(counterCFValueSchema)
+
+ private def getRowCounterCFName(stateName: String) = "$rowCounter_" +
stateName
+
+ stateStore.createColFamilyIfAbsent(getRowCounterCFName(baseStateName),
exprEncSchema,
Review Comment:
Is this (V)CF also tracked by state metadata? Do we have tests which perform
restart of the query and see the value to be in sync?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateMetricsImpl.scala:
##########
@@ -0,0 +1,85 @@
+/*
+ * 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.execution.streaming
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow,
UnsafeProjection, UnsafeRow}
+import
org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec,
StateStore}
+import org.apache.spark.sql.types._
+
+/**
+ * Trait that provides helper methods to maintain metrics for a list state.
+ * For list state, we keep track of the count of entries in the list in a
separate column family
+ * to get an accurate view of the number of entries that are updated/removed
from the list and
+ * reported as part of the query progress metrics.
+ */
+trait ListStateMetricsImpl {
+ def stateStore: StateStore
+
+ def baseStateName: String
+
+ def exprEncSchema: StructType
+
+ // We keep track of the count of entries in the list in a separate column
family
+ // to avoid scanning the entire list to get the count.
+ private val counterCFValueSchema: StructType =
+ StructType(Seq(StructField("count", LongType, nullable = false)))
+
+ private val counterCFProjection =
UnsafeProjection.create(counterCFValueSchema)
+
+ private def getRowCounterCFName(stateName: String) = "$rowCounter_" +
stateName
+
+ stateStore.createColFamilyIfAbsent(getRowCounterCFName(baseStateName),
exprEncSchema,
+ counterCFValueSchema, NoPrefixKeyStateEncoderSpec(exprEncSchema),
isInternal = true)
+
+ /**
+ * Function to get the number of entries in the list state for a given
grouping key
+ * @param encodedKey - encoded grouping key
+ * @return - number of entries in the list state
+ */
+ def getEntryCount(encodedKey: UnsafeRow): Long = {
+ val countRow = stateStore.get(encodedKey,
getRowCounterCFName(baseStateName))
+ if (countRow != null) {
+ countRow.getLong(0)
+ } else {
+ 0L
+ }
+ }
+
+ /**
+ * Function to update the number of entries in the list state for a given
grouping key
+ * @param encodedKey - encoded grouping key
+ * @param updatedCount - updated count of entries in the list state
+ */
+ def updateEntryCount(
+ encodedKey: UnsafeRow,
+ updatedCount: Long): Unit = {
+ val updatedCountRow = new GenericInternalRow(1)
Review Comment:
If we assume there is no concurrency for this trait (class), the Unsafe row
instance can be created in prior and be reused - long is a fixed length type in
UnsafeRow, so you can perform in-place update (setLong).
Btw, if we assume there is concurrency, counterCFProjection may be also
dangerous, though we are copying row very soon hence the window to race is
quite small.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MapStateImpl.scala:
##########
@@ -98,6 +112,7 @@ class MapStateImpl[K, V](
StateStoreErrors.requireNonNullStateValue(key, stateName)
val compositeKey = stateTypesEncoder.encodeCompositeKey(key)
store.remove(compositeKey, stateName)
Review Comment:
We probably want to document that the number does not represent the "actual"
state rows to delete?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
