[GitHub] [spark] zsxwing commented on a change in pull request #34502: [SPARK-37224][SS] Optimize write path on RocksDB state store provider

[GitBox <gi...@apache.org>]

zsxwing commented on a change in pull request #34502:
URL: https://github.com/apache/spark/pull/34502#discussion_r751780630



##########
File path: sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala
##########
@@ -112,8 +112,20 @@ class RocksDB(
         closeDB()
         val metadata = fileManager.loadCheckpointFromDfs(version, workingDir)
         openDB()
-        numKeysOnWritingVersion = metadata.numKeys
-        numKeysOnLoadedVersion = metadata.numKeys
+
+        val numKeys = if (!conf.trackTotalNumberOfRows) {
+          // we don't track the total number of rows - discard the number being track
+          -1L

Review comment:
       Could you point out where we turn `-1` to `0`? I don't find it.

##########
File path: sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala
##########
@@ -196,6 +210,29 @@ class RocksDB(
     }
   }
 
+  private def countKeys(): Long = {
+    // This is being called when opening DB, so doesn't need to deal with writeBatch.
+    val iter = db.newIterator()
+    logInfo(s"Counting keys - getting iterator from version $loadedVersion")
+    iter.seekToFirst()
+
+    // Attempt to close this iterator if there is a task failure, or a task interruption.
+    // This is a hack because it assumes that the RocksDB is running inside a task.
+    Option(TaskContext.get()).foreach { tc =>
+      tc.addTaskCompletionListener[Unit] { _ => iter.close() }
+    }

Review comment:
       We can put `iter.close()` in `finally` instead. This method doesn't return an Iterator to the caller.

##########
File path: docs/structured-streaming-programming-guide.md
##########
@@ -1956,8 +1956,21 @@ Here are the configs regarding to RocksDB instance of the state store provider:
     <td>Whether we resets all ticker and histogram stats for RocksDB on load.</td>
     <td>True</td>
   </tr>
+  <tr>
+    <td>spark.sql.streaming.stateStore.rocksdb.trackTotalNumberOfRows</td>
+    <td>Whether we track the total number of rows in state store. Please refer the details in "Performance-aspect considerations".</td>

Review comment:
       It's better to use a link such as `in [Performance-aspect considerations](#performance-aspect-considerations)` (the link syntax may be wrong. Could you try to build the docs locally to check it?)




-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org