Second round of self review comments

Author: HeartSaVioR

sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala

@@ -1633,7 +1633,7 @@ class RocksDBStateStoreSuite extends StateStoreSuiteBase[RocksDBStateStoreProvid
     -230L, -14569L, -92L, -7434253L, 35L, 6L, 9L, -323L, 5L,
     -32L, -64L, -256L, 64L, 32L, 1024L, 4096L, 0L)
 
-  testWithColumnFamiliesAndEncodingTypes("rocksdb range scan",
+  testWithColumnFamiliesAndEncodingTypes("rocksdb range scan - scan",
     TestWithBothChangelogCheckpointingEnabledAndDisabled) { colFamiliesEnabled =>
 
     tryWithProviderResource(newStoreProvider(keySchemaWithRangeScan,
@@ -1664,6 +1664,18 @@ class RocksDBStateStoreSuite extends StateStoreSuiteBase[RocksDBStateStoreProvid
         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.scan(
+          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.scan(
           None, Some(dataToKeyRowWithRangeScan(0L, "a")), cfName)
@@ -1728,8 +1740,10 @@ class RocksDBStateStoreSuite extends StateStoreSuiteBase[RocksDBStateStoreProvid
         }.toList
         iter.close()
 
-        assert(results.map(_._1).distinct.sorted === Seq(100L, 200L))
-        assert(results.length === 6) // 3 key2 values x 2 key1 values
+        val expectedResults = Seq(
+          (100L, "a"), (100L, "b"), (100L, "c"),
+          (200L, "a"), (200L, "b"), (200L, "c"))
+        assert(results === expectedResults)
       } finally {
         if (!store.hasCommitted) store.abort()
       }
@@ -1777,13 +1791,48 @@ class RocksDBStateStoreSuite extends StateStoreSuiteBase[RocksDBStateStoreProvid
             .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.scanWithMultiValues(
+            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.scanWithMultiValues(
             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.scanWithMultiValues(
+            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.scanWithMultiValues(
+            Some(dataToKeyRowWithRangeScan(10L, "a")),
+            Some(dataToKeyRowWithRangeScan(31L, "a")), cfName)
+          assert(!emptyIter.hasNext)
+          emptyIter.close()
+
+          // Bounded negative range [-300, 0)
+          val negIter = store.scanWithMultiValues(
+            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)
         } finally {
           if (!store.hasCommitted) store.abort()
         }

sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBTimestampEncoderOperationsSuite.scala

@@ -591,17 +591,36 @@ class RocksDBTimestampEncoderOperationsSuite extends SharedSparkSession
             Some(keyAndTimestampToRow("key1", 1, 0L)),
             Some(keyAndTimestampToRow("key1", 1, 1001L)))
           val boundedResults = boundedIter.map { pair =>
-            (pair.key.getLong(2), pair.value.getInt(0))
+            (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(_._1).distinct === expectedTimestamps)
+          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(_._2) === expectedValues)
+          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.scanWithMultiValues(
+            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.scanWithMultiValues(
@@ -612,6 +631,15 @@ class RocksDBTimestampEncoderOperationsSuite extends SharedSparkSession
 
           assert(fullResults.distinct === diverseTimestamps.sorted)
 
+          // Bounded negative range [-300, 0)
+          val negIter = store.scanWithMultiValues(
+            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.scanWithMultiValues(
             Some(keyAndTimestampToRow("key1", 1, 10L)),
@@ -624,6 +652,9 @@ class RocksDBTimestampEncoderOperationsSuite extends SharedSparkSession
       }
     }
 
+    // Sanity test for prefix encoder scan. Full scan coverage is in RocksDBStateStoreSuite's
+    // "rocksdb range scan - scan" and "rocksdb range scan - scanWithMultiValues" tests.
+    // This test verifies the timestamp prefix encoder integration works correctly.
     test(s"scan with prefix encoder (encoding = $encoding)") {
       tryWithProviderResource(
         newStoreProviderWithTimestampEncoder(
@@ -643,10 +674,13 @@ class RocksDBTimestampEncoderOperationsSuite extends SharedSparkSession
           // None startKey scans from beginning up to 301 (exclusive)
           val iter1 = store.scanWithMultiValues(None,
             Some(keyAndTimestampToRow("key1", 1, 301L)))
-          val results1 = iter1.map(_.key.getLong(2)).toList
+          val results1 = iter1.map { pair =>
+            (pair.key.getString(0), pair.key.getLong(2))
+          }.toList
           iter1.close()
 
-          assert(results1 === Seq(100L, 150L, 200L, 300L))
+          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
@@ -657,54 +691,14 @@ class RocksDBTimestampEncoderOperationsSuite extends SharedSparkSession
           }.toList
           iter2.close()
 
-          assert(results2.map(_._2) === Seq(100L, 150L, 200L))
+          assert(results2 === Seq(
+            ("key1", 100L), ("key2", 150L), ("key1", 200L)))
         } finally {
           store.abort()
         }
       }
     }
 
-    test(s"scan single-value variant (encoding = $encoding)") {
-      tryWithProviderResource(
-        newStoreProviderWithTimestampEncoder(
-          encoderType = "postfix",
-          useMultipleValuesPerKey = false,
-          dataEncoding = encoding)
-      ) { provider =>
-        val store = provider.getStore(0)
-
-        try {
-          diverseTimestamps.foreach { ts =>
-            store.put(keyAndTimestampToRow("key1", 1, ts), valueToRow(ts.toInt))
-          }
-
-          // Bounded positive range [0, 100)
-          val posIter = store.scan(
-            Some(keyAndTimestampToRow("key1", 1, 0L)),
-            Some(keyAndTimestampToRow("key1", 1, 100L)))
-          val posResults = posIter.map { pair =>
-            (pair.key.getLong(2), pair.value.getInt(0))
-          }.toList
-          posIter.close()
-
-          val expectedPosTs = diverseTimestamps.filter(ts => ts >= 0 && ts < 100).sorted
-          assert(posResults.map(_._1) === expectedPosTs)
-          assert(posResults.map(_._2) === expectedPosTs.map(_.toInt))
-
-          // Bounded negative range [-300, 0)
-          val negIter = store.scan(
-            Some(keyAndTimestampToRow("key1", 1, -300L)),
-            Some(keyAndTimestampToRow("key1", 1, 0L)))
-          val negResults = negIter.map(_.key.getLong(2)).toList
-          negIter.close()
-
-          val expectedNegTs = diverseTimestamps.filter(ts => ts >= -300 && ts < 0).sorted
-          assert(negResults === expectedNegTs)
-        } finally {
-          store.abort()
-        }
-      }
-    }
   }
 
   // Helper methods to create test data