zeruibao commented on code in PR #53583:
URL: https://github.com/apache/spark/pull/53583#discussion_r2692257441
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala:
##########
@@ -2401,6 +2401,80 @@ class RocksDBStateStoreSuite extends
StateStoreSuiteBase[RocksDBStateStoreProvid
}
}
+ test("multiGet - batch retrieval of multiple keys") {
+ tryWithProviderResource(newStoreProvider(useColumnFamilies = false)) {
provider =>
+ val store = provider.getStore(0)
+ try {
+ // Put multiple key-value pairs
+ put(store, "a", 1, 10)
+ put(store, "b", 2, 20)
+ put(store, "c", 3, 30)
+ put(store, "d", 4, 40)
+
+ // Create keys array for multiGet
+ val keys = Array(
+ dataToKeyRow("a", 1),
+ dataToKeyRow("b", 2),
+ dataToKeyRow("c", 3),
+ dataToKeyRow("nonexistent", 999) // Key that doesn't exist
+ )
+
+ // Perform multiGet
+ // Note: multiGet returns an iterator that reuses the underlying
UnsafeRow,
+ // so we must copy each row when collecting to an array
+ val results = store.multiGet(keys, StateStore.DEFAULT_COL_FAMILY_NAME)
+ .map(row => if (row != null) row.copy() else null).toArray
+
+ // Verify results
+ assert(results.length === 4)
+ assert(valueRowToData(results(0)) === 10)
+ assert(valueRowToData(results(1)) === 20)
+ assert(valueRowToData(results(2)) === 30)
+ assert(results(3) === null) // Non-existent key should return null
+ } finally {
+ if (!store.hasCommitted) store.abort()
+ }
+ }
+ }
+
+ test("deleteRange - bulk deletion of keys in range") {
+ tryWithProviderResource(
+ newStoreProvider(
+ keySchemaWithRangeScan,
+ RangeKeyScanStateEncoderSpec(keySchemaWithRangeScan, Seq(0)),
Review Comment:
sure, just enforce
--
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]
