You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by za...@apache.org on 2024/07/25 02:31:29 UTC
(flink) branch master updated (95c617249c7 -> d4294c59e6f)
This is an automated email from the ASF dual-hosted git repository.
zakelly pushed a change to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
from 95c617249c7 [FLINK-35835][test] Make file-merging test tolerate the scenario that source does not produce any record (#25116)
new b07b0b4d994 [FLINK-35858][Runtime/State] Add namespace in internal async states
new 055e11e6b46 [FLINK-35858][Runtime/State] Create state with namespaces
new d4294c59e6f [FLINK-35858][State/ForSt] Support state namespace
The 3 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
.../asyncprocessing/AsyncExecutionController.java | 8 ++
.../runtime/asyncprocessing/RecordContext.java | 18 +++
.../asyncprocessing/StateRequestHandler.java | 8 ++
.../runtime/state/AsyncKeyedStateBackend.java | 11 +-
.../runtime/state/v2/DefaultKeyedStateStoreV2.java | 16 ++-
.../runtime/state/v2/InternalAggregatingState.java | 2 +-
.../flink/runtime/state/v2/InternalKeyedState.java | 7 +-
.../flink/runtime/state/v2/InternalListState.java | 3 +-
.../flink/runtime/state/v2/InternalMapState.java | 2 +-
.../runtime/state/v2/InternalPartitionedState.java | 14 ++-
.../runtime/state/v2/InternalReducingState.java | 2 +-
.../flink/runtime/state/v2/InternalValueState.java | 3 +-
.../AsyncExecutionControllerTest.java | 133 +++++++++++++++++++--
.../flink/runtime/state/StateBackendTestUtils.java | 3 +-
.../state/v2/InternalAggregatingStateTest.java | 2 +-
.../state/v2/InternalKeyedStateTestBase.java | 5 +-
.../runtime/state/v2/InternalListStateTest.java | 3 +-
.../runtime/state/v2/InternalMapStateTest.java | 2 +-
.../state/v2/InternalReducingStateTest.java | 2 +-
.../runtime/state/v2/InternalValueStateTest.java | 3 +-
.../org/apache/flink/state/forst/ContextKey.java | 11 +-
.../flink/state/forst/ForStDBGetRequest.java | 13 +-
.../flink/state/forst/ForStDBPutRequest.java | 17 +--
.../state/forst/ForStGeneralMultiGetOperation.java | 6 +-
.../apache/flink/state/forst/ForStInnerTable.java | 8 +-
.../flink/state/forst/ForStKeyedStateBackend.java | 4 +-
.../flink/state/forst/ForStStateExecutor.java | 4 +-
.../state/forst/ForStStateRequestClassifier.java | 20 ++--
.../apache/flink/state/forst/ForStValueState.java | 24 ++--
.../state/forst/ForStWriteBatchOperation.java | 10 +-
.../state/forst/ForStDBOperationTestBase.java | 14 ++-
.../forst/ForStGeneralMultiGetOperationTest.java | 13 +-
.../flink/state/forst/ForStStateExecutorTest.java | 20 ++--
.../state/forst/ForStWriteBatchOperationTest.java | 17 +--
.../api/operators/StreamOperatorStateHandler.java | 16 +++
.../AbstractAsyncStateStreamOperator.java | 9 ++
.../AbstractAsyncStateStreamOperatorV2.java | 9 ++
.../api/operators/StreamingRuntimeContextTest.java | 6 +-
38 files changed, 357 insertions(+), 111 deletions(-)
copy flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/dataview/StateDataView.java => flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalPartitionedState.java (71%)
(flink) 01/03: [FLINK-35858][Runtime/State] Add namespace in internal async states
Posted by za...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
zakelly pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit b07b0b4d994d961c9f3d5d2def1372cbbbc85cf3
Author: Zakelly <za...@gmail.com>
AuthorDate: Wed Jul 17 17:42:44 2024 +0800
[FLINK-35858][Runtime/State] Add namespace in internal async states
---
.../asyncprocessing/AsyncExecutionController.java | 8 ++
.../runtime/asyncprocessing/RecordContext.java | 18 +++
.../asyncprocessing/StateRequestHandler.java | 8 ++
.../runtime/state/v2/InternalAggregatingState.java | 2 +-
.../flink/runtime/state/v2/InternalKeyedState.java | 7 +-
.../flink/runtime/state/v2/InternalListState.java | 3 +-
.../flink/runtime/state/v2/InternalMapState.java | 2 +-
.../v2/InternalPartitionedState.java} | 20 +---
.../runtime/state/v2/InternalReducingState.java | 2 +-
.../flink/runtime/state/v2/InternalValueState.java | 3 +-
.../AsyncExecutionControllerTest.java | 128 +++++++++++++++++++--
.../state/v2/InternalAggregatingStateTest.java | 2 +-
.../runtime/state/v2/InternalListStateTest.java | 3 +-
.../runtime/state/v2/InternalMapStateTest.java | 2 +-
.../state/v2/InternalReducingStateTest.java | 2 +-
.../runtime/state/v2/InternalValueStateTest.java | 3 +-
.../state/forst/ForStStateRequestClassifier.java | 12 +-
.../apache/flink/state/forst/ForStValueState.java | 2 +-
.../state/forst/ForStDBOperationTestBase.java | 10 +-
.../forst/ForStGeneralMultiGetOperationTest.java | 9 +-
.../flink/state/forst/ForStStateExecutorTest.java | 20 ++--
.../state/forst/ForStWriteBatchOperationTest.java | 9 +-
22 files changed, 219 insertions(+), 56 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionController.java b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionController.java
index d06888ab6c4..d48208e452f 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionController.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionController.java
@@ -25,12 +25,14 @@ import org.apache.flink.core.state.InternalStateFuture;
import org.apache.flink.core.state.StateFutureImpl.AsyncFrameworkExceptionHandler;
import org.apache.flink.runtime.asyncprocessing.EpochManager.ParallelMode;
import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
+import org.apache.flink.runtime.state.v2.InternalPartitionedState;
import org.apache.flink.util.Preconditions;
import org.apache.flink.util.function.ThrowingRunnable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.util.Optional;
@@ -255,6 +257,12 @@ public class AsyncExecutionController<K> implements StateRequestHandler {
return stateFuture;
}
+ @Override
+ public <N> void setCurrentNamespaceForState(
+ @Nonnull InternalPartitionedState<N> state, N namespace) {
+ currentContext.setNamespace(state, namespace);
+ }
+
<IN, OUT> void insertActiveBuffer(StateRequest<K, IN, OUT> request) {
stateRequestsBuffer.enqueueToActive(request);
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/RecordContext.java b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/RecordContext.java
index 86352ce1cf9..b9f48125c29 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/RecordContext.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/RecordContext.java
@@ -19,9 +19,12 @@
package org.apache.flink.runtime.asyncprocessing;
import org.apache.flink.runtime.asyncprocessing.EpochManager.Epoch;
+import org.apache.flink.runtime.state.v2.InternalPartitionedState;
import javax.annotation.Nullable;
+import java.util.HashMap;
+import java.util.Map;
import java.util.Objects;
import java.util.function.Consumer;
@@ -56,6 +59,9 @@ public class RecordContext<K> extends ReferenceCounted<RecordContext.DisposerRun
/** The keyGroup to which key belongs. */
private final int keyGroup;
+ /** The namespaces of states. Lazy initialization for saving memory. */
+ private Map<InternalPartitionedState<?>, Object> namespaces = null;
+
/**
* The extra context info which is used to hold customized data defined by state backend. The
* state backend can use this field to cache some data that can be used multiple times in
@@ -111,6 +117,18 @@ public class RecordContext<K> extends ReferenceCounted<RecordContext.DisposerRun
return keyGroup;
}
+ @SuppressWarnings("unchecked")
+ public <N> N getNamespace(InternalPartitionedState<N> state) {
+ return namespaces == null ? null : (N) namespaces.get(state);
+ }
+
+ public <N> void setNamespace(InternalPartitionedState<N> state, N namespace) {
+ if (namespaces == null) {
+ namespaces = new HashMap<>();
+ }
+ namespaces.put(state, namespace);
+ }
+
public void setExtra(Object extra) {
this.extra = extra;
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequestHandler.java b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequestHandler.java
index 1288ad8a01b..fc81f87f9fb 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequestHandler.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequestHandler.java
@@ -21,7 +21,9 @@ package org.apache.flink.runtime.asyncprocessing;
import org.apache.flink.annotation.Internal;
import org.apache.flink.api.common.state.v2.State;
import org.apache.flink.core.state.InternalStateFuture;
+import org.apache.flink.runtime.state.v2.InternalPartitionedState;
+import javax.annotation.Nonnull;
import javax.annotation.Nullable;
/** The handler which can process {@link StateRequest}. */
@@ -39,4 +41,10 @@ public interface StateRequestHandler {
*/
<IN, OUT> InternalStateFuture<OUT> handleRequest(
@Nullable State state, StateRequestType type, @Nullable IN payload);
+
+ /**
+ * Set current namespace for a state. See {@link
+ * InternalPartitionedState#setCurrentNamespace(Object)}.
+ */
+ <N> void setCurrentNamespaceForState(@Nonnull InternalPartitionedState<N> state, N namespace);
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalAggregatingState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalAggregatingState.java
index 7f05cdbd4d2..971af6f7f7d 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalAggregatingState.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalAggregatingState.java
@@ -33,7 +33,7 @@ import org.apache.flink.runtime.asyncprocessing.StateRequestType;
* @param <ACC> TThe type of the accumulator (intermediate aggregation state).
* @param <OUT> The type of the values that are returned from the state.
*/
-public class InternalAggregatingState<K, IN, ACC, OUT> extends InternalKeyedState<K, ACC>
+public class InternalAggregatingState<K, N, IN, ACC, OUT> extends InternalKeyedState<K, N, ACC>
implements AggregatingState<IN, OUT> {
protected final AggregateFunction<IN, ACC, OUT> aggregateFunction;
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalKeyedState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalKeyedState.java
index 2b1be2a6764..408ca80f6c8 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalKeyedState.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalKeyedState.java
@@ -37,7 +37,7 @@ import org.apache.flink.runtime.asyncprocessing.StateRequestType;
* @param <V> The type of values kept internally in state.
*/
@Internal
-public abstract class InternalKeyedState<K, V> implements State {
+public abstract class InternalKeyedState<K, N, V> implements InternalPartitionedState<N> {
private final StateRequestHandler stateRequestHandler;
@@ -64,6 +64,11 @@ public abstract class InternalKeyedState<K, V> implements State {
return stateRequestHandler.handleRequest(this, stateRequestType, payload);
}
+ @Override
+ public void setCurrentNamespace(N namespace) {
+ stateRequestHandler.setCurrentNamespaceForState(this, namespace);
+ }
+
@Override
public final StateFuture<Void> asyncClear() {
return handleRequest(StateRequestType.CLEAR, null);
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalListState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalListState.java
index 7bcc652c2b5..ac1c43005c2 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalListState.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalListState.java
@@ -32,7 +32,8 @@ import java.util.List;
* @param <K> The type of key the state is associated to.
* @param <V> The type of values kept internally in state.
*/
-public class InternalListState<K, V> extends InternalKeyedState<K, V> implements ListState<V> {
+public class InternalListState<K, N, V> extends InternalKeyedState<K, N, V>
+ implements ListState<V> {
public InternalListState(
StateRequestHandler stateRequestHandler, ListStateDescriptor<V> stateDescriptor) {
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalMapState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalMapState.java
index b718a3eb687..d028f9d611c 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalMapState.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalMapState.java
@@ -34,7 +34,7 @@ import java.util.Map;
* @param <UK> The type of user key of this state.
* @param <V> The type of values kept internally in state.
*/
-public class InternalMapState<K, UK, V> extends InternalKeyedState<K, V>
+public class InternalMapState<K, N, UK, V> extends InternalKeyedState<K, N, V>
implements MapState<UK, V> {
public InternalMapState(
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequestHandler.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalPartitionedState.java
similarity index 57%
copy from flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequestHandler.java
copy to flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalPartitionedState.java
index 1288ad8a01b..46e35b89838 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequestHandler.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalPartitionedState.java
@@ -16,27 +16,19 @@
* limitations under the License.
*/
-package org.apache.flink.runtime.asyncprocessing;
+package org.apache.flink.runtime.state.v2;
import org.apache.flink.annotation.Internal;
import org.apache.flink.api.common.state.v2.State;
-import org.apache.flink.core.state.InternalStateFuture;
-import javax.annotation.Nullable;
-
-/** The handler which can process {@link StateRequest}. */
+/** A state that is partitioned into different namespaces. */
@Internal
-public interface StateRequestHandler {
+public interface InternalPartitionedState<N> extends State {
/**
- * Submit a {@link StateRequest} to this StateRequestHandler.
+ * Set current namespace and access state under specified namespace afterward.
*
- * @param state the state to request. Could be {@code null} if the type is {@link
- * StateRequestType#SYNC_POINT}.
- * @param type the type of this request.
- * @param payload the payload input for this request.
- * @return the state future.
+ * @param namespace the specified namespace
*/
- <IN, OUT> InternalStateFuture<OUT> handleRequest(
- @Nullable State state, StateRequestType type, @Nullable IN payload);
+ void setCurrentNamespace(N namespace);
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalReducingState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalReducingState.java
index eed6eae0dec..64d538d6bb1 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalReducingState.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalReducingState.java
@@ -30,7 +30,7 @@ import org.apache.flink.runtime.asyncprocessing.StateRequestType;
* @param <K> The type of key the state is associated to.
* @param <V> The type of values kept internally in state.
*/
-public class InternalReducingState<K, V> extends InternalKeyedState<K, V>
+public class InternalReducingState<K, N, V> extends InternalKeyedState<K, N, V>
implements ReducingState<V> {
protected final ReduceFunction<V> reduceFunction;
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalValueState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalValueState.java
index 2a9c643d17a..8c87685a62b 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalValueState.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/InternalValueState.java
@@ -30,7 +30,8 @@ import org.apache.flink.runtime.asyncprocessing.StateRequestType;
* @param <K> The type of key the state is associated to.
* @param <V> The type of values kept internally in state.
*/
-public class InternalValueState<K, V> extends InternalKeyedState<K, V> implements ValueState<V> {
+public class InternalValueState<K, N, V> extends InternalKeyedState<K, N, V>
+ implements ValueState<V> {
public InternalValueState(
StateRequestHandler stateRequestHandler, ValueStateDescriptor<V> valueStateDescriptor) {
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionControllerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionControllerTest.java
index e98a66b82a0..53337d71740 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionControllerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionControllerTest.java
@@ -22,6 +22,7 @@ import org.apache.flink.api.common.operators.MailboxExecutor;
import org.apache.flink.api.common.state.v2.State;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
import org.apache.flink.api.common.typeutils.base.IntSerializer;
+import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.core.fs.CloseableRegistry;
import org.apache.flink.core.state.StateFutureImpl.AsyncFrameworkExceptionHandler;
import org.apache.flink.core.state.StateFutureUtils;
@@ -45,6 +46,7 @@ import java.util.HashMap;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
import java.util.function.Supplier;
import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
@@ -232,6 +234,115 @@ class AsyncExecutionControllerTest {
resourceRegistry.close();
}
+ @Test
+ void testNamespace() throws IOException {
+ final Consumer<String> userCode =
+ (r) -> {
+ valueState.setCurrentNamespace(r);
+ valueState
+ .asyncValue()
+ .thenCompose(
+ val -> {
+ int updated = (val == null ? 1 : (val + 1));
+ return valueState
+ .asyncUpdate(updated)
+ .thenCompose(
+ o ->
+ StateFutureUtils.completedFuture(
+ updated));
+ })
+ .thenAccept(val -> output.set(val));
+ };
+ CloseableRegistry resourceRegistry = new CloseableRegistry();
+ setup(
+ 100,
+ 10000L,
+ 1000,
+ new SyncMailboxExecutor(),
+ new TestAsyncFrameworkExceptionHandler(),
+ resourceRegistry);
+ // ============================ element1 ============================
+ String record1 = "key1-r1";
+ String key1 = "key1";
+ // Simulate the wrapping in {@link RecordProcessorUtils#getRecordProcessor()}, wrapping the
+ // record and key with RecordContext.
+ RecordContext<String> recordContext1 = aec.buildContext(record1, key1);
+ aec.setCurrentContext(recordContext1);
+ // execute user code
+ userCode.accept(record1);
+
+ // Single-step run.
+ // Firstly, the user code generates value get in active buffer.
+ assertThat(aec.stateRequestsBuffer.activeQueueSize()).isEqualTo(1);
+ assertThat(aec.keyAccountingUnit.occupiedCount()).isEqualTo(1);
+ assertThat(aec.inFlightRecordNum.get()).isEqualTo(1);
+ aec.triggerIfNeeded(true);
+ // After running, the value update is in active buffer.
+ assertThat(aec.stateRequestsBuffer.activeQueueSize()).isEqualTo(1);
+ assertThat(aec.keyAccountingUnit.occupiedCount()).isEqualTo(1);
+ aec.triggerIfNeeded(true);
+ // Value update finishes.
+ assertThat(aec.stateRequestsBuffer.activeQueueSize()).isEqualTo(0);
+ assertThat(aec.keyAccountingUnit.occupiedCount()).isEqualTo(0);
+ assertThat(output.get()).isEqualTo(1);
+ assertThat(recordContext1.getReferenceCount()).isEqualTo(0);
+ assertThat(aec.inFlightRecordNum.get()).isEqualTo(0);
+
+ // ============================ element 2 & 3(1) ============================
+ String record2 = "key1-r2";
+ String key2 = "key1";
+ RecordContext<String> recordContext2 = aec.buildContext(record2, key2);
+ aec.setCurrentContext(recordContext2);
+ // execute user code
+ userCode.accept(record2);
+
+ String record3 = "key1-r1";
+ String key3 = "key1";
+ RecordContext<String> recordContext3 = aec.buildContext(record3, key3);
+ aec.setCurrentContext(recordContext3);
+ // execute user code
+ userCode.accept(record3);
+
+ // Single-step run.
+ // Firstly, the user code for record2 generates value get in active buffer,
+ // while user code for record3 generates value get in blocking buffer.
+ assertThat(aec.stateRequestsBuffer.activeQueueSize()).isEqualTo(1);
+ assertThat(aec.stateRequestsBuffer.blockingQueueSize()).isEqualTo(1);
+ assertThat(aec.keyAccountingUnit.occupiedCount()).isEqualTo(1);
+ assertThat(aec.inFlightRecordNum.get()).isEqualTo(2);
+ aec.triggerIfNeeded(true);
+ // After running, the value update for record2 is in active buffer.
+ assertThat(aec.stateRequestsBuffer.activeQueueSize()).isEqualTo(1);
+ assertThat(aec.stateRequestsBuffer.blockingQueueSize()).isEqualTo(1);
+ assertThat(aec.keyAccountingUnit.occupiedCount()).isEqualTo(1);
+ assertThat(aec.inFlightRecordNum.get()).isEqualTo(2);
+ aec.triggerIfNeeded(true);
+ // Value update for record2 finishes. The value get for record3 is migrated from blocking
+ // buffer to active buffer actively.
+ assertThat(aec.stateRequestsBuffer.activeQueueSize()).isEqualTo(1);
+ assertThat(aec.keyAccountingUnit.occupiedCount()).isEqualTo(1);
+ assertThat(aec.inFlightRecordNum.get()).isEqualTo(1);
+ assertThat(output.get()).isEqualTo(1);
+ assertThat(recordContext2.getReferenceCount()).isEqualTo(0);
+ assertThat(aec.stateRequestsBuffer.blockingQueueSize()).isEqualTo(0);
+
+ // Let value get for record3 to run.
+ aec.triggerIfNeeded(true);
+ // After running, the value update for record3 is in active buffer.
+ assertThat(aec.stateRequestsBuffer.activeQueueSize()).isEqualTo(1);
+ assertThat(aec.keyAccountingUnit.occupiedCount()).isEqualTo(1);
+ assertThat(aec.inFlightRecordNum.get()).isEqualTo(1);
+ aec.triggerIfNeeded(true);
+ // Value update for record3 finishes.
+ assertThat(aec.stateRequestsBuffer.activeQueueSize()).isEqualTo(0);
+ assertThat(aec.keyAccountingUnit.occupiedCount()).isEqualTo(0);
+ assertThat(aec.inFlightRecordNum.get()).isEqualTo(0);
+ assertThat(output.get()).isEqualTo(2);
+ assertThat(recordContext3.getReferenceCount()).isEqualTo(0);
+
+ resourceRegistry.close();
+ }
+
@Test
void testRecordsRunInOrder() throws IOException {
CloseableRegistry resourceRegistry = new CloseableRegistry();
@@ -727,22 +838,22 @@ class AsyncExecutionControllerTest {
/** Simulate the underlying state that is actually used to execute the request. */
static class TestUnderlyingState {
- private final HashMap<String, Integer> hashMap;
+ private final HashMap<Tuple2<String, String>, Integer> hashMap;
public TestUnderlyingState() {
this.hashMap = new HashMap<>();
}
- public Integer get(String key) {
- return hashMap.get(key);
+ public Integer get(String key, String namespace) {
+ return hashMap.get(Tuple2.of(key, namespace));
}
- public void update(String key, Integer val) {
- hashMap.put(key, val);
+ public void update(String key, String namespace, Integer val) {
+ hashMap.put(Tuple2.of(key, namespace), val);
}
}
- static class TestValueState extends InternalValueState<String, Integer> {
+ static class TestValueState extends InternalValueState<String, String, Integer> {
private final TestUnderlyingState underlyingState;
@@ -776,7 +887,9 @@ class AsyncExecutionControllerTest {
Preconditions.checkState(request.getState() != null);
TestValueState state = (TestValueState) request.getState();
Integer val =
- state.underlyingState.get((String) request.getRecordContext().getKey());
+ state.underlyingState.get(
+ (String) request.getRecordContext().getKey(),
+ (String) request.getRecordContext().getNamespace(state));
request.getFuture().complete(val);
} else if (request.getRequestType() == StateRequestType.VALUE_UPDATE) {
Preconditions.checkState(request.getState() != null);
@@ -784,6 +897,7 @@ class AsyncExecutionControllerTest {
state.underlyingState.update(
(String) request.getRecordContext().getKey(),
+ (String) request.getRecordContext().getNamespace(state),
(Integer) request.getPayload());
request.getFuture().complete(null);
} else {
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalAggregatingStateTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalAggregatingStateTest.java
index 6e9724a6155..97410748189 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalAggregatingStateTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalAggregatingStateTest.java
@@ -55,7 +55,7 @@ class InternalAggregatingStateTest extends InternalKeyedStateTestBase {
AggregatingStateDescriptor<Integer, Integer, Integer> descriptor =
new AggregatingStateDescriptor<>(
"testAggState", aggregator, BasicTypeInfo.INT_TYPE_INFO);
- InternalAggregatingState<String, Integer, Integer, Integer> state =
+ InternalAggregatingState<String, Void, Integer, Integer, Integer> state =
new InternalAggregatingState<>(aec, descriptor);
aec.setCurrentContext(aec.buildContext("test", "test"));
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalListStateTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalListStateTest.java
index 941092fcdb3..ccb4c977233 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalListStateTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalListStateTest.java
@@ -34,7 +34,8 @@ public class InternalListStateTest extends InternalKeyedStateTestBase {
public void testEachOperation() {
ListStateDescriptor<Integer> descriptor =
new ListStateDescriptor<>("testState", BasicTypeInfo.INT_TYPE_INFO);
- InternalListState<String, Integer> listState = new InternalListState<>(aec, descriptor);
+ InternalListState<String, Void, Integer> listState =
+ new InternalListState<>(aec, descriptor);
aec.setCurrentContext(aec.buildContext("test", "test"));
listState.asyncClear();
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalMapStateTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalMapStateTest.java
index 3eb7c652cce..3cfd214c0ae 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalMapStateTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalMapStateTest.java
@@ -36,7 +36,7 @@ public class InternalMapStateTest extends InternalKeyedStateTestBase {
MapStateDescriptor<String, Integer> descriptor =
new MapStateDescriptor<>(
"testState", BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.INT_TYPE_INFO);
- InternalMapState<String, String, Integer> mapState =
+ InternalMapState<String, Void, String, Integer> mapState =
new InternalMapState<>(aec, descriptor);
aec.setCurrentContext(aec.buildContext("test", "test"));
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalReducingStateTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalReducingStateTest.java
index 80ccc1ede9a..6340952aad3 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalReducingStateTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalReducingStateTest.java
@@ -33,7 +33,7 @@ public class InternalReducingStateTest extends InternalKeyedStateTestBase {
ReduceFunction<Integer> reducer = Integer::sum;
ReducingStateDescriptor<Integer> descriptor =
new ReducingStateDescriptor<>("testState", reducer, BasicTypeInfo.INT_TYPE_INFO);
- InternalReducingState<String, Integer> reducingState =
+ InternalReducingState<String, Void, Integer> reducingState =
new InternalReducingState<>(aec, descriptor);
aec.setCurrentContext(aec.buildContext("test", "test"));
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalValueStateTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalValueStateTest.java
index 801f264573e..d4754ce7f55 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalValueStateTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalValueStateTest.java
@@ -31,7 +31,8 @@ public class InternalValueStateTest extends InternalKeyedStateTestBase {
public void testEachOperation() {
ValueStateDescriptor<Integer> descriptor =
new ValueStateDescriptor<>("testState", BasicTypeInfo.INT_TYPE_INFO);
- InternalValueState<String, Integer> valueState = new InternalValueState<>(aec, descriptor);
+ InternalValueState<String, Void, Integer> valueState =
+ new InternalValueState<>(aec, descriptor);
aec.setCurrentContext(aec.buildContext("test", "test"));
valueState.asyncClear();
diff --git a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStStateRequestClassifier.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStStateRequestClassifier.java
index f32b3b665af..42324eccb67 100644
--- a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStStateRequestClassifier.java
+++ b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStStateRequestClassifier.java
@@ -56,23 +56,23 @@ public class ForStStateRequestClassifier implements StateRequestContainer {
switch (stateRequestType) {
case VALUE_GET:
{
- ForStValueState<?, ?> forStValueState =
- (ForStValueState<?, ?>) stateRequest.getState();
+ ForStValueState<?, ?, ?> forStValueState =
+ (ForStValueState<?, ?, ?>) stateRequest.getState();
dbGetRequests.add(forStValueState.buildDBGetRequest(stateRequest));
return;
}
case VALUE_UPDATE:
{
- ForStValueState<?, ?> forStValueState =
- (ForStValueState<?, ?>) stateRequest.getState();
+ ForStValueState<?, ?, ?> forStValueState =
+ (ForStValueState<?, ?, ?>) stateRequest.getState();
dbPutRequests.add(forStValueState.buildDBPutRequest(stateRequest));
return;
}
case CLEAR:
{
if (stateRequest.getState() instanceof ForStValueState) {
- ForStValueState<?, ?> forStValueState =
- (ForStValueState<?, ?>) stateRequest.getState();
+ ForStValueState<?, ?, ?> forStValueState =
+ (ForStValueState<?, ?, ?>) stateRequest.getState();
dbPutRequests.add(forStValueState.buildDBPutRequest(stateRequest));
return;
} else {
diff --git a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStValueState.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStValueState.java
index 52f2e6d848c..153cdbc0fa1 100644
--- a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStValueState.java
+++ b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStValueState.java
@@ -42,7 +42,7 @@ import java.util.function.Supplier;
* @param <K> The type of the key.
* @param <V> The type of the value.
*/
-public class ForStValueState<K, V> extends InternalValueState<K, V>
+public class ForStValueState<K, N, V> extends InternalValueState<K, N, V>
implements ValueState<V>, ForStInnerTable<ContextKey<K>, V> {
/** The column family which this internal value state belongs to. */
diff --git a/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStDBOperationTestBase.java b/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStDBOperationTestBase.java
index 73b914ebe52..fa09bfc46d7 100644
--- a/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStDBOperationTestBase.java
+++ b/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStDBOperationTestBase.java
@@ -32,6 +32,7 @@ import org.apache.flink.runtime.asyncprocessing.StateRequestHandler;
import org.apache.flink.runtime.asyncprocessing.StateRequestType;
import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
import org.apache.flink.runtime.state.SerializedCompositeKeyBuilder;
+import org.apache.flink.runtime.state.v2.InternalPartitionedState;
import org.apache.flink.runtime.state.v2.ValueStateDescriptor;
import org.apache.flink.util.function.BiFunctionWithException;
import org.apache.flink.util.function.FunctionWithException;
@@ -45,6 +46,7 @@ import org.rocksdb.ColumnFamilyHandle;
import org.rocksdb.ColumnFamilyOptions;
import org.rocksdb.RocksDB;
+import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.nio.file.Path;
@@ -84,6 +86,12 @@ public class ForStDBOperationTestBase {
@Nullable State state, StateRequestType type, @Nullable IN payload) {
throw new UnsupportedOperationException();
}
+
+ @Override
+ public <N> void setCurrentNamespaceForState(
+ @Nonnull InternalPartitionedState<N> state, N namespace) {
+ throw new UnsupportedOperationException();
+ }
};
}
@@ -94,7 +102,7 @@ public class ForStDBOperationTestBase {
return new ContextKey<>(recordContext);
}
- protected ForStValueState<Integer, String> buildForStValueState(String stateName)
+ protected ForStValueState<Integer, Void, String> buildForStValueState(String stateName)
throws Exception {
ColumnFamilyHandle cf = createColumnFamilyHandle(stateName);
ValueStateDescriptor<String> valueStateDescriptor =
diff --git a/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStGeneralMultiGetOperationTest.java b/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStGeneralMultiGetOperationTest.java
index 5b8263ec25e..2b24a65d650 100644
--- a/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStGeneralMultiGetOperationTest.java
+++ b/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStGeneralMultiGetOperationTest.java
@@ -34,15 +34,18 @@ public class ForStGeneralMultiGetOperationTest extends ForStDBOperationTestBase
@Test
public void testValueStateMultiGet() throws Exception {
- ForStValueState<Integer, String> valueState1 = buildForStValueState("test-multiGet-1");
- ForStValueState<Integer, String> valueState2 = buildForStValueState("test-multiGet-2");
+ ForStValueState<Integer, Void, String> valueState1 =
+ buildForStValueState("test-multiGet-1");
+ ForStValueState<Integer, Void, String> valueState2 =
+ buildForStValueState("test-multiGet-2");
List<ForStDBGetRequest<?, ?>> batchGetRequest = new ArrayList<>();
List<Tuple2<String, TestStateFuture<String>>> resultCheckList = new ArrayList<>();
int keyNum = 1000;
for (int i = 0; i < keyNum; i++) {
TestStateFuture<String> future = new TestStateFuture<>();
- ForStValueState<Integer, String> table = ((i % 2 == 0) ? valueState1 : valueState2);
+ ForStValueState<Integer, Void, String> table =
+ ((i % 2 == 0) ? valueState1 : valueState2);
ForStDBGetRequest<ContextKey<Integer>, String> request =
ForStDBGetRequest.of(buildContextKey(i), table, future);
batchGetRequest.add(request);
diff --git a/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStStateExecutorTest.java b/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStStateExecutorTest.java
index e8717750329..95ea527fa2a 100644
--- a/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStStateExecutorTest.java
+++ b/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStStateExecutorTest.java
@@ -42,8 +42,8 @@ public class ForStStateExecutorTest extends ForStDBOperationTestBase {
@SuppressWarnings("unchecked")
public void testExecuteValueStateRequest() throws Exception {
ForStStateExecutor forStStateExecutor = new ForStStateExecutor(4, db, new WriteOptions());
- ForStValueState<Integer, String> state1 = buildForStValueState("value-state-1");
- ForStValueState<Integer, String> state2 = buildForStValueState("value-state-2");
+ ForStValueState<Integer, Void, String> state1 = buildForStValueState("value-state-1");
+ ForStValueState<Integer, Void, String> state2 = buildForStValueState("value-state-2");
StateRequestContainer stateRequestContainer =
forStStateExecutor.createStateRequestContainer();
@@ -52,7 +52,7 @@ public class ForStStateExecutorTest extends ForStDBOperationTestBase {
// 1. Update value state: keyRange [0, keyNum)
int keyNum = 1000;
for (int i = 0; i < keyNum; i++) {
- ForStValueState<Integer, String> state = (i % 2 == 0 ? state1 : state2);
+ ForStValueState<Integer, Void, String> state = (i % 2 == 0 ? state1 : state2);
stateRequestContainer.offer(
buildStateRequest(state, StateRequestType.VALUE_UPDATE, i, "test-" + i, i * 2));
}
@@ -64,14 +64,14 @@ public class ForStStateExecutorTest extends ForStDBOperationTestBase {
// 2. Get value state: keyRange [0, keyNum)
// Update value state: keyRange [keyNum, keyNum + 100]
for (int i = 0; i < keyNum; i++) {
- ForStValueState<Integer, String> state = (i % 2 == 0 ? state1 : state2);
+ ForStValueState<Integer, Void, String> state = (i % 2 == 0 ? state1 : state2);
StateRequest<?, ?, ?> getRequest =
buildStateRequest(state, StateRequestType.VALUE_GET, i, null, i * 2);
stateRequestContainer.offer(getRequest);
checkList.add(getRequest);
}
for (int i = keyNum; i < keyNum + 100; i++) {
- ForStValueState<Integer, String> state = (i % 2 == 0 ? state1 : state2);
+ ForStValueState<Integer, Void, String> state = (i % 2 == 0 ? state1 : state2);
stateRequestContainer.offer(
buildStateRequest(state, StateRequestType.VALUE_UPDATE, i, "test-" + i, i * 2));
}
@@ -90,12 +90,12 @@ public class ForStStateExecutorTest extends ForStDBOperationTestBase {
// Update state with null-value : keyRange [keyNum, keyNum + 100]
stateRequestContainer = forStStateExecutor.createStateRequestContainer();
for (int i = keyNum - 100; i < keyNum; i++) {
- ForStValueState<Integer, String> state = (i % 2 == 0 ? state1 : state2);
+ ForStValueState<Integer, Void, String> state = (i % 2 == 0 ? state1 : state2);
stateRequestContainer.offer(
buildStateRequest(state, StateRequestType.CLEAR, i, null, i * 2));
}
for (int i = keyNum; i < keyNum + 100; i++) {
- ForStValueState<Integer, String> state = (i % 2 == 0 ? state1 : state2);
+ ForStValueState<Integer, Void, String> state = (i % 2 == 0 ? state1 : state2);
stateRequestContainer.offer(
buildStateRequest(state, StateRequestType.VALUE_UPDATE, i, null, i * 2));
}
@@ -105,7 +105,7 @@ public class ForStStateExecutorTest extends ForStDBOperationTestBase {
stateRequestContainer = forStStateExecutor.createStateRequestContainer();
checkList.clear();
for (int i = keyNum - 100; i < keyNum + 100; i++) {
- ForStValueState<Integer, String> state = (i % 2 == 0 ? state1 : state2);
+ ForStValueState<Integer, Void, String> state = (i % 2 == 0 ? state1 : state2);
StateRequest<?, ?, ?> getRequest =
buildStateRequest(state, StateRequestType.VALUE_GET, i, null, i * 2);
stateRequestContainer.offer(getRequest);
@@ -121,8 +121,8 @@ public class ForStStateExecutorTest extends ForStDBOperationTestBase {
}
@SuppressWarnings({"rawtypes", "unchecked"})
- private <K, V, R> StateRequest<?, ?, ?> buildStateRequest(
- InternalKeyedState<K, V> innerTable,
+ private <K, N, V, R> StateRequest<?, ?, ?> buildStateRequest(
+ InternalKeyedState<K, N, V> innerTable,
StateRequestType requestType,
K key,
V value,
diff --git a/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStWriteBatchOperationTest.java b/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStWriteBatchOperationTest.java
index 19e11684fbf..dce11d12d77 100644
--- a/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStWriteBatchOperationTest.java
+++ b/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStWriteBatchOperationTest.java
@@ -34,8 +34,10 @@ public class ForStWriteBatchOperationTest extends ForStDBOperationTestBase {
@Test
public void testValueStateWriteBatch() throws Exception {
- ForStValueState<Integer, String> valueState1 = buildForStValueState("test-write-batch-1");
- ForStValueState<Integer, String> valueState2 = buildForStValueState("test-write-batch-2");
+ ForStValueState<Integer, Void, String> valueState1 =
+ buildForStValueState("test-write-batch-1");
+ ForStValueState<Integer, Void, String> valueState2 =
+ buildForStValueState("test-write-batch-2");
List<ForStDBPutRequest<?, ?>> batchPutRequest = new ArrayList<>();
int keyNum = 100;
for (int i = 0; i < keyNum; i++) {
@@ -62,7 +64,8 @@ public class ForStWriteBatchOperationTest extends ForStDBOperationTestBase {
@Test
public void testWriteBatchWithNullValue() throws Exception {
- ForStValueState<Integer, String> valueState = buildForStValueState("test-write-batch");
+ ForStValueState<Integer, Void, String> valueState =
+ buildForStValueState("test-write-batch");
List<ForStDBPutRequest<?, ?>> batchPutRequest = new ArrayList<>();
// 1. write some data without null value
int keyNum = 100;
(flink) 03/03: [FLINK-35858][State/ForSt] Support state namespace
Posted by za...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
zakelly pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit d4294c59e6f2ec8702f53916ea49cf23f6db8961
Author: Zakelly <za...@gmail.com>
AuthorDate: Thu Jul 18 19:13:20 2024 +0800
[FLINK-35858][State/ForSt] Support state namespace
---
.../org/apache/flink/state/forst/ContextKey.java | 11 ++++++++---
.../flink/state/forst/ForStDBGetRequest.java | 13 +++++++------
.../flink/state/forst/ForStDBPutRequest.java | 17 ++++++++++-------
.../state/forst/ForStGeneralMultiGetOperation.java | 6 +++---
.../apache/flink/state/forst/ForStInnerTable.java | 8 ++++----
.../flink/state/forst/ForStKeyedStateBackend.java | 1 +
.../flink/state/forst/ForStStateExecutor.java | 4 ++--
.../state/forst/ForStStateRequestClassifier.java | 8 ++++----
.../apache/flink/state/forst/ForStValueState.java | 22 +++++++++++++---------
.../state/forst/ForStWriteBatchOperation.java | 10 +++++-----
.../state/forst/ForStDBOperationTestBase.java | 4 +++-
.../forst/ForStGeneralMultiGetOperationTest.java | 4 ++--
.../state/forst/ForStWriteBatchOperationTest.java | 8 ++++----
13 files changed, 66 insertions(+), 50 deletions(-)
diff --git a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ContextKey.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ContextKey.java
index 05dd2f58796..9f73cefcc6e 100644
--- a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ContextKey.java
+++ b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ContextKey.java
@@ -19,6 +19,7 @@
package org.apache.flink.state.forst;
import org.apache.flink.runtime.asyncprocessing.RecordContext;
+import org.apache.flink.runtime.state.v2.InternalPartitionedState;
import org.apache.flink.util.function.FunctionWithException;
import javax.annotation.concurrent.ThreadSafe;
@@ -32,7 +33,7 @@ import java.util.Objects;
* @param <K> The type of the raw key.
*/
@ThreadSafe
-public class ContextKey<K> {
+public class ContextKey<K, N> {
private final RecordContext<K> recordContext;
@@ -48,6 +49,10 @@ public class ContextKey<K> {
return recordContext.getKeyGroup();
}
+ public N getNamespace(InternalPartitionedState<N> state) {
+ return recordContext.getNamespace(state);
+ }
+
/**
* Get the serialized key. If the cached serialized key within {@code RecordContext#payload} is
* null, the provided serialization function will be called, and the serialization result will
@@ -57,7 +62,7 @@ public class ContextKey<K> {
* @return the serialized bytes.
*/
public byte[] getOrCreateSerializedKey(
- FunctionWithException<ContextKey<K>, byte[], IOException> serializeKeyFunc)
+ FunctionWithException<ContextKey<K, N>, byte[], IOException> serializeKeyFunc)
throws IOException {
if (recordContext.getExtra() != null) {
return (byte[]) recordContext.getExtra();
@@ -84,7 +89,7 @@ public class ContextKey<K> {
if (o == null || getClass() != o.getClass()) {
return false;
}
- ContextKey<?> that = (ContextKey<?>) o;
+ ContextKey<?, ?> that = (ContextKey<?, ?>) o;
return Objects.equals(recordContext, that.recordContext);
}
}
diff --git a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStDBGetRequest.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStDBGetRequest.java
index 11868defb17..24c699cc018 100644
--- a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStDBGetRequest.java
+++ b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStDBGetRequest.java
@@ -30,13 +30,14 @@ import java.io.IOException;
* @param <K> The type of key in get access request.
* @param <V> The type of value returned by get request.
*/
-public class ForStDBGetRequest<K, V> {
+public class ForStDBGetRequest<K, N, V> {
- private final K key;
- private final ForStInnerTable<K, V> table;
+ private final ContextKey<K, N> key;
+ private final ForStInnerTable<K, N, V> table;
private final InternalStateFuture<V> future;
- private ForStDBGetRequest(K key, ForStInnerTable<K, V> table, InternalStateFuture<V> future) {
+ private ForStDBGetRequest(
+ ContextKey<K, N> key, ForStInnerTable<K, N, V> table, InternalStateFuture<V> future) {
this.key = key;
this.table = table;
this.future = future;
@@ -63,8 +64,8 @@ public class ForStDBGetRequest<K, V> {
future.completeExceptionally(message, ex);
}
- static <K, V> ForStDBGetRequest<K, V> of(
- K key, ForStInnerTable<K, V> table, InternalStateFuture<V> future) {
+ static <K, N, V> ForStDBGetRequest<K, N, V> of(
+ ContextKey<K, N> key, ForStInnerTable<K, N, V> table, InternalStateFuture<V> future) {
return new ForStDBGetRequest<>(key, table, future);
}
}
diff --git a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStDBPutRequest.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStDBPutRequest.java
index cc868b8ebb3..bfe1b15d8cc 100644
--- a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStDBPutRequest.java
+++ b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStDBPutRequest.java
@@ -32,15 +32,18 @@ import java.io.IOException;
* @param <K> The type of key in put access request.
* @param <V> The type of value in put access request.
*/
-public class ForStDBPutRequest<K, V> {
+public class ForStDBPutRequest<K, N, V> {
- private final K key;
+ private final ContextKey<K, N> key;
@Nullable private final V value;
- private final ForStInnerTable<K, V> table;
+ private final ForStInnerTable<K, N, V> table;
private final InternalStateFuture<Void> future;
private ForStDBPutRequest(
- K key, V value, ForStInnerTable<K, V> table, InternalStateFuture<Void> future) {
+ ContextKey<K, N> key,
+ V value,
+ ForStInnerTable<K, N, V> table,
+ InternalStateFuture<Void> future) {
this.key = key;
this.value = value;
this.table = table;
@@ -76,10 +79,10 @@ public class ForStDBPutRequest<K, V> {
* If the value of the ForStDBPutRequest is null, then the request will signify the deletion of
* the data associated with that key.
*/
- static <K, V> ForStDBPutRequest<K, V> of(
- K key,
+ static <K, N, V> ForStDBPutRequest<K, N, V> of(
+ ContextKey<K, N> key,
@Nullable V value,
- ForStInnerTable<K, V> table,
+ ForStInnerTable<K, N, V> table,
InternalStateFuture<Void> future) {
return new ForStDBPutRequest<>(key, value, table, future);
}
diff --git a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStGeneralMultiGetOperation.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStGeneralMultiGetOperation.java
index 77885d1a55a..ed7abbf7df8 100644
--- a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStGeneralMultiGetOperation.java
+++ b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStGeneralMultiGetOperation.java
@@ -38,12 +38,12 @@ public class ForStGeneralMultiGetOperation implements ForStDBOperation {
private final RocksDB db;
- private final List<ForStDBGetRequest<?, ?>> batchRequest;
+ private final List<ForStDBGetRequest<?, ?, ?>> batchRequest;
private final Executor executor;
ForStGeneralMultiGetOperation(
- RocksDB db, List<ForStDBGetRequest<?, ?>> batchRequest, Executor executor) {
+ RocksDB db, List<ForStDBGetRequest<?, ?, ?>> batchRequest, Executor executor) {
this.db = db;
this.batchRequest = batchRequest;
this.executor = executor;
@@ -58,7 +58,7 @@ public class ForStGeneralMultiGetOperation implements ForStDBOperation {
AtomicReference<Exception> error = new AtomicReference<>();
AtomicInteger counter = new AtomicInteger(batchRequest.size());
for (int i = 0; i < batchRequest.size(); i++) {
- ForStDBGetRequest<?, ?> request = batchRequest.get(i);
+ ForStDBGetRequest<?, ?, ?> request = batchRequest.get(i);
executor.execute(
() -> {
try {
diff --git a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStInnerTable.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStInnerTable.java
index 5d36b72097b..e8b763ee816 100644
--- a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStInnerTable.java
+++ b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStInnerTable.java
@@ -34,7 +34,7 @@ import java.io.IOException;
* @param <K> The key type of the table.
* @param <V> The value type of the table.
*/
-public interface ForStInnerTable<K, V> {
+public interface ForStInnerTable<K, N, V> {
/** Get the columnFamily handle corresponding to table. */
ColumnFamilyHandle getColumnFamilyHandle();
@@ -46,7 +46,7 @@ public interface ForStInnerTable<K, V> {
* @return the key bytes
* @throws IOException Thrown if the serialization encountered an I/O related error.
*/
- byte[] serializeKey(K key) throws IOException;
+ byte[] serializeKey(ContextKey<K, N> key) throws IOException;
/**
* Serialize the given value to the outputView.
@@ -73,7 +73,7 @@ public interface ForStInnerTable<K, V> {
* @param stateRequest The given stateRequest.
* @return The corresponding ForSt GetRequest.
*/
- ForStDBGetRequest<K, V> buildDBGetRequest(StateRequest<?, ?, ?> stateRequest);
+ ForStDBGetRequest<K, N, V> buildDBGetRequest(StateRequest<?, ?, ?> stateRequest);
/**
* Build a {@link ForStDBPutRequest} that belong to {@code ForStInnerTable} with the given
@@ -82,5 +82,5 @@ public interface ForStInnerTable<K, V> {
* @param stateRequest The given stateRequest.
* @return The corresponding ForSt PutRequest.
*/
- ForStDBPutRequest<K, V> buildDBPutRequest(StateRequest<?, ?, ?> stateRequest);
+ ForStDBPutRequest<K, N, V> buildDBPutRequest(StateRequest<?, ?, ?> stateRequest);
}
diff --git a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStKeyedStateBackend.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStKeyedStateBackend.java
index e80cd18b2d2..10bfb819cf3 100644
--- a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStKeyedStateBackend.java
+++ b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStKeyedStateBackend.java
@@ -154,6 +154,7 @@ public class ForStKeyedStateBackend<K> implements AsyncKeyedStateBackend {
columnFamilyHandle,
(ValueStateDescriptor<SV>) stateDesc,
serializedKeyBuilder,
+ namespaceSerializer::duplicate,
valueSerializerView,
valueDeserializerView);
}
diff --git a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStStateExecutor.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStStateExecutor.java
index e96ee32010d..7973ab05f3b 100644
--- a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStStateExecutor.java
+++ b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStStateExecutor.java
@@ -83,7 +83,7 @@ public class ForStStateExecutor implements StateExecutor {
() -> {
long startTime = System.currentTimeMillis();
List<CompletableFuture<Void>> futures = new ArrayList<>(2);
- List<ForStDBPutRequest<?, ?>> putRequests =
+ List<ForStDBPutRequest<?, ?, ?>> putRequests =
stateRequestClassifier.pollDbPutRequests();
if (!putRequests.isEmpty()) {
ForStWriteBatchOperation writeOperations =
@@ -92,7 +92,7 @@ public class ForStStateExecutor implements StateExecutor {
futures.add(writeOperations.process());
}
- List<ForStDBGetRequest<?, ?>> getRequests =
+ List<ForStDBGetRequest<?, ?, ?>> getRequests =
stateRequestClassifier.pollDbGetRequests();
if (!getRequests.isEmpty()) {
ForStGeneralMultiGetOperation getOperations =
diff --git a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStStateRequestClassifier.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStStateRequestClassifier.java
index 42324eccb67..1cfc3ca4c8f 100644
--- a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStStateRequestClassifier.java
+++ b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStStateRequestClassifier.java
@@ -31,9 +31,9 @@ import java.util.List;
*/
public class ForStStateRequestClassifier implements StateRequestContainer {
- private final List<ForStDBGetRequest<?, ?>> dbGetRequests;
+ private final List<ForStDBGetRequest<?, ?, ?>> dbGetRequests;
- private final List<ForStDBPutRequest<?, ?>> dbPutRequests;
+ private final List<ForStDBPutRequest<?, ?, ?>> dbPutRequests;
public ForStStateRequestClassifier() {
this.dbGetRequests = new ArrayList<>();
@@ -88,11 +88,11 @@ public class ForStStateRequestClassifier implements StateRequestContainer {
}
}
- public List<ForStDBGetRequest<?, ?>> pollDbGetRequests() {
+ public List<ForStDBGetRequest<?, ?, ?>> pollDbGetRequests() {
return dbGetRequests;
}
- public List<ForStDBPutRequest<?, ?>> pollDbPutRequests() {
+ public List<ForStDBPutRequest<?, ?, ?>> pollDbPutRequests() {
return dbPutRequests;
}
}
diff --git a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStValueState.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStValueState.java
index 153cdbc0fa1..d6bfb97b184 100644
--- a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStValueState.java
+++ b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStValueState.java
@@ -19,6 +19,7 @@
package org.apache.flink.state.forst;
import org.apache.flink.api.common.state.v2.ValueState;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.core.memory.DataInputDeserializer;
import org.apache.flink.core.memory.DataOutputSerializer;
import org.apache.flink.core.state.InternalStateFuture;
@@ -43,7 +44,7 @@ import java.util.function.Supplier;
* @param <V> The type of the value.
*/
public class ForStValueState<K, N, V> extends InternalValueState<K, N, V>
- implements ValueState<V>, ForStInnerTable<ContextKey<K>, V> {
+ implements ValueState<V>, ForStInnerTable<K, N, V> {
/** The column family which this internal value state belongs to. */
private final ColumnFamilyHandle columnFamilyHandle;
@@ -51,6 +52,8 @@ public class ForStValueState<K, N, V> extends InternalValueState<K, N, V>
/** The serialized key builder which should be thread-safe. */
private final ThreadLocal<SerializedCompositeKeyBuilder<K>> serializedKeyBuilder;
+ private final ThreadLocal<TypeSerializer<N>> namespaceSerializer;
+
/** The data outputStream used for value serializer, which should be thread-safe. */
private final ThreadLocal<DataOutputSerializer> valueSerializerView;
@@ -62,11 +65,13 @@ public class ForStValueState<K, N, V> extends InternalValueState<K, N, V>
ColumnFamilyHandle columnFamily,
ValueStateDescriptor<V> valueStateDescriptor,
Supplier<SerializedCompositeKeyBuilder<K>> serializedKeyBuilderInitializer,
+ Supplier<TypeSerializer<N>> namespaceSerializerInitializer,
Supplier<DataOutputSerializer> valueSerializerViewInitializer,
Supplier<DataInputDeserializer> valueDeserializerViewInitializer) {
super(stateRequestHandler, valueStateDescriptor);
this.columnFamilyHandle = columnFamily;
this.serializedKeyBuilder = ThreadLocal.withInitial(serializedKeyBuilderInitializer);
+ this.namespaceSerializer = ThreadLocal.withInitial(namespaceSerializerInitializer);
this.valueSerializerView = ThreadLocal.withInitial(valueSerializerViewInitializer);
this.valueDeserializerView = ThreadLocal.withInitial(valueDeserializerViewInitializer);
}
@@ -77,12 +82,13 @@ public class ForStValueState<K, N, V> extends InternalValueState<K, N, V>
}
@Override
- public byte[] serializeKey(ContextKey<K> contextKey) throws IOException {
+ public byte[] serializeKey(ContextKey<K, N> contextKey) throws IOException {
return contextKey.getOrCreateSerializedKey(
ctxKey -> {
SerializedCompositeKeyBuilder<K> builder = serializedKeyBuilder.get();
builder.setKeyAndKeyGroup(ctxKey.getRawKey(), ctxKey.getKeyGroup());
- return builder.build();
+ return builder.buildCompositeKeyNamespace(
+ contextKey.getNamespace(this), namespaceSerializer.get());
});
}
@@ -103,10 +109,9 @@ public class ForStValueState<K, N, V> extends InternalValueState<K, N, V>
@SuppressWarnings("unchecked")
@Override
- public ForStDBGetRequest<ContextKey<K>, V> buildDBGetRequest(
- StateRequest<?, ?, ?> stateRequest) {
+ public ForStDBGetRequest<K, N, V> buildDBGetRequest(StateRequest<?, ?, ?> stateRequest) {
Preconditions.checkArgument(stateRequest.getRequestType() == StateRequestType.VALUE_GET);
- ContextKey<K> contextKey =
+ ContextKey<K, N> contextKey =
new ContextKey<>((RecordContext<K>) stateRequest.getRecordContext());
return ForStDBGetRequest.of(
contextKey, this, (InternalStateFuture<V>) stateRequest.getFuture());
@@ -114,12 +119,11 @@ public class ForStValueState<K, N, V> extends InternalValueState<K, N, V>
@SuppressWarnings("unchecked")
@Override
- public ForStDBPutRequest<ContextKey<K>, V> buildDBPutRequest(
- StateRequest<?, ?, ?> stateRequest) {
+ public ForStDBPutRequest<K, N, V> buildDBPutRequest(StateRequest<?, ?, ?> stateRequest) {
Preconditions.checkArgument(
stateRequest.getRequestType() == StateRequestType.VALUE_UPDATE
|| stateRequest.getRequestType() == StateRequestType.CLEAR);
- ContextKey<K> contextKey =
+ ContextKey<K, N> contextKey =
new ContextKey<>((RecordContext<K>) stateRequest.getRecordContext());
V value =
(stateRequest.getRequestType() == StateRequestType.CLEAR)
diff --git a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStWriteBatchOperation.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStWriteBatchOperation.java
index 3d4e59e3247..188d5c49d20 100644
--- a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStWriteBatchOperation.java
+++ b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStWriteBatchOperation.java
@@ -34,7 +34,7 @@ public class ForStWriteBatchOperation implements ForStDBOperation {
private final RocksDB db;
- private final List<ForStDBPutRequest<?, ?>> batchRequest;
+ private final List<ForStDBPutRequest<?, ?, ?>> batchRequest;
private final WriteOptions writeOptions;
@@ -42,7 +42,7 @@ public class ForStWriteBatchOperation implements ForStDBOperation {
ForStWriteBatchOperation(
RocksDB db,
- List<ForStDBPutRequest<?, ?>> batchRequest,
+ List<ForStDBPutRequest<?, ?, ?>> batchRequest,
WriteOptions writeOptions,
Executor executor) {
this.db = db;
@@ -57,7 +57,7 @@ public class ForStWriteBatchOperation implements ForStDBOperation {
() -> {
try (WriteBatch writeBatch =
new WriteBatch(batchRequest.size() * PER_RECORD_ESTIMATE_BYTES)) {
- for (ForStDBPutRequest<?, ?> request : batchRequest) {
+ for (ForStDBPutRequest<?, ?, ?> request : batchRequest) {
if (request.valueIsNull()) {
// put(key, null) == delete(key)
writeBatch.delete(
@@ -71,12 +71,12 @@ public class ForStWriteBatchOperation implements ForStDBOperation {
}
}
db.write(writeOptions, writeBatch);
- for (ForStDBPutRequest<?, ?> request : batchRequest) {
+ for (ForStDBPutRequest<?, ?, ?> request : batchRequest) {
request.completeStateFuture();
}
} catch (Exception e) {
String msg = "Error while write batch data to ForStDB.";
- for (ForStDBPutRequest<?, ?> request : batchRequest) {
+ for (ForStDBPutRequest<?, ?, ?> request : batchRequest) {
// fail every state request in this batch
request.completeStateFutureExceptionally(msg, e);
}
diff --git a/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStDBOperationTestBase.java b/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStDBOperationTestBase.java
index fa09bfc46d7..3f31a110dab 100644
--- a/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStDBOperationTestBase.java
+++ b/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStDBOperationTestBase.java
@@ -22,6 +22,7 @@ import org.apache.flink.api.common.state.v2.State;
import org.apache.flink.api.common.state.v2.StateFuture;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
import org.apache.flink.api.common.typeutils.base.IntSerializer;
+import org.apache.flink.api.common.typeutils.base.VoidSerializer;
import org.apache.flink.configuration.ConfigConstants;
import org.apache.flink.core.memory.DataInputDeserializer;
import org.apache.flink.core.memory.DataOutputSerializer;
@@ -95,7 +96,7 @@ public class ForStDBOperationTestBase {
};
}
- protected ContextKey<Integer> buildContextKey(int i) {
+ protected ContextKey<Integer, Void> buildContextKey(int i) {
int keyGroup = KeyGroupRangeAssignment.assignToKeyGroup(i, 128);
RecordContext<Integer> recordContext =
new RecordContext<>(i, i, t -> {}, keyGroup, new Epoch(0));
@@ -117,6 +118,7 @@ public class ForStDBOperationTestBase {
cf,
valueStateDescriptor,
serializedKeyBuilder,
+ () -> VoidSerializer.INSTANCE,
valueSerializerView,
valueDeserializerView);
}
diff --git a/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStGeneralMultiGetOperationTest.java b/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStGeneralMultiGetOperationTest.java
index 2b24a65d650..910fab50643 100644
--- a/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStGeneralMultiGetOperationTest.java
+++ b/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStGeneralMultiGetOperationTest.java
@@ -38,7 +38,7 @@ public class ForStGeneralMultiGetOperationTest extends ForStDBOperationTestBase
buildForStValueState("test-multiGet-1");
ForStValueState<Integer, Void, String> valueState2 =
buildForStValueState("test-multiGet-2");
- List<ForStDBGetRequest<?, ?>> batchGetRequest = new ArrayList<>();
+ List<ForStDBGetRequest<?, ?, ?>> batchGetRequest = new ArrayList<>();
List<Tuple2<String, TestStateFuture<String>>> resultCheckList = new ArrayList<>();
int keyNum = 1000;
@@ -46,7 +46,7 @@ public class ForStGeneralMultiGetOperationTest extends ForStDBOperationTestBase
TestStateFuture<String> future = new TestStateFuture<>();
ForStValueState<Integer, Void, String> table =
((i % 2 == 0) ? valueState1 : valueState2);
- ForStDBGetRequest<ContextKey<Integer>, String> request =
+ ForStDBGetRequest<Integer, Void, String> request =
ForStDBGetRequest.of(buildContextKey(i), table, future);
batchGetRequest.add(request);
diff --git a/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStWriteBatchOperationTest.java b/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStWriteBatchOperationTest.java
index dce11d12d77..9e377debd88 100644
--- a/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStWriteBatchOperationTest.java
+++ b/flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStWriteBatchOperationTest.java
@@ -38,7 +38,7 @@ public class ForStWriteBatchOperationTest extends ForStDBOperationTestBase {
buildForStValueState("test-write-batch-1");
ForStValueState<Integer, Void, String> valueState2 =
buildForStValueState("test-write-batch-2");
- List<ForStDBPutRequest<?, ?>> batchPutRequest = new ArrayList<>();
+ List<ForStDBPutRequest<?, ?, ?>> batchPutRequest = new ArrayList<>();
int keyNum = 100;
for (int i = 0; i < keyNum; i++) {
batchPutRequest.add(
@@ -55,7 +55,7 @@ public class ForStWriteBatchOperationTest extends ForStDBOperationTestBase {
writeBatchOperation.process().get();
// check data correctness
- for (ForStDBPutRequest<?, ?> request : batchPutRequest) {
+ for (ForStDBPutRequest<?, ?, ?> request : batchPutRequest) {
byte[] keyBytes = request.buildSerializedKey();
byte[] valueBytes = db.get(request.getColumnFamilyHandle(), keyBytes);
assertArrayEquals(valueBytes, request.buildSerializedValue());
@@ -66,7 +66,7 @@ public class ForStWriteBatchOperationTest extends ForStDBOperationTestBase {
public void testWriteBatchWithNullValue() throws Exception {
ForStValueState<Integer, Void, String> valueState =
buildForStValueState("test-write-batch");
- List<ForStDBPutRequest<?, ?>> batchPutRequest = new ArrayList<>();
+ List<ForStDBPutRequest<?, ?, ?>> batchPutRequest = new ArrayList<>();
// 1. write some data without null value
int keyNum = 100;
for (int i = 0; i < keyNum; i++) {
@@ -103,7 +103,7 @@ public class ForStWriteBatchOperationTest extends ForStDBOperationTestBase {
writeBatchOperation2.process().get();
// 3. check data correctness
- for (ForStDBPutRequest<?, ?> request : batchPutRequest) {
+ for (ForStDBPutRequest<?, ?, ?> request : batchPutRequest) {
byte[] keyBytes = request.buildSerializedKey();
byte[] valueBytes = db.get(request.getColumnFamilyHandle(), keyBytes);
if (valueBytes == null) {
(flink) 02/03: [FLINK-35858][Runtime/State] Create state with namespaces
Posted by za...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
zakelly pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 055e11e6b4673d8fd1e62f54497a3553d7a21e4e
Author: Zakelly <za...@gmail.com>
AuthorDate: Thu Jul 18 17:59:52 2024 +0800
[FLINK-35858][Runtime/State] Create state with namespaces
---
.../flink/runtime/state/AsyncKeyedStateBackend.java | 11 ++++++++---
.../flink/runtime/state/v2/DefaultKeyedStateStoreV2.java | 16 +++++++++++-----
.../asyncprocessing/AsyncExecutionControllerTest.java | 5 ++++-
.../flink/runtime/state/StateBackendTestUtils.java | 3 ++-
.../runtime/state/v2/InternalKeyedStateTestBase.java | 5 ++++-
.../apache/flink/state/forst/ForStKeyedStateBackend.java | 3 ++-
.../api/operators/StreamOperatorStateHandler.java | 16 ++++++++++++++++
.../AbstractAsyncStateStreamOperator.java | 9 +++++++++
.../AbstractAsyncStateStreamOperatorV2.java | 9 +++++++++
.../api/operators/StreamingRuntimeContextTest.java | 6 ++++--
10 files changed, 69 insertions(+), 14 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AsyncKeyedStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AsyncKeyedStateBackend.java
index 3f49984bd7e..5348dc7cb5a 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AsyncKeyedStateBackend.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AsyncKeyedStateBackend.java
@@ -20,6 +20,7 @@ package org.apache.flink.runtime.state;
import org.apache.flink.annotation.Internal;
import org.apache.flink.api.common.state.v2.State;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.runtime.asyncprocessing.StateExecutor;
import org.apache.flink.runtime.asyncprocessing.StateRequestHandler;
import org.apache.flink.runtime.state.v2.StateDescriptor;
@@ -46,13 +47,17 @@ public interface AsyncKeyedStateBackend extends Disposable, Closeable {
/**
* Creates and returns a new state.
*
- * @param stateDesc The {@code StateDescriptor} that contains the name of the state.
- * @param <SV> The type of the stored state value.
+ * @param <N> the type of namespace for partitioning.
* @param <S> The type of the public API state.
+ * @param <SV> The type of the stored state value.
+ * @param namespaceSerializer the serializer for namespace.
+ * @param stateDesc The {@code StateDescriptor} that contains the name of the state.
* @throws Exception Exceptions may occur during initialization of the state.
*/
@Nonnull
- <SV, S extends State> S createState(@Nonnull StateDescriptor<SV> stateDesc) throws Exception;
+ <N, S extends State, SV> S createState(
+ TypeSerializer<N> namespaceSerializer, @Nonnull StateDescriptor<SV> stateDesc)
+ throws Exception;
/**
* Creates a {@code StateExecutor} which supports to execute a batch of state requests
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/DefaultKeyedStateStoreV2.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/DefaultKeyedStateStoreV2.java
index 002d5963786..803ad89c04d 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/DefaultKeyedStateStoreV2.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/DefaultKeyedStateStoreV2.java
@@ -24,6 +24,7 @@ import org.apache.flink.api.common.state.v2.MapState;
import org.apache.flink.api.common.state.v2.ReducingState;
import org.apache.flink.api.common.state.v2.ValueState;
import org.apache.flink.runtime.state.AsyncKeyedStateBackend;
+import org.apache.flink.runtime.state.VoidNamespaceSerializer;
import org.apache.flink.util.Preconditions;
import javax.annotation.Nonnull;
@@ -41,7 +42,8 @@ public class DefaultKeyedStateStoreV2 implements KeyedStateStoreV2 {
public <T> ValueState<T> getValueState(@Nonnull ValueStateDescriptor<T> stateProperties) {
Preconditions.checkNotNull(stateProperties, "The state properties must not be null");
try {
- return asyncKeyedStateBackend.createState(stateProperties);
+ return asyncKeyedStateBackend.createState(
+ VoidNamespaceSerializer.INSTANCE, stateProperties);
} catch (Exception e) {
throw new RuntimeException("Error while getting state", e);
}
@@ -51,7 +53,8 @@ public class DefaultKeyedStateStoreV2 implements KeyedStateStoreV2 {
public <T> ListState<T> getListState(@Nonnull ListStateDescriptor<T> stateProperties) {
Preconditions.checkNotNull(stateProperties, "The state properties must not be null");
try {
- return asyncKeyedStateBackend.createState(stateProperties);
+ return asyncKeyedStateBackend.createState(
+ VoidNamespaceSerializer.INSTANCE, stateProperties);
} catch (Exception e) {
throw new RuntimeException("Error while getting state", e);
}
@@ -62,7 +65,8 @@ public class DefaultKeyedStateStoreV2 implements KeyedStateStoreV2 {
@Nonnull MapStateDescriptor<UK, UV> stateProperties) {
Preconditions.checkNotNull(stateProperties, "The state properties must not be null");
try {
- return asyncKeyedStateBackend.createState(stateProperties);
+ return asyncKeyedStateBackend.createState(
+ VoidNamespaceSerializer.INSTANCE, stateProperties);
} catch (Exception e) {
throw new RuntimeException("Error while getting state", e);
}
@@ -73,7 +77,8 @@ public class DefaultKeyedStateStoreV2 implements KeyedStateStoreV2 {
@Nonnull ReducingStateDescriptor<T> stateProperties) {
Preconditions.checkNotNull(stateProperties, "The state properties must not be null");
try {
- return asyncKeyedStateBackend.createState(stateProperties);
+ return asyncKeyedStateBackend.createState(
+ VoidNamespaceSerializer.INSTANCE, stateProperties);
} catch (Exception e) {
throw new RuntimeException("Error while getting state", e);
}
@@ -84,7 +89,8 @@ public class DefaultKeyedStateStoreV2 implements KeyedStateStoreV2 {
@Nonnull AggregatingStateDescriptor<IN, ACC, OUT> stateProperties) {
Preconditions.checkNotNull(stateProperties, "The state properties must not be null");
try {
- return asyncKeyedStateBackend.createState(stateProperties);
+ return asyncKeyedStateBackend.createState(
+ VoidNamespaceSerializer.INSTANCE, stateProperties);
} catch (Exception e) {
throw new RuntimeException("Error while getting state", e);
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionControllerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionControllerTest.java
index 53337d71740..09e4cc5f0cb 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionControllerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionControllerTest.java
@@ -33,6 +33,7 @@ import org.apache.flink.runtime.mailbox.SyncMailboxExecutor;
import org.apache.flink.runtime.state.AsyncKeyedStateBackend;
import org.apache.flink.runtime.state.StateBackend;
import org.apache.flink.runtime.state.StateBackendTestUtils;
+import org.apache.flink.runtime.state.VoidNamespaceSerializer;
import org.apache.flink.runtime.state.v2.InternalValueState;
import org.apache.flink.runtime.state.v2.ValueStateDescriptor;
import org.apache.flink.util.FlinkRuntimeException;
@@ -108,7 +109,9 @@ class AsyncExecutionControllerTest {
asyncKeyedStateBackend.setup(aec);
try {
- valueState = asyncKeyedStateBackend.createState(stateDescriptor);
+ valueState =
+ asyncKeyedStateBackend.createState(
+ VoidNamespaceSerializer.INSTANCE, stateDescriptor);
} catch (Exception e) {
throw new RuntimeException(e);
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestUtils.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestUtils.java
index e1e0e5cb49c..a9981763f25 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestUtils.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestUtils.java
@@ -130,7 +130,8 @@ public class StateBackendTestUtils {
@Nonnull
@Override
@SuppressWarnings("unchecked")
- public <SV, S extends org.apache.flink.api.common.state.v2.State> S createState(
+ public <N, S extends org.apache.flink.api.common.state.v2.State, SV> S createState(
+ TypeSerializer<N> namespaceSerializer,
@Nonnull org.apache.flink.runtime.state.v2.StateDescriptor<SV> stateDesc) {
return (S) innerStateSupplier.get();
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalKeyedStateTestBase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalKeyedStateTestBase.java
index 8e098a055be..f5b5668003b 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalKeyedStateTestBase.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/v2/InternalKeyedStateTestBase.java
@@ -19,6 +19,7 @@
package org.apache.flink.runtime.state.v2;
import org.apache.flink.api.common.state.v2.State;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.runtime.asyncprocessing.AsyncExecutionController;
import org.apache.flink.runtime.asyncprocessing.StateExecutor;
import org.apache.flink.runtime.asyncprocessing.StateRequest;
@@ -126,7 +127,9 @@ public class InternalKeyedStateTestBase {
@Nonnull
@Override
- public <SV, S extends State> S createState(@Nonnull StateDescriptor<SV> stateDesc)
+ public <N, S extends State, SV> S createState(
+ TypeSerializer<N> namespaceSerializer,
+ @Nonnull StateDescriptor<SV> stateDesc)
throws Exception {
return null;
}
diff --git a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStKeyedStateBackend.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStKeyedStateBackend.java
index c5ec58bccc0..e80cd18b2d2 100644
--- a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStKeyedStateBackend.java
+++ b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStKeyedStateBackend.java
@@ -139,7 +139,8 @@ public class ForStKeyedStateBackend<K> implements AsyncKeyedStateBackend {
@Nonnull
@Override
@SuppressWarnings("unchecked")
- public <SV, S extends State> S createState(@Nonnull StateDescriptor<SV> stateDesc) {
+ public <N, S extends State, SV> S createState(
+ TypeSerializer<N> namespaceSerializer, @Nonnull StateDescriptor<SV> stateDesc) {
Preconditions.checkNotNull(
stateRequestHandler,
"A non-null stateRequestHandler must be setup before createState");
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamOperatorStateHandler.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamOperatorStateHandler.java
index 2033ba66ac8..ae035a67f49 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamOperatorStateHandler.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamOperatorStateHandler.java
@@ -367,6 +367,22 @@ public class StreamOperatorStateHandler {
}
}
+ /** Create new state (v2) based on new state descriptor. */
+ public <N, S extends org.apache.flink.api.common.state.v2.State, T> S getOrCreateKeyedState(
+ TypeSerializer<N> namespaceSerializer,
+ org.apache.flink.runtime.state.v2.StateDescriptor<T> stateDescriptor)
+ throws Exception {
+
+ if (asyncKeyedStateBackend != null) {
+ return asyncKeyedStateBackend.createState(namespaceSerializer, stateDescriptor);
+ } else {
+ throw new IllegalStateException(
+ "Cannot create partitioned state. "
+ + "The keyed state backend has not been set."
+ + "This indicates that the operator is not partitioned/keyed.");
+ }
+ }
+
/**
* Creates a partitioned state handle, using the state backend configured for this task.
*
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperator.java
index a87a9049583..4b96e8e60ca 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperator.java
@@ -21,6 +21,7 @@ package org.apache.flink.streaming.runtime.operators.asyncprocessing;
import org.apache.flink.annotation.Internal;
import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.api.common.operators.MailboxExecutor;
+import org.apache.flink.api.common.state.v2.State;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.runtime.asyncprocessing.AsyncExecutionController;
@@ -31,6 +32,7 @@ import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.state.AsyncKeyedStateBackend;
import org.apache.flink.runtime.state.CheckpointStreamFactory;
import org.apache.flink.runtime.state.KeyedStateBackend;
+import org.apache.flink.runtime.state.v2.StateDescriptor;
import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
import org.apache.flink.streaming.api.operators.Input;
import org.apache.flink.streaming.api.operators.InternalTimeServiceManager;
@@ -175,6 +177,13 @@ public abstract class AbstractAsyncStateStreamOperator<OUT> extends AbstractStre
getClass().getName(), inputId));
}
+ /** Create new state (v2) based on new state descriptor. */
+ protected <N, S extends State, T> S getOrCreateKeyedState(
+ TypeSerializer<N> namespaceSerializer, StateDescriptor<T> stateDescriptor)
+ throws Exception {
+ return stateHandler.getOrCreateKeyedState(namespaceSerializer, stateDescriptor);
+ }
+
@Override
public final OperatorSnapshotFutures snapshotState(
long checkpointId,
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2.java
index a142a59ccac..82787703216 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2.java
@@ -20,6 +20,7 @@ package org.apache.flink.streaming.runtime.operators.asyncprocessing;
import org.apache.flink.annotation.Internal;
import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.api.common.state.v2.State;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.runtime.asyncprocessing.AsyncExecutionController;
@@ -30,6 +31,7 @@ import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.state.AsyncKeyedStateBackend;
import org.apache.flink.runtime.state.CheckpointStreamFactory;
import org.apache.flink.runtime.state.KeyedStateBackend;
+import org.apache.flink.runtime.state.v2.StateDescriptor;
import org.apache.flink.streaming.api.operators.AbstractStreamOperatorV2;
import org.apache.flink.streaming.api.operators.InternalTimeServiceManager;
import org.apache.flink.streaming.api.operators.InternalTimerService;
@@ -150,6 +152,13 @@ public abstract class AbstractAsyncStateStreamOperatorV2<OUT> extends AbstractSt
+ " since this part is handled by the Input.");
}
+ /** Create new state (v2) based on new state descriptor. */
+ protected <N, S extends State, T> S getOrCreateKeyedState(
+ TypeSerializer<N> namespaceSerializer, StateDescriptor<T> stateDescriptor)
+ throws Exception {
+ return stateHandler.getOrCreateKeyedState(namespaceSerializer, stateDescriptor);
+ }
+
@Override
public final OperatorSnapshotFutures snapshotState(
long checkpointId,
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamingRuntimeContextTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamingRuntimeContextTest.java
index 144a22a21f5..97543d314f7 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamingRuntimeContextTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamingRuntimeContextTest.java
@@ -480,11 +480,13 @@ class StreamingRuntimeContextTest {
doAnswer(
(Answer<Object>)
invocationOnMock -> {
- ref.set(invocationOnMock.getArguments()[0]);
+ ref.set(invocationOnMock.getArguments()[1]);
return null;
})
.when(asyncKeyedStateBackend)
- .createState(any(org.apache.flink.runtime.state.v2.StateDescriptor.class));
+ .createState(
+ any(TypeSerializer.class),
+ any(org.apache.flink.runtime.state.v2.StateDescriptor.class));
operator.initializeState(streamTaskStateManager);
operator.getRuntimeContext().setKeyedStateStore(keyedStateStore);