move error to StateStoreErrors

eason-yuchen-liu · Jun 26, 2024 · 40b6dc6 · 40b6dc6
1 parent 23639f4
commit 40b6dc6
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 10 deletions.
diff --git a/...main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala b/...main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala
@@ -16,7 +16,6 @@
  */
 package org.apache.spark.sql.execution.datasources.v2.state
 
-import org.apache.spark.SparkException
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeRow}
@@ -99,10 +98,8 @@ class StatePartitionReader(
 
       case Some(snapshotStartBatchId) =>
         if (!provider.isInstanceOf[SupportsFineGrainedReplayFromSnapshot]) {
-          throw new SparkException(
-            errorClass = "STATE_STORE_PROVIDER_DOES_NOT_SUPPORT_FINE_GRAINED_STATE_REPLAY",
-            messageParameters = Map("inputClass" -> provider.getClass.toString),
-            cause = null)
+          StateStoreErrors.stateStoreProviderDoesNotSupportFineGrainedReplay(
+            provider.getClass.toString)
         }
         provider.asInstanceOf[SupportsFineGrainedReplayFromSnapshot]
           .replayReadStateFromSnapshot(

diff --git a/...core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreErrors.scala b/...core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreErrors.scala
@@ -160,6 +160,11 @@ object StateStoreErrors {
     StateStoreSnapshotPartitionNotFound = {
     new StateStoreSnapshotPartitionNotFound(snapshotPartitionId, operatorId, checkpointLocation)
   }
+
+  def stateStoreProviderDoesNotSupportFineGrainedReplay(inputClass: String):
+    StateStoreProviderDoesNotSupportFineGrainedReplay = {
+    new StateStoreProviderDoesNotSupportFineGrainedReplay(inputClass)
+  }
 }
 
 class StateStoreMultipleColumnFamiliesNotSupportedException(stateStoreProvider: String)
@@ -291,3 +296,8 @@ class StateStoreValueRowFormatValidationFailure(errorMsg: String)
   extends SparkRuntimeException(
     errorClass = "STATE_STORE_VALUE_ROW_FORMAT_VALIDATION_FAILURE",
     messageParameters = Map("errorMsg" -> errorMsg))
+
+class StateStoreProviderDoesNotSupportFineGrainedReplay(inputClass: String)
+ extends SparkUnsupportedOperationException(
+   errorClass = "STATE_STORE_PROVIDER_DOES_NOT_SUPPORT_FINE_GRAINED_STATE_REPLAY",
+   messageParameters = Map("inputClass" -> inputClass))
diff --git a/.../scala/org/apache/spark/sql/execution/streaming/state/SymmetricHashJoinStateManager.scala b/.../scala/org/apache/spark/sql/execution/streaming/state/SymmetricHashJoinStateManager.scala
@@ -23,7 +23,6 @@ import scala.annotation.tailrec
 
 import org.apache.hadoop.conf.Configuration
 
-import org.apache.spark.SparkException
 import org.apache.spark.TaskContext
 import org.apache.spark.internal.{Logging, MDC}
 import org.apache.spark.internal.LogKeys.{END_INDEX, START_INDEX, STATE_STORE_ID}
@@ -495,10 +494,8 @@ class SymmetricHashJoinStateManager(
           useMultipleValuesPerKey = false)
         if (snapshotStartVersion.isDefined) {
           if (!stateStoreProvider.isInstanceOf[SupportsFineGrainedReplayFromSnapshot]) {
-            throw new SparkException(
-              errorClass = "STATE_STORE_PROVIDER_DOES_NOT_SUPPORT_FINE_GRAINED_STATE_REPLAY",
-              messageParameters = Map("inputClass" -> stateStoreProvider.getClass.toString),
-              cause = null)
+            StateStoreErrors.stateStoreProviderDoesNotSupportFineGrainedReplay(
+              stateStoreProvider.getClass.toString)
           }
           stateStoreProvider.asInstanceOf[SupportsFineGrainedReplayFromSnapshot]
             .replayStateFromSnapshot(snapshotStartVersion.get, stateInfo.get.storeVersion)