Re: [PR] [SPARK-52989][SS] Add explicit close() API to State Store iterators [spark]

["anishshri-db (via GitHub)" <gi...@apache.org>]

anishshri-db commented on code in PR #51701:
URL: https://github.com/apache/spark/pull/51701#discussion_r2248771281


##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala:
##########
@@ -261,14 +280,15 @@ class WrappedReadStateStore(store: StateStore) extends ReadStateStore {
     colFamilyName)
 
   override def iterator(colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME):
-    Iterator[UnsafeRowPair] = store.iterator(colFamilyName)
+  StateStoreIterator[UnsafeRowPair] = store.iterator(colFamilyName)
 
   override def abort(): Unit = store.abort()
 
   override def release(): Unit = store.release()
 
   override def prefixScan(prefixKey: UnsafeRow,
-    colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Iterator[UnsafeRowPair] =
+    colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME):
+  StateStoreIterator[UnsafeRowPair] =

Review Comment:
   nit: indent seems off ?



##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala:
##########
@@ -77,12 +77,14 @@ case class CkptIdCollectingStateStoreWrapper(innerStore: StateStore) extends Sta
 
   override def prefixScan(
       prefixKey: UnsafeRow,
-      colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Iterator[UnsafeRowPair] = {
+      colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME):
+  StateStoreIterator[UnsafeRowPair] = {

Review Comment:
   nit: indent seems off ?



##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/statevariables/MapStateImpl.scala:
##########
@@ -56,7 +56,10 @@ class MapStateImpl[K, V](
 
   /** Whether state exists or not. */
   override def exists(): Boolean = {
-    store.prefixScan(stateTypesEncoder.encodeGroupingKey(), stateName).nonEmpty
+    val iter = store.prefixScan(stateTypesEncoder.encodeGroupingKey(), stateName)
+    val nonEmpty = iter.nonEmpty

Review Comment:
   nit: `val result` ?



##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala:
##########
@@ -360,11 +368,16 @@ private[sql] class RocksDBStateStoreProvider
 
       val rowPair = new UnsafeRowPair()
       val prefix = kvEncoder._1.encodePrefixKey(prefixKey)
-      rocksDB.prefixScan(prefix, colFamilyName).map { kv =>
+
+      val rocksDBiter = rocksDB.prefixScan(prefix, colFamilyName)
+
+      val iter = rocksDBiter.map { kv =>
         rowPair.withRows(kvEncoder._1.decodeKey(kv.key),
           kvEncoder._2.decodeValue(kv.value))
         rowPair
       }
+
+      new StateStoreIterator(iter, rocksDBiter.closeIfNeeded)

Review Comment:
   nit: `rocksDbIter`



-- 
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