anishshri-db commented on code in PR #50742:
URL: https://github.com/apache/spark/pull/50742#discussion_r2067623925
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreRDD.scala:
##########
@@ -82,19 +107,47 @@ class ReadStateStoreRDD[T: ClassTag, U: ClassTag](
useColumnFamilies: Boolean = false,
extraOptions: Map[String, String] = Map.empty)
extends BaseStateStoreRDD[T, U](dataRDD, checkpointLocation, queryRunId,
operatorId,
- sessionState, storeCoordinator, extraOptions) {
+ sessionState, storeCoordinator, extraOptions) with StateStoreRDDProvider {
+
+ // Using a ConcurrentHashMap to track state stores by partition ID
+ // and whether this store was used to create a write store or not.
+ @transient private lazy val partitionStores =
+ new java.util.concurrent.ConcurrentHashMap[StateStoreProviderId,
(ReadStateStore, Boolean)]()
+
+ override def getStateStoreForPartition(
+ providerId: StateStoreProviderId): Option[ReadStateStore] = {
+ val (readStore, _) = partitionStores.get(providerId)
+ partitionStores.put(providerId, (readStore, true))
+ Option(readStore)
+ }
override protected def getPartitions: Array[Partition] = dataRDD.partitions
override def compute(partition: Partition, ctxt: TaskContext): Iterator[U] =
{
val storeProviderId = getStateProviderId(partition)
+ val partitionId = partition.index
val inputIter = dataRDD.iterator(partition, ctxt)
val store = StateStore.getReadOnly(
storeProviderId, keySchema, valueSchema, keyStateEncoderSpec,
storeVersion,
- stateStoreCkptIds.map(_.apply(partition.index).head),
+ stateStoreCkptIds.map(_.apply(partitionId).head),
Review Comment:
Was this a bug before ?
--
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-unsubscr...@spark.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
