[hotfix][state] Fix the wrong type cast in ForStMapState

apache · Jan 13, 2025 · d63a7a4 · d63a7a4
1 parent 105720f
commit d63a7a4
Showing 1 changed file with 19 additions and 8 deletions.
diff --git a/...ds/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStMapState.java b/...ds/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStMapState.java
@@ -175,15 +175,26 @@ public UK deserializeUserKey(byte[] userKeyBytes, int userKeyOffset) throws IOEx
     @Override
     @SuppressWarnings("unchecked")
     public ForStDBPutRequest<K, N, UV> buildDBPutRequest(StateRequest<?, ?, ?, ?> stateRequest) {
-        Preconditions.checkArgument(
-                stateRequest.getRequestType() == StateRequestType.MAP_PUT
-                        || stateRequest.getRequestType() == StateRequestType.MAP_REMOVE);
-        ContextKey<K, N> contextKey =
-                new ContextKey<>(
-                        (RecordContext<K>) stateRequest.getRecordContext(),
-                        (N) stateRequest.getNamespace(),
-                        ((Tuple2<UK, UV>) stateRequest.getPayload()).f0);
         Preconditions.checkNotNull(stateRequest.getPayload());
+        ContextKey<K, N> contextKey;
+        if (stateRequest.getRequestType() == StateRequestType.MAP_PUT) {
+            contextKey =
+                    new ContextKey<>(
+                            (RecordContext<K>) stateRequest.getRecordContext(),
+                            (N) stateRequest.getNamespace(),
+                            ((Tuple2<UK, UV>) stateRequest.getPayload()).f0);
+        } else if (stateRequest.getRequestType() == StateRequestType.MAP_REMOVE) {
+            contextKey =
+                    new ContextKey<>(
+                            (RecordContext<K>) stateRequest.getRecordContext(),
+                            (N) stateRequest.getNamespace(),
+                            stateRequest.getPayload());
+        } else {
+            throw new IllegalArgumentException(
+                    "The State type is: "
+                            + stateRequest.getRequestType().name()
+                            + ", which is not a valid put request.");
+        }
         UV value = null;
         if (stateRequest.getRequestType() == StateRequestType.MAP_PUT) {
             value = ((Tuple2<UK, UV>) stateRequest.getPayload()).f1;