You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ka...@apache.org on 2026/04/18 09:16:15 UTC
(spark) branch master updated: [SPARK-56369][SS] Introduce a new API "rangeScan/rangeScanWithMultiValues" in StateStore
This is an automated email from the ASF dual-hosted git repository.
HeartSaVioR pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new d7df1920cfb8 [SPARK-56369][SS] Introduce a new API "rangeScan/rangeScanWithMultiValues" in StateStore
d7df1920cfb8 is described below
commit d7df1920cfb8577c17a9b1555592d3772f061417
Author: Jungtaek Lim <ka...@gmail.com>
AuthorDate: Sat Apr 18 18:16:02 2026 +0900
[SPARK-56369][SS] Introduce a new API "rangeScan/rangeScanWithMultiValues" in StateStore
### What changes were proposed in this pull request?
This PR proposes to introduce a new API "rangeScan/rangeScanWithMultiValues" in StateStore. This new API is mostly an optimization focused on RocksDB state store provider.
The new API receives startKey (inclusive) and endKey (exclusive), and provides valid entries in the range. The new API requires the column family to use the state key encoder which can support range scan. This PR adds a new flag on state key encoder to represent the case. At this point, the new API supports three state key encoders - range scan encoder, timestamp prefix/postfix encoder.
Worth noting that the new API supports the combination of "prefix + range (+ remaining)" along with "range + remaining", hence the callers should reason about the ordering for the keys based on the state key encoder they are using for the CF before calling the API. The API implementation may not prevent the incorrect orderliness of the given startKey and endKey (e.g. startKey being later than endKey, some unexpected keys in binary ordering of the keys between startKey and endKey), so [...]
### Why are the changes needed?
This will enable the pattern of range scan to optimize further, effectively skipping tombstones and valid keys in both sides (before the lower bound, after the upper bound) within CF.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
New UTs.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude 4.6 Opus
Closes #55226 from HeartSaVioR/SPARK-56369.
Authored-by: Jungtaek Lim <ka...@gmail.com>
Signed-off-by: Jungtaek Lim <ka...@gmail.com>
---
.../sql/execution/streaming/state/RocksDB.scala | 87 ++++++++
.../streaming/state/RocksDBStateEncoder.scala | 11 +
.../state/RocksDBStateStoreProvider.scala | 62 ++++++
.../sql/execution/streaming/state/StateStore.scala | 59 ++++++
.../RocksDBStateStoreCheckpointFormatV2Suite.scala | 14 ++
.../streaming/state/RocksDBStateStoreSuite.scala | 223 +++++++++++++++++++++
.../RocksDBTimestampEncoderOperationsSuite.scala | 134 +++++++++++++
7 files changed, 590 insertions(+)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala
index 5b4f1d3084ea..89e7a3058f3d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala
@@ -1749,6 +1749,93 @@ class RocksDB(
}
}
+ /**
+ * Scan key-value pairs in the range [startKey, endKey).
+ *
+ * @param startKey None to seek to the beginning of the column family,
+ * or Some(key) to seek to the given start position (inclusive).
+ * @param endKey None to scan to the end of the column family,
+ * or Some(key) as the exclusive upper bound for the scan (encoded key bytes).
+ * @param cfName The column family name.
+ * @return An iterator of ByteArrayPairs in the given range.
+ */
+ def rangeScan(
+ startKey: Option[Array[Byte]],
+ endKey: Option[Array[Byte]],
+ cfName: String = StateStore.DEFAULT_COL_FAMILY_NAME): NextIterator[ByteArrayPair] = {
+ updateMemoryUsageIfNeeded()
+
+ val upperBoundBytes: Option[Array[Byte]] = endKey match {
+ case Some(key) =>
+ Some(if (useColumnFamilies) encodeStateRowWithPrefix(key, cfName) else key)
+ case None =>
+ if (useColumnFamilies) {
+ val cfPrefix = encodeStateRowWithPrefix(Array.emptyByteArray, cfName)
+ RocksDB.prefixUpperBound(cfPrefix)
+ } else {
+ None
+ }
+ }
+
+ val seekTarget = startKey match {
+ case Some(key) =>
+ if (useColumnFamilies) encodeStateRowWithPrefix(key, cfName) else key
+ case None =>
+ if (useColumnFamilies) encodeStateRowWithPrefix(Array.emptyByteArray, cfName)
+ else null
+ }
+
+ val upperBoundSlice = upperBoundBytes.map(new Slice(_))
+ val scanReadOptions = new ReadOptions()
+ upperBoundSlice.foreach(scanReadOptions.setIterateUpperBound)
+
+ val iter = db.newIterator(scanReadOptions)
+ if (seekTarget != null) {
+ iter.seek(seekTarget)
+ } else {
+ iter.seekToFirst()
+ }
+
+ def closeResources(): Unit = {
+ iter.close()
+ scanReadOptions.close()
+ upperBoundSlice.foreach(_.close())
+ }
+
+ Option(TaskContext.get()).foreach { tc =>
+ tc.addTaskCompletionListener[Unit] { _ => closeResources() }
+ }
+
+ new NextIterator[ByteArrayPair] {
+ override protected def getNext(): ByteArrayPair = {
+ if (iter.isValid) {
+ val key = if (useColumnFamilies) {
+ decodeStateRowWithPrefix(iter.key)._1
+ } else {
+ iter.key
+ }
+
+ val value = if (conf.rowChecksumEnabled) {
+ KeyValueChecksumEncoder.decodeAndVerifyValueRowWithChecksum(
+ readVerifier, iter.key, iter.value, delimiterSize)
+ } else {
+ iter.value
+ }
+
+ byteArrayPair.set(key, value)
+ iter.next()
+ byteArrayPair
+ } else {
+ finished = true
+ closeResources()
+ null
+ }
+ }
+
+ override protected def close(): Unit = closeResources()
+ }
+ }
+
def release(): Unit = {}
/**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
index 5acf7cdc9b97..4d9c77348b49 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala
@@ -46,6 +46,7 @@ import org.apache.spark.unsafe.Platform
sealed trait RocksDBKeyStateEncoder {
def supportPrefixKeyScan: Boolean
def supportsDeleteRange: Boolean
+ def supportsRangeScan: Boolean
def encodePrefixKey(prefixKey: UnsafeRow): Array[Byte]
def encodeKey(row: UnsafeRow): Array[Byte]
def decodeKey(keyBytes: Array[Byte]): UnsafeRow
@@ -1500,6 +1501,8 @@ class PrefixKeyScanStateEncoder(
override def supportPrefixKeyScan: Boolean = true
override def supportsDeleteRange: Boolean = false
+
+ override def supportsRangeScan: Boolean = false
}
/**
@@ -1699,6 +1702,8 @@ class RangeKeyScanStateEncoder(
override def supportPrefixKeyScan: Boolean = true
override def supportsDeleteRange: Boolean = true
+
+ override def supportsRangeScan: Boolean = true
}
/**
@@ -1731,6 +1736,8 @@ class NoPrefixKeyStateEncoder(
override def supportsDeleteRange: Boolean = false
+ override def supportsRangeScan: Boolean = false
+
override def encodePrefixKey(prefixKey: UnsafeRow): Array[Byte] = {
throw new IllegalStateException("This encoder doesn't support prefix key!")
}
@@ -1884,6 +1891,8 @@ class TimestampAsPrefixKeyStateEncoder(
// TODO: [SPARK-55491] Revisit this to support delete range if needed.
override def supportsDeleteRange: Boolean = false
+
+ override def supportsRangeScan: Boolean = true
}
/**
@@ -1932,6 +1941,8 @@ class TimestampAsPostfixKeyStateEncoder(
}
override def supportsDeleteRange: Boolean = false
+
+ override def supportsRangeScan: Boolean = true
}
/**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala
index 647b2e27f8c6..45168f407132 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala
@@ -556,6 +556,68 @@ private[sql] class RocksDBStateStoreProvider
new StateStoreIterator(iter, rocksDbIter.closeIfNeeded)
}
+ override def rangeScan(
+ startKey: Option[UnsafeRow],
+ endKey: Option[UnsafeRow],
+ colFamilyName: String): StateStoreIterator[UnsafeRowPair] = {
+ validateAndTransitionState(UPDATE)
+ verifyColFamilyOperations("rangeScan", colFamilyName)
+
+ val kvEncoder = keyValueEncoderMap.get(colFamilyName)
+ require(kvEncoder._1.supportsRangeScan,
+ "Range scan requires an encoder that supports range scanning!")
+
+ val encodedStartKey = startKey.map(kvEncoder._1.encodeKey)
+ val encodedEndKey = endKey.map(kvEncoder._1.encodeKey)
+
+ val rowPair = new UnsafeRowPair()
+ val rocksDbIter = rocksDB.rangeScan(encodedStartKey, encodedEndKey, colFamilyName)
+ val iter = rocksDbIter.map { kv =>
+ rowPair.withRows(kvEncoder._1.decodeKey(kv.key),
+ kvEncoder._2.decodeValue(kv.value))
+ rowPair
+ }
+
+ new StateStoreIterator(iter, rocksDbIter.closeIfNeeded)
+ }
+
+ override def rangeScanWithMultiValues(
+ startKey: Option[UnsafeRow],
+ endKey: Option[UnsafeRow],
+ colFamilyName: String): StateStoreIterator[UnsafeRowPair] = {
+ validateAndTransitionState(UPDATE)
+ verifyColFamilyOperations("rangeScanWithMultiValues", colFamilyName)
+
+ val kvEncoder = keyValueEncoderMap.get(colFamilyName)
+ require(kvEncoder._1.supportsRangeScan,
+ "Range scan with multiple values requires an encoder that supports range scanning!")
+ verify(
+ kvEncoder._2.supportsMultipleValuesPerKey,
+ "Multi-value iterator operation requires an encoder" +
+ " which supports multiple values for a single key")
+
+ val encodedStartKey = startKey.map(kvEncoder._1.encodeKey)
+ val encodedEndKey = endKey.map(kvEncoder._1.encodeKey)
+ val rocksDbIter = rocksDB.rangeScan(encodedStartKey, encodedEndKey, colFamilyName)
+
+ val rowPair = new UnsafeRowPair()
+ val iter = rocksDbIter.flatMap { kv =>
+ val keyRow = kvEncoder._1.decodeKey(kv.key)
+ val valueRows = kvEncoder._2.decodeValues(kv.value)
+ valueRows.iterator.map { valueRow =>
+ rowPair.withRows(keyRow, valueRow)
+ if (!isValidated && rowPair.value != null && !useColumnFamilies) {
+ StateStoreProvider.validateStateRowFormat(
+ rowPair.key, keySchema, rowPair.value, valueSchema, stateStoreId, storeConf)
+ isValidated = true
+ }
+ rowPair
+ }
+ }
+
+ new StateStoreIterator(iter, rocksDbIter.closeIfNeeded)
+ }
+
var checkpointInfo: Option[StateStoreCheckpointInfo] = None
private var storedMetrics: Option[RocksDBMetrics] = None
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
index 6e08c10476ce..e3601f1ef224 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
@@ -183,6 +183,51 @@ trait ReadStateStore {
prefixKey: UnsafeRow,
colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): StateStoreIterator[UnsafeRowPair]
+ /**
+ * Scan key-value pairs in the range [startKey, endKey).
+ *
+ * @param startKey None to scan from the beginning of the column family,
+ * or Some(key) to seek to the given start position (inclusive).
+ * @param endKey None to scan to the end of the column family,
+ * or Some(key) as the exclusive upper bound for the scan.
+ * @param colFamilyName The column family name.
+ *
+ * Callers must ensure the column family's key encoder produces lexicographically ordered
+ * bytes for the scan range to be meaningful (e.g., timestamp-based encoders or
+ * RangeKeyScanStateEncoder).
+ */
+ def rangeScan(
+ startKey: Option[UnsafeRow],
+ endKey: Option[UnsafeRow],
+ colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME)
+ : StateStoreIterator[UnsafeRowPair] = {
+ throw StateStoreErrors.unsupportedOperationException("rangeScan", "")
+ }
+
+ /**
+ * Scan key-value pairs in the range [startKey, endKey), expanding multi-valued entries.
+ *
+ * @param startKey None to scan from the beginning of the column family,
+ * or Some(key) to seek to the given start position (inclusive).
+ * @param endKey None to scan to the end of the column family,
+ * or Some(key) as the exclusive upper bound for the scan.
+ * @param colFamilyName The column family name.
+ *
+ * Callers must ensure the column family's key encoder produces lexicographically ordered
+ * bytes for the scan range to be meaningful (e.g., timestamp-based encoders or
+ * RangeKeyScanStateEncoder).
+ *
+ * It is expected to throw exception if Spark calls this method without setting
+ * multipleValuesPerKey as true for the column family.
+ */
+ def rangeScanWithMultiValues(
+ startKey: Option[UnsafeRow],
+ endKey: Option[UnsafeRow],
+ colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME)
+ : StateStoreIterator[UnsafeRowPair] = {
+ throw StateStoreErrors.unsupportedOperationException("rangeScanWithMultiValues", "")
+ }
+
/** Return an iterator containing all the key-value pairs in the StateStore. */
def iterator(
colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): StateStoreIterator[UnsafeRowPair]
@@ -411,6 +456,20 @@ class WrappedReadStateStore(store: StateStore) extends ReadStateStore {
store.prefixScanWithMultiValues(prefixKey, colFamilyName)
}
+ override def rangeScan(
+ startKey: Option[UnsafeRow],
+ endKey: Option[UnsafeRow],
+ colFamilyName: String): StateStoreIterator[UnsafeRowPair] = {
+ store.rangeScan(startKey, endKey, colFamilyName)
+ }
+
+ override def rangeScanWithMultiValues(
+ startKey: Option[UnsafeRow],
+ endKey: Option[UnsafeRow],
+ colFamilyName: String): StateStoreIterator[UnsafeRowPair] = {
+ store.rangeScanWithMultiValues(startKey, endKey, colFamilyName)
+ }
+
override def iteratorWithMultiValues(
colFamilyName: String): StateStoreIterator[UnsafeRowPair] = {
store.iteratorWithMultiValues(colFamilyName)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala
index 14aa43d3234f..fe09506023cc 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala
@@ -172,6 +172,20 @@ case class CkptIdCollectingStateStoreWrapper(innerStore: StateStore) extends Sta
innerStore.prefixScanWithMultiValues(prefixKey, colFamilyName)
}
+ override def rangeScan(
+ startKey: Option[UnsafeRow],
+ endKey: Option[UnsafeRow],
+ colFamilyName: String): StateStoreIterator[UnsafeRowPair] = {
+ innerStore.rangeScan(startKey, endKey, colFamilyName)
+ }
+
+ override def rangeScanWithMultiValues(
+ startKey: Option[UnsafeRow],
+ endKey: Option[UnsafeRow],
+ colFamilyName: String): StateStoreIterator[UnsafeRowPair] = {
+ innerStore.rangeScanWithMultiValues(startKey, endKey, colFamilyName)
+ }
+
override def iteratorWithMultiValues(
colFamilyName: String): StateStoreIterator[UnsafeRowPair] = {
innerStore.iteratorWithMultiValues(colFamilyName)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala
index ea6a41420ebd..ed8c22740b32 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala
@@ -1633,6 +1633,229 @@ class RocksDBStateStoreSuite extends StateStoreSuiteBase[RocksDBStateStoreProvid
}
}
+ private val diverseTimestamps = Seq(931L, 8000L, 452300L, 4200L, -1L, 90L, 1L, 2L, 8L,
+ -230L, -14569L, -92L, -7434253L, 35L, 6L, 9L, -323L, 5L,
+ -32L, -64L, -256L, 64L, 32L, 1024L, 4096L, 0L)
+
+ testWithColumnFamiliesAndEncodingTypes("rocksdb range scan - rangeScan",
+ TestWithChangelogCheckpointingDisabled) { colFamiliesEnabled =>
+
+ tryWithProviderResource(newStoreProvider(keySchemaWithRangeScan,
+ RangeKeyScanStateEncoderSpec(keySchemaWithRangeScan, Seq(0)),
+ colFamiliesEnabled)) { provider =>
+ val store = provider.getStore(0)
+ try {
+ val cfName = if (colFamiliesEnabled) "testColFamily" else "default"
+ if (colFamiliesEnabled) {
+ store.createColFamilyIfAbsent(cfName,
+ keySchemaWithRangeScan, valueSchema,
+ RangeKeyScanStateEncoderSpec(keySchemaWithRangeScan, Seq(0)))
+ }
+
+ diverseTimestamps.foreach { ts =>
+ store.put(dataToKeyRowWithRangeScan(ts, "a"), dataToValueRow(ts.toInt), cfName)
+ }
+
+ // Bounded positive range [0, 100)
+ val boundedIter = store.rangeScan(
+ Some(dataToKeyRowWithRangeScan(0L, "a")),
+ Some(dataToKeyRowWithRangeScan(100L, "a")), cfName)
+ val boundedResults = boundedIter.map { pair =>
+ (pair.key.getLong(0), pair.value.getInt(0))
+ }.toList
+ boundedIter.close()
+ val expectedBoundedTs = diverseTimestamps.filter(ts => ts >= 0 && ts < 100).sorted
+ assert(boundedResults.map(_._1) === expectedBoundedTs)
+ assert(boundedResults.map(_._2) === expectedBoundedTs.map(_.toInt))
+
+ // Exact bound: startKey is inclusive, endKey is exclusive.
+ // 9 exists in diverseTimestamps, 90 exists in diverseTimestamps.
+ // Scan [9, 90) should include 9 but exclude 90.
+ val exactIter = store.rangeScan(
+ Some(dataToKeyRowWithRangeScan(9L, "a")),
+ Some(dataToKeyRowWithRangeScan(90L, "a")), cfName)
+ val exactResults = exactIter.map(_.key.getLong(0)).toList
+ exactIter.close()
+ assert(exactResults === diverseTimestamps.filter(ts => ts >= 9 && ts < 90).sorted)
+ assert(exactResults.contains(9L))
+ assert(!exactResults.contains(90L))
+
+ // None startKey scans from beginning to 0
+ val noneStartIter = store.rangeScan(
+ None, Some(dataToKeyRowWithRangeScan(0L, "a")), cfName)
+ val noneStartResults = noneStartIter.map(_.key.getLong(0)).toList
+ noneStartIter.close()
+ assert(noneStartResults === diverseTimestamps.filter(_ < 0).sorted)
+
+ // None endKey scans from 1000 to end
+ val noneEndIter = store.rangeScan(
+ Some(dataToKeyRowWithRangeScan(1000L, "a")), None, cfName)
+ val noneEndResults = noneEndIter.map(_.key.getLong(0)).toList
+ noneEndIter.close()
+ assert(noneEndResults === diverseTimestamps.filter(_ >= 1000).sorted)
+
+ // Empty range [10, 31) - no entries between 9 and 32
+ val emptyIter = store.rangeScan(
+ Some(dataToKeyRowWithRangeScan(10L, "a")),
+ Some(dataToKeyRowWithRangeScan(31L, "a")), cfName)
+ assert(!emptyIter.hasNext)
+ emptyIter.close()
+
+ // Bounded negative range [-300, 0)
+ val negIter = store.rangeScan(
+ Some(dataToKeyRowWithRangeScan(-300L, "a")),
+ Some(dataToKeyRowWithRangeScan(0L, "a")), cfName)
+ val negResults = negIter.map(_.key.getLong(0)).toList
+ negIter.close()
+ assert(negResults === diverseTimestamps.filter(ts => ts >= -300 && ts < 0).sorted)
+
+ // Both None: scan entire column family
+ val allIter = store.rangeScan(None, None, cfName)
+ val allResults = allIter.map(_.key.getLong(0)).toList
+ allIter.close()
+ assert(allResults === diverseTimestamps.sorted)
+ } finally {
+ if (!store.hasCommitted) store.abort()
+ }
+ }
+ }
+
+ testWithColumnFamiliesAndEncodingTypes(
+ "rocksdb range scan - scan with multiple key2 values within same key1 range",
+ TestWithChangelogCheckpointingDisabled) { colFamiliesEnabled =>
+
+ tryWithProviderResource(newStoreProvider(keySchemaWithRangeScan,
+ RangeKeyScanStateEncoderSpec(keySchemaWithRangeScan, Seq(0)),
+ colFamiliesEnabled)) { provider =>
+ val store = provider.getStore(0)
+ try {
+ val cfName = if (colFamiliesEnabled) "testColFamily" else "default"
+ if (colFamiliesEnabled) {
+ store.createColFamilyIfAbsent(cfName,
+ keySchemaWithRangeScan, valueSchema,
+ RangeKeyScanStateEncoderSpec(keySchemaWithRangeScan, Seq(0)))
+ }
+
+ Seq("a", "b", "c").foreach { key2 =>
+ Seq(100L, 200L, 300L).foreach { ts =>
+ store.put(dataToKeyRowWithRangeScan(ts, key2), dataToValueRow(ts.toInt), cfName)
+ }
+ }
+
+ val startKey = dataToKeyRowWithRangeScan(100L, "a")
+ val endKey = dataToKeyRowWithRangeScan(201L, "a")
+ val iter = store.rangeScan(Some(startKey), Some(endKey), cfName)
+ val results = iter.map { pair =>
+ (pair.key.getLong(0), pair.key.getUTF8String(1).toString)
+ }.toList
+ iter.close()
+
+ val expectedResults = Seq(
+ (100L, "a"), (100L, "b"), (100L, "c"),
+ (200L, "a"), (200L, "b"), (200L, "c"))
+ assert(results === expectedResults)
+ } finally {
+ if (!store.hasCommitted) store.abort()
+ }
+ }
+ }
+
+ testWithColumnFamiliesAndEncodingTypes(
+ "rocksdb range scan - rangeScanWithMultiValues",
+ TestWithChangelogCheckpointingDisabled) { colFamiliesEnabled =>
+
+ if (colFamiliesEnabled) {
+ tryWithProviderResource(newStoreProvider(
+ StateStoreId(newDir(), Random.nextInt(), 0),
+ RangeKeyScanStateEncoderSpec(keySchemaWithRangeScan, Seq(0)),
+ keySchema = keySchemaWithRangeScan,
+ useColumnFamilies = colFamiliesEnabled,
+ useMultipleValuesPerKey = true)) { provider =>
+ val store = provider.getStore(0)
+ try {
+ val cfName = "testColFamily"
+ store.createColFamilyIfAbsent(cfName,
+ keySchemaWithRangeScan, valueSchema,
+ RangeKeyScanStateEncoderSpec(keySchemaWithRangeScan, Seq(0)),
+ useMultipleValuesPerKey = true)
+
+ diverseTimestamps.zipWithIndex.foreach { case (ts, idx) =>
+ store.putList(dataToKeyRowWithRangeScan(ts, "a"),
+ Array(dataToValueRow(idx * 10), dataToValueRow(idx * 10 + 1)), cfName)
+ }
+
+ // Bounded range [0, 1001)
+ val boundedIter = store.rangeScanWithMultiValues(
+ Some(dataToKeyRowWithRangeScan(0L, "a")),
+ Some(dataToKeyRowWithRangeScan(1001L, "a")), cfName)
+ val boundedResults = boundedIter.map { pair =>
+ (pair.key.getLong(0), pair.value.getInt(0))
+ }.toList
+ boundedIter.close()
+
+ val expectedTimestamps = diverseTimestamps.filter(ts => ts >= 0 && ts <= 1000).sorted
+ assert(boundedResults.map(_._1).distinct === expectedTimestamps)
+ val expectedValues = diverseTimestamps.zipWithIndex
+ .filter { case (ts, _) => ts >= 0 && ts <= 1000 }
+ .sortBy(_._1)
+ .flatMap { case (_, idx) => Seq(idx * 10, idx * 10 + 1) }
+ assert(boundedResults.map(_._2) === expectedValues)
+
+ // Exact bound: startKey is inclusive, endKey is exclusive.
+ // 9 exists in diverseTimestamps, 90 exists in diverseTimestamps.
+ val exactIter = store.rangeScanWithMultiValues(
+ Some(dataToKeyRowWithRangeScan(9L, "a")),
+ Some(dataToKeyRowWithRangeScan(90L, "a")), cfName)
+ val exactResults = exactIter.map(_.key.getLong(0)).toList
+ exactIter.close()
+ val exactResultsDistinct = exactResults.distinct
+ assert(exactResultsDistinct === diverseTimestamps
+ .filter(ts => ts >= 9 && ts < 90).sorted)
+ assert(exactResultsDistinct.contains(9L))
+ assert(!exactResultsDistinct.contains(90L))
+
+ // None startKey scans from beginning to 0
+ val noneStartIter = store.rangeScanWithMultiValues(
+ None, Some(dataToKeyRowWithRangeScan(0L, "a")), cfName)
+ val noneStartResults = noneStartIter.map(_.key.getLong(0)).toList
+ noneStartIter.close()
+ assert(noneStartResults.distinct === diverseTimestamps.filter(_ < 0).sorted)
+
+ // None endKey scans from 1000 to end
+ val noneEndIter = store.rangeScanWithMultiValues(
+ Some(dataToKeyRowWithRangeScan(1000L, "a")), None, cfName)
+ val noneEndResults = noneEndIter.map(_.key.getLong(0)).toList
+ noneEndIter.close()
+ assert(noneEndResults.distinct === diverseTimestamps.filter(_ >= 1000).sorted)
+
+ // Empty range [10, 31) - no entries between 9 and 32
+ val emptyIter = store.rangeScanWithMultiValues(
+ Some(dataToKeyRowWithRangeScan(10L, "a")),
+ Some(dataToKeyRowWithRangeScan(31L, "a")), cfName)
+ assert(!emptyIter.hasNext)
+ emptyIter.close()
+
+ // Bounded negative range [-300, 0)
+ val negIter = store.rangeScanWithMultiValues(
+ Some(dataToKeyRowWithRangeScan(-300L, "a")),
+ Some(dataToKeyRowWithRangeScan(0L, "a")), cfName)
+ val negResults = negIter.map(_.key.getLong(0)).toList
+ negIter.close()
+ assert(negResults.distinct === diverseTimestamps
+ .filter(ts => ts >= -300 && ts < 0).sorted)
+
+ // Both None: scan entire column family
+ val allIter = store.rangeScanWithMultiValues(None, None, cfName)
+ val allResults = allIter.map(_.key.getLong(0)).toList
+ allIter.close()
+ assert(allResults.distinct === diverseTimestamps.sorted)
+ } finally {
+ if (!store.hasCommitted) store.abort()
+ }
+ }
+ }
+ }
+
testWithColumnFamiliesAndEncodingTypes(
"rocksdb key and value schema encoders for column families",
TestWithBothChangelogCheckpointingEnabledAndDisabled) { colFamiliesEnabled =>
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBTimestampEncoderOperationsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBTimestampEncoderOperationsSuite.scala
index a9540a4ad623..6823cf7db406 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBTimestampEncoderOperationsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBTimestampEncoderOperationsSuite.scala
@@ -566,6 +566,140 @@ class RocksDBTimestampEncoderOperationsSuite extends SharedSparkSession
}
}
+ Seq("unsaferow", "avro").foreach { encoding =>
+ test(s"rangeScan with postfix encoder (encoding = $encoding)") {
+ tryWithProviderResource(
+ newStoreProviderWithTimestampEncoder(
+ encoderType = "postfix",
+ useMultipleValuesPerKey = true,
+ dataEncoding = encoding)
+ ) { provider =>
+ val store = provider.getStore(0)
+
+ try {
+ diverseTimestamps.zipWithIndex.foreach { case (ts, idx) =>
+ val keyRow = keyAndTimestampToRow("key1", 1, ts)
+ store.putList(keyRow, Array(valueToRow(idx * 10), valueToRow(idx * 10 + 1)))
+ }
+
+ // key2 entry to verify prefix isolation
+ store.putList(keyAndTimestampToRow("key2", 1, 500L),
+ Array(valueToRow(999)))
+
+ // Bounded range [0, 1001)
+ val boundedIter = store.rangeScanWithMultiValues(
+ Some(keyAndTimestampToRow("key1", 1, 0L)),
+ Some(keyAndTimestampToRow("key1", 1, 1001L)))
+ val boundedResults = boundedIter.map { pair =>
+ (pair.key.getString(0), pair.key.getLong(2), pair.value.getInt(0))
+ }.toList
+ boundedIter.close()
+
+ val expectedTimestamps = diverseTimestamps.filter(ts => ts >= 0 && ts <= 1000).sorted
+ assert(boundedResults.map(_._2).distinct === expectedTimestamps)
+ val expectedValues = diverseTimestamps.zipWithIndex
+ .filter { case (ts, _) => ts >= 0 && ts <= 1000 }
+ .sortBy(_._1)
+ .flatMap { case (_, idx) => Seq(idx * 10, idx * 10 + 1) }
+ assert(boundedResults.map(_._3) === expectedValues)
+ assert(boundedResults.forall(_._1 == "key1"))
+
+ // Exact bound: startKey is inclusive, endKey is exclusive.
+ // 9 exists in diverseTimestamps, 90 exists in diverseTimestamps.
+ val exactIter = store.rangeScanWithMultiValues(
+ Some(keyAndTimestampToRow("key1", 1, 9L)),
+ Some(keyAndTimestampToRow("key1", 1, 90L)))
+ val exactResults = exactIter.map(_.key.getLong(2)).toList
+ exactIter.close()
+ val exactResultsDistinct = exactResults.distinct
+ assert(exactResultsDistinct === diverseTimestamps
+ .filter(ts => ts >= 9 && ts < 90).sorted)
+ assert(exactResultsDistinct.contains(9L))
+ assert(!exactResultsDistinct.contains(90L))
+
+ // Postfix timestamp encoder places the timestamp after the key prefix.
+ // With different key prefixes, None in startKey or endKey would scan across
+ // key boundaries, which is not meaningful for postfix encoding. Hence we only
+ // test bounded ranges with explicit keys here.
+
+ // Full range [MinValue, MaxValue)
+ val fullIter = store.rangeScanWithMultiValues(
+ Some(keyAndTimestampToRow("key1", 1, Long.MinValue)),
+ Some(keyAndTimestampToRow("key1", 1, Long.MaxValue)))
+ val fullResults = fullIter.map(_.key.getLong(2)).toList
+ fullIter.close()
+
+ assert(fullResults.distinct === diverseTimestamps.sorted)
+
+ // Bounded negative range [-300, 0)
+ val negIter = store.rangeScanWithMultiValues(
+ Some(keyAndTimestampToRow("key1", 1, -300L)),
+ Some(keyAndTimestampToRow("key1", 1, 0L)))
+ val negResults = negIter.map(_.key.getLong(2)).toList
+ negIter.close()
+ assert(negResults.distinct === diverseTimestamps
+ .filter(ts => ts >= -300 && ts < 0).sorted)
+
+ // Empty range [10, 31) - no diverseTimestamps entries between 9 and 32
+ val emptyIter = store.rangeScanWithMultiValues(
+ Some(keyAndTimestampToRow("key1", 1, 10L)),
+ Some(keyAndTimestampToRow("key1", 1, 31L)))
+ assert(!emptyIter.hasNext)
+ emptyIter.close()
+ } finally {
+ store.abort()
+ }
+ }
+ }
+
+ // Sanity test for prefix encoder scan. Full scan coverage is in RocksDBStateStoreSuite's
+ // "rocksdb range scan - rangeScan" and "rocksdb range scan - rangeScanWithMultiValues" tests.
+ // This test verifies the timestamp prefix encoder integration works correctly.
+ test(s"rangeScan with prefix encoder (encoding = $encoding)") {
+ tryWithProviderResource(
+ newStoreProviderWithTimestampEncoder(
+ encoderType = "prefix",
+ useMultipleValuesPerKey = true,
+ dataEncoding = encoding)
+ ) { provider =>
+ val store = provider.getStore(0)
+
+ try {
+ val timestamps = Seq(100L, 200L, 300L, 400L, 500L)
+ timestamps.foreach { ts =>
+ store.merge(keyAndTimestampToRow("key1", 1, ts), valueToRow(ts.toInt))
+ }
+ store.merge(keyAndTimestampToRow("key2", 2, 150L), valueToRow(150))
+
+ // None startKey scans from beginning up to 301 (exclusive)
+ val iter1 = store.rangeScanWithMultiValues(None,
+ Some(keyAndTimestampToRow("key1", 1, 301L)))
+ val results1 = iter1.map { pair =>
+ (pair.key.getString(0), pair.key.getLong(2))
+ }.toList
+ iter1.close()
+
+ assert(results1 === Seq(
+ ("key1", 100L), ("key2", 150L), ("key1", 200L), ("key1", 300L)))
+
+ // Boundary safety: endKey at 201, includes everything up to 200
+ // regardless of join key
+ val iter2 = store.rangeScanWithMultiValues(None,
+ Some(keyAndTimestampToRow("key1", 1, 201L)))
+ val results2 = iter2.map { pair =>
+ (pair.key.getString(0), pair.key.getLong(2))
+ }.toList
+ iter2.close()
+
+ assert(results2 === Seq(
+ ("key1", 100L), ("key2", 150L), ("key1", 200L)))
+ } finally {
+ store.abort()
+ }
+ }
+ }
+ }
+
// Helper methods to create test data
private val keyProjection = UnsafeProjection.create(keySchema)
private val keyAndTimestampProjection = UnsafeProjection.create(
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org