[SPARK-54264][SS] DeDup Operator can use RocksDB's keyExists()

### What changes were proposed in this pull request?
Add a keyExists() function in StateStore. In RocksDB state store, it would use RocksDB's keyExists(), while by default, it will check whether get() to return any value. Change deduplication to use keyExists(), rather than get() to see whether the key exists in the state store.

### Why are the changes needed?
Dedup operation now uses RocksDB's Get() only to check if there is result. Rather than doing that, it can use keyExists(), which serves the exact purpose but is faster.

### Does this PR introduce _any_ user-facing change?
No.

### How was this patch tested?
Existing RocksDBStateStoreStreamingDeduplicationSuite should already cover this case.

### Was this patch authored or co-authored using generative AI tooling?
No.

Closes #52960 from siying/dedup_keyexists.

Authored-by: Siying Dong <[email protected]>
Signed-off-by: Anish Shrigondekar <[email protected]>

Author: siying

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala

@@ -1317,8 +1317,8 @@ abstract class BaseStreamingDeduplicateExec
       val result = baseIterator.filter { r =>
         val row = r.asInstanceOf[UnsafeRow]
         val key = getKey(row)
-        val value = store.get(key)
-        if (value == null) {
+        val keyExists = store.keyExists(key)
+        if (!keyExists) {
           putDupInfoIntoState(store, row, key, reusedDupInfoRow)
           numUpdatedStateRows += 1
           numOutputRows += 1

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala

@@ -983,6 +983,26 @@ class RocksDB(
     }
   }
 
+  /**
+   * This method should gives a 100% guarantee of a correct result, whether the key exists or
+   * not.
+   *
+   * @param key The key to check
+   * @param cfName The column family name
+   * @return true if the key exists, false otherwise
+   */
+  def keyExists(
+      key: Array[Byte],
+      cfName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Boolean = {
+    updateMemoryUsageIfNeeded()
+    val keyWithPrefix = if (useColumnFamilies) {
+      encodeStateRowWithPrefix(key, cfName)
+    } else {
+      key
+    }
+    db.keyExists(keyWithPrefix)
+  }
+
   /**
    * Get the values for a given key if present, that were merged (via merge).
    * This returns the values as an iterator of index range, to allow inline access

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala

@@ -251,6 +251,15 @@ private[sql] class RocksDBStateStoreProvider
       value
     }
 
+    override def keyExists(key: UnsafeRow, colFamilyName: String): Boolean = {
+      validateAndTransitionState(UPDATE)
+      verify(key != null, "Key cannot be null")
+      verifyColFamilyOperations("keyExists", colFamilyName)
+
+      val kvEncoder = keyValueEncoderMap.get(colFamilyName)
+      rocksDB.keyExists(kvEncoder._1.encodeKey(key), colFamilyName)
+    }
+
     /**
      * Provides an iterator containing all values of a non-null key.
      *

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

@@ -115,6 +115,23 @@ trait ReadStateStore {
       key: UnsafeRow,
       colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): UnsafeRow
 
+  /**
+   * Check if a key exists in the store, with 100% guarantee of a correct result.
+   *
+   * Default implementation calls get() and checks if the result is null.
+   * Implementations backed by RocksDB should override this to use the native
+   * keyExists() method for better performance.
+   *
+   * @param key The key to check
+   * @param colFamilyName The column family name
+   * @return true if the key exists, false if it doesn't exist
+   */
+  def keyExists(
+      key: UnsafeRow,
+      colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Boolean = {
+    get(key, colFamilyName) != null
+  }
+
   /**
    * Provides an iterator containing all values of a non-null key. If key does not exist,
    * an empty iterator is returned. Implementations should make sure to return an empty
@@ -305,6 +322,12 @@ class WrappedReadStateStore(store: StateStore) extends ReadStateStore {
     colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): UnsafeRow = store.get(key,
     colFamilyName)
 
+  override def keyExists(
+      key: UnsafeRow,
+      colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Boolean = {
+    store.keyExists(key, colFamilyName)
+  }
+
   override def iterator(colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME)
     : StateStoreIterator[UnsafeRowPair] = store.iterator(colFamilyName)
 

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

@@ -70,6 +70,12 @@ case class CkptIdCollectingStateStoreWrapper(innerStore: StateStore) extends Sta
     innerStore.get(key, colFamilyName)
   }
 
+  override def keyExists(
+      key: UnsafeRow,
+      colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Boolean = {
+    innerStore.keyExists(key, colFamilyName)
+  }
+
   override def valuesIterator(
       key: UnsafeRow,
       colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Iterator[UnsafeRow] = {

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

@@ -1393,6 +1393,79 @@ class RocksDBSuite extends AlsoTestWithRocksDBFeatures with SharedSparkSession
     encodeMethod.invoke(db, key.getBytes, cfName).asInstanceOf[Array[Byte]]
   }
 
+  testWithStateStoreCheckpointIdsAndColumnFamilies(
+    "RocksDB: keyExists over 1000 random keys across CFs",
+    TestWithBothChangelogCheckpointingEnabledAndDisabled) {
+    case (enableStateStoreCheckpointIds, colFamiliesEnabled) =>
+      val remoteDir = Utils.createTempDir().toString
+      new File(remoteDir).delete()
+
+      val conf = dbConf.copy(compactOnCommit = false)
+      withDB(
+        remoteDir,
+        conf = conf,
+        useColumnFamilies = colFamiliesEnabled,
+        enableStateStoreCheckpointIds = enableStateStoreCheckpointIds) { db =>
+        val totalPresent = 500
+        val totalAbsent = 500
+
+        // Generate present and absent keys using simple disjoint prefixes
+        val presentKeysAll = (0 until totalPresent).map(i => s"present_$i")
+
+        // Insert present keys
+        db.load(0)
+        // If column families are enabled, create a CF and use it uniformly (after load)
+        val cfNameOpt =
+          if (colFamiliesEnabled) {
+            val cf = "test_cf_random"
+            db.createColFamilyIfAbsent(cf, isInternal = false)
+            Some(cf)
+          } else {
+            None
+          }
+        cfNameOpt match {
+          case Some(cf) =>
+            presentKeysAll.foreach { k => db.put(k, s"v_$k", cf) }
+          case None =>
+            presentKeysAll.foreach { k => db.put(k, s"v_$k") }
+        }
+
+        // Generate absent keys using a different prefix to avoid overlap
+        val absentKeysAll = (0 until totalAbsent).map(i => s"absent_$i")
+
+        // Validation helper to avoid duplication
+        def validate(label: String): Unit = {
+          cfNameOpt match {
+            case Some(cf) =>
+              presentKeysAll.foreach { k =>
+                assert(db.keyExists(k, cf),
+                  s"$label Expected keyExists(true) for present CF key $k")
+              }
+              absentKeysAll.foreach { k =>
+                assert(!db.keyExists(k, cf),
+                  s"$label Expected keyExists(false) for absent CF key $k")
+              }
+            case None =>
+              presentKeysAll.foreach { k =>
+                assert(db.keyExists(k),
+                  s"$label Expected keyExists(true) for present default key $k")
+              }
+              absentKeysAll.foreach { k =>
+                assert(!db.keyExists(k),
+                  s"$label Expected keyExists(false) for absent default key $k")
+              }
+          }
+        }
+
+        // First check before commit
+        validate("(pre-commit)")
+
+        // Commit and re-check
+        db.commit()
+        validate("(post-commit)")
+      }
+  }
+
   testWithStateStoreCheckpointIdsAndColumnFamilies(s"RocksDB: get, put, iterator, commit, load",
     TestWithBothChangelogCheckpointingEnabledAndDisabled) {
     case (enableStateStoreCheckpointIds, colFamiliesEnabled) =>