You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by tz...@apache.org on 2019/04/10 07:58:43 UTC
[flink] branch master updated: [FLINK-12121] [State Backends] Use
composition instead of inheritance for InternalKeyContext in backends
This is an automated email from the ASF dual-hosted git repository.
tzulitai pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new 5dc8afd [FLINK-12121] [State Backends] Use composition instead of inheritance for InternalKeyContext in backends
5dc8afd is described below
commit 5dc8afd4c8bfea1b2aea0eb6f83ffce38bad3955
Author: Yu Li <li...@apache.org>
AuthorDate: Mon Apr 8 17:32:48 2019 +0800
[FLINK-12121] [State Backends] Use composition instead of inheritance for InternalKeyContext in backends
This closes #8122.
---
.../runtime/state/AbstractKeyedStateBackend.java | 43 +++++------
.../flink/runtime/state/KeyedStateBackend.java | 13 +++-
...AsyncSnapshotStrategySynchronicityBehavior.java | 6 +-
.../runtime/state/heap/CopyOnWriteStateTable.java | 27 ++++---
.../state/heap/CopyOnWriteStateTableSnapshot.java | 2 +-
.../runtime/state/heap/HeapKeyedStateBackend.java | 20 +++--
.../state/heap/HeapKeyedStateBackendBuilder.java | 35 +++++----
.../runtime/state/heap/HeapRestoreOperation.java | 11 ++-
.../runtime/state/heap/HeapSnapshotStrategy.java | 5 +-
.../runtime/state/heap/InternalKeyContext.java | 16 +++-
.../runtime/state/heap/InternalKeyContextImpl.java | 77 +++++++++++++++++++
.../runtime/state/heap/NestedMapsStateTable.java | 16 ++--
.../SnapshotStrategySynchronicityBehavior.java | 4 +-
.../flink/runtime/state/heap/StateTable.java | 17 ++++-
.../state/heap/StateTableByKeyGroupReaders.java | 16 ++--
.../SyncSnapshotStrategySynchronicityBehavior.java | 6 +-
.../state/heap/CopyOnWriteStateTableTest.java | 89 ++++------------------
.../heap/StateTableSnapshotCompatibilityTest.java | 14 ++--
.../state/ttl/mock/MockKeyedStateBackend.java | 8 +-
.../ttl/mock/MockKeyedStateBackendBuilder.java | 9 ++-
.../streaming/state/RocksDBKeyedStateBackend.java | 20 +++--
.../state/RocksDBKeyedStateBackendBuilder.java | 12 ++-
22 files changed, 274 insertions(+), 192 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractKeyedStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractKeyedStateBackend.java
index ab4cbf7..606e827 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractKeyedStateBackend.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractKeyedStateBackend.java
@@ -26,6 +26,7 @@ import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.core.fs.CloseableRegistry;
import org.apache.flink.runtime.checkpoint.CheckpointOptions;
import org.apache.flink.runtime.query.TaskKvStateRegistry;
+import org.apache.flink.runtime.state.heap.InternalKeyContext;
import org.apache.flink.runtime.state.internal.InternalKvState;
import org.apache.flink.runtime.state.ttl.TtlStateFactory;
import org.apache.flink.runtime.state.ttl.TtlTimeProvider;
@@ -55,15 +56,9 @@ public abstract class AbstractKeyedStateBackend<K> implements
/** {@link StateSerializerProvider} for our key serializer. */
protected final StateSerializerProvider<K> keySerializerProvider;
- /** The currently active key. */
- private K currentKey;
-
- /** Listeners to changes of keyed context ({@link #currentKey}). */
+ /** Listeners to changes of ({@link #keyContext}). */
private final ArrayList<KeySelectionListener<K>> keySelectionListeners;
- /** The key group of the currently active key. */
- private int currentKeyGroup;
-
/** So that we can give out state when the user uses the same key. */
private final HashMap<String, InternalKvState<K, ?, ?>> keyValueStatesByName;
@@ -94,25 +89,26 @@ public abstract class AbstractKeyedStateBackend<K> implements
/** Decorates the input and output streams to write key-groups compressed. */
protected final StreamCompressionDecorator keyGroupCompressionDecorator;
+ /** The key context for this backend. */
+ protected final InternalKeyContext<K> keyContext;
+
public AbstractKeyedStateBackend(
TaskKvStateRegistry kvStateRegistry,
TypeSerializer<K> keySerializer,
ClassLoader userCodeClassLoader,
- int numberOfKeyGroups,
- KeyGroupRange keyGroupRange,
ExecutionConfig executionConfig,
TtlTimeProvider ttlTimeProvider,
- CloseableRegistry cancelStreamRegistry) {
+ CloseableRegistry cancelStreamRegistry,
+ InternalKeyContext<K> keyContext) {
this(
kvStateRegistry,
StateSerializerProvider.fromNewRegisteredSerializer(keySerializer),
userCodeClassLoader,
- numberOfKeyGroups,
- keyGroupRange,
executionConfig,
ttlTimeProvider,
cancelStreamRegistry,
- determineStreamCompression(executionConfig)
+ determineStreamCompression(executionConfig),
+ keyContext
);
}
@@ -120,20 +116,20 @@ public abstract class AbstractKeyedStateBackend<K> implements
TaskKvStateRegistry kvStateRegistry,
StateSerializerProvider<K> keySerializerProvider,
ClassLoader userCodeClassLoader,
- int numberOfKeyGroups,
- KeyGroupRange keyGroupRange,
ExecutionConfig executionConfig,
TtlTimeProvider ttlTimeProvider,
CloseableRegistry cancelStreamRegistry,
- StreamCompressionDecorator keyGroupCompressionDecorator) {
+ StreamCompressionDecorator keyGroupCompressionDecorator,
+ InternalKeyContext<K> keyContext) {
+ this.keyContext = Preconditions.checkNotNull(keyContext);
+ this.numberOfKeyGroups = keyContext.getNumberOfKeyGroups();
+ this.keyGroupRange = Preconditions.checkNotNull(keyContext.getKeyGroupRange());
Preconditions.checkArgument(numberOfKeyGroups >= 1, "NumberOfKeyGroups must be a positive number");
Preconditions.checkArgument(numberOfKeyGroups >= keyGroupRange.getNumberOfKeyGroups(), "The total number of key groups must be at least the number in the key group range assigned to this backend");
this.kvStateRegistry = kvStateRegistry;
this.keySerializerProvider = keySerializerProvider;
- this.numberOfKeyGroups = numberOfKeyGroups;
this.userCodeClassLoader = Preconditions.checkNotNull(userCodeClassLoader);
- this.keyGroupRange = Preconditions.checkNotNull(keyGroupRange);
this.cancelStreamRegistry = cancelStreamRegistry;
this.keyValueStatesByName = new HashMap<>();
this.executionConfig = executionConfig;
@@ -175,8 +171,8 @@ public abstract class AbstractKeyedStateBackend<K> implements
@Override
public void setCurrentKey(K newKey) {
notifyKeySelected(newKey);
- this.currentKey = newKey;
- this.currentKeyGroup = KeyGroupRangeAssignment.assignToKeyGroup(newKey, numberOfKeyGroups);
+ this.keyContext.setCurrentKey(newKey);
+ this.keyContext.setCurrentKeyGroupIndex(KeyGroupRangeAssignment.assignToKeyGroup(newKey, numberOfKeyGroups));
}
private void notifyKeySelected(K newKey) {
@@ -209,21 +205,19 @@ public abstract class AbstractKeyedStateBackend<K> implements
*/
@Override
public K getCurrentKey() {
- return currentKey;
+ return this.keyContext.getCurrentKey();
}
/**
* @see KeyedStateBackend
*/
- @Override
public int getCurrentKeyGroupIndex() {
- return currentKeyGroup;
+ return this.keyContext.getCurrentKeyGroupIndex();
}
/**
* @see KeyedStateBackend
*/
- @Override
public int getNumberOfKeyGroups() {
return numberOfKeyGroups;
}
@@ -231,7 +225,6 @@ public abstract class AbstractKeyedStateBackend<K> implements
/**
* @see KeyedStateBackend
*/
- @Override
public KeyGroupRange getKeyGroupRange() {
return keyGroupRange;
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java
index 997e68b..4a59ead 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java
@@ -21,7 +21,6 @@ package org.apache.flink.runtime.state;
import org.apache.flink.api.common.state.State;
import org.apache.flink.api.common.state.StateDescriptor;
import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.runtime.state.heap.InternalKeyContext;
import org.apache.flink.util.Disposable;
import java.util.stream.Stream;
@@ -32,7 +31,7 @@ import java.util.stream.Stream;
* @param <K> The key by which state is keyed.
*/
public interface KeyedStateBackend<K>
- extends InternalKeyContext<K>, KeyedStateFactory, PriorityQueueSetFactory, Disposable {
+ extends KeyedStateFactory, PriorityQueueSetFactory, Disposable {
/**
* Sets the current key that is used for partitioned state.
@@ -41,6 +40,16 @@ public interface KeyedStateBackend<K>
void setCurrentKey(K newKey);
/**
+ * @return Current key.
+ */
+ K getCurrentKey();
+
+ /**
+ * @return Serializer of the key.
+ */
+ TypeSerializer<K> getKeySerializer();
+
+ /**
* Applies the provided {@link KeyedStateFunction} to the state with the provided
* {@link StateDescriptor} of all the currently active keys.
*
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/AsyncSnapshotStrategySynchronicityBehavior.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/AsyncSnapshotStrategySynchronicityBehavior.java
index 05d6475..bb6801f 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/AsyncSnapshotStrategySynchronicityBehavior.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/AsyncSnapshotStrategySynchronicityBehavior.java
@@ -18,6 +18,7 @@
package org.apache.flink.runtime.state.heap;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.runtime.state.RegisteredKeyValueStateBackendMetaInfo;
/**
@@ -35,7 +36,8 @@ class AsyncSnapshotStrategySynchronicityBehavior<K> implements SnapshotStrategyS
@Override
public <N, V> StateTable<K, N, V> newStateTable(
InternalKeyContext<K> keyContext,
- RegisteredKeyValueStateBackendMetaInfo<N, V> newMetaInfo) {
- return new CopyOnWriteStateTable<>(keyContext, newMetaInfo);
+ RegisteredKeyValueStateBackendMetaInfo<N, V> newMetaInfo,
+ TypeSerializer<K> keySerializer) {
+ return new CopyOnWriteStateTable<>(keyContext, newMetaInfo, keySerializer);
}
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/CopyOnWriteStateTable.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/CopyOnWriteStateTable.java
index 9c360da..5852bc2 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/CopyOnWriteStateTable.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/CopyOnWriteStateTable.java
@@ -211,24 +211,33 @@ public class CopyOnWriteStateTable<K, N, S> extends StateTable<K, N, S> implemen
/**
* Constructs a new {@code StateTable} with default capacity of {@code DEFAULT_CAPACITY}.
*
- * @param keyContext the key context.
- * @param metaInfo the meta information, including the type serializer for state copy-on-write.
+ * @param keyContext the key context.
+ * @param metaInfo the meta information, including the type serializer for state copy-on-write.
+ * @param keySerializer the serializer of the key.
*/
- CopyOnWriteStateTable(InternalKeyContext<K> keyContext, RegisteredKeyValueStateBackendMetaInfo<N, S> metaInfo) {
- this(keyContext, metaInfo, DEFAULT_CAPACITY);
+ CopyOnWriteStateTable(
+ InternalKeyContext<K> keyContext,
+ RegisteredKeyValueStateBackendMetaInfo<N, S> metaInfo,
+ TypeSerializer<K> keySerializer) {
+ this(keyContext, metaInfo, DEFAULT_CAPACITY, keySerializer);
}
/**
* Constructs a new {@code StateTable} instance with the specified capacity.
*
- * @param keyContext the key context.
- * @param metaInfo the meta information, including the type serializer for state copy-on-write.
- * @param capacity the initial capacity of this hash map.
+ * @param keyContext the key context.
+ * @param metaInfo the meta information, including the type serializer for state copy-on-write.
+ * @param capacity the initial capacity of this hash map.
+ * @param keySerializer the serializer of the key.
* @throws IllegalArgumentException when the capacity is less than zero.
*/
@SuppressWarnings("unchecked")
- private CopyOnWriteStateTable(InternalKeyContext<K> keyContext, RegisteredKeyValueStateBackendMetaInfo<N, S> metaInfo, int capacity) {
- super(keyContext, metaInfo);
+ private CopyOnWriteStateTable(
+ InternalKeyContext<K> keyContext,
+ RegisteredKeyValueStateBackendMetaInfo<N, S> metaInfo,
+ int capacity,
+ TypeSerializer<K> keySerializer) {
+ super(keyContext, metaInfo, keySerializer);
// initialized tables to EMPTY_TABLE.
this.primaryTable = (StateTableEntry<K, N, S>[]) EMPTY_TABLE;
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/CopyOnWriteStateTableSnapshot.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/CopyOnWriteStateTableSnapshot.java
index 12afcbc..86cd263 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/CopyOnWriteStateTableSnapshot.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/CopyOnWriteStateTableSnapshot.java
@@ -112,7 +112,7 @@ public class CopyOnWriteStateTableSnapshot<K, N, S>
// We create duplicates of the serializers for the async snapshot, because TypeSerializer
// might be stateful and shared with the event processing thread.
- this.localKeySerializer = owningStateTable.keyContext.getKeySerializer().duplicate();
+ this.localKeySerializer = owningStateTable.keySerializer.duplicate();
this.localNamespaceSerializer = owningStateTable.metaInfo.getNamespaceSerializer().duplicate();
this.localStateSerializer = owningStateTable.metaInfo.getStateSerializer().duplicate();
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java
index d3dee0d..25c3470 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java
@@ -37,7 +37,6 @@ import org.apache.flink.runtime.query.TaskKvStateRegistry;
import org.apache.flink.runtime.state.AbstractKeyedStateBackend;
import org.apache.flink.runtime.state.CheckpointStreamFactory;
import org.apache.flink.runtime.state.KeyExtractorFunction;
-import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.KeyGroupedInternalPriorityQueue;
import org.apache.flink.runtime.state.Keyed;
import org.apache.flink.runtime.state.KeyedStateFunction;
@@ -117,8 +116,6 @@ public class HeapKeyedStateBackend<K> extends AbstractKeyedStateBackend<K> {
TaskKvStateRegistry kvStateRegistry,
StateSerializerProvider<K> keySerializerProvider,
ClassLoader userCodeClassLoader,
- int numberOfKeyGroups,
- KeyGroupRange keyGroupRange,
ExecutionConfig executionConfig,
TtlTimeProvider ttlTimeProvider,
CloseableRegistry cancelStreamRegistry,
@@ -127,10 +124,17 @@ public class HeapKeyedStateBackend<K> extends AbstractKeyedStateBackend<K> {
Map<String, HeapPriorityQueueSnapshotRestoreWrapper> registeredPQStates,
LocalRecoveryConfig localRecoveryConfig,
HeapPriorityQueueSetFactory priorityQueueSetFactory,
- HeapSnapshotStrategy<K> snapshotStrategy
- ) {
- super(kvStateRegistry, keySerializerProvider, userCodeClassLoader, numberOfKeyGroups,
- keyGroupRange, executionConfig, ttlTimeProvider, cancelStreamRegistry, keyGroupCompressionDecorator);
+ HeapSnapshotStrategy<K> snapshotStrategy,
+ InternalKeyContext<K> keyContext) {
+ super(
+ kvStateRegistry,
+ keySerializerProvider,
+ userCodeClassLoader,
+ executionConfig,
+ ttlTimeProvider,
+ cancelStreamRegistry,
+ keyGroupCompressionDecorator,
+ keyContext);
this.registeredKVStates = registeredKVStates;
this.registeredPQStates = registeredPQStates;
this.localRecoveryConfig = localRecoveryConfig;
@@ -236,7 +240,7 @@ public class HeapKeyedStateBackend<K> extends AbstractKeyedStateBackend<K> {
newStateSerializer,
snapshotTransformFactory);
- stateTable = snapshotStrategy.newStateTable(this, newMetaInfo);
+ stateTable = snapshotStrategy.newStateTable(keyContext, newMetaInfo, keySerializerProvider.currentSchemaSerializer());
registeredKVStates.put(stateDesc.getName(), stateTable);
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackendBuilder.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackendBuilder.java
index c1bd0bb..71d048c 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackendBuilder.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackendBuilder.java
@@ -31,7 +31,6 @@ import org.apache.flink.runtime.state.StreamCompressionDecorator;
import org.apache.flink.runtime.state.ttl.TtlTimeProvider;
import javax.annotation.Nonnull;
-
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
@@ -95,21 +94,9 @@ public class HeapKeyedStateBackendBuilder<K> extends AbstractKeyedStateBackendBu
CloseableRegistry cancelStreamRegistryForBackend = new CloseableRegistry();
HeapSnapshotStrategy<K> snapshotStrategy = initSnapshotStrategy(
asynchronousSnapshots, registeredKVStates, registeredPQStates, cancelStreamRegistryForBackend);
- HeapKeyedStateBackend<K> backend = new HeapKeyedStateBackend<>(
- kvStateRegistry,
- keySerializerProvider,
- userCodeClassLoader,
- numberOfKeyGroups,
+ InternalKeyContext<K> keyContext = new InternalKeyContextImpl<>(
keyGroupRange,
- executionConfig,
- ttlTimeProvider,
- cancelStreamRegistryForBackend,
- keyGroupCompressionDecorator,
- registeredKVStates,
- registeredPQStates,
- localRecoveryConfig,
- priorityQueueSetFactory,
- snapshotStrategy
+ numberOfKeyGroups
);
HeapRestoreOperation<K> restoreOperation = new HeapRestoreOperation<>(
restoreStateHandles,
@@ -122,14 +109,26 @@ public class HeapKeyedStateBackendBuilder<K> extends AbstractKeyedStateBackendBu
keyGroupRange,
numberOfKeyGroups,
snapshotStrategy,
- backend);
+ keyContext);
try {
restoreOperation.restore();
} catch (Exception e) {
- backend.dispose();
throw new BackendBuildingException("Failed when trying to restore heap backend", e);
}
- return backend;
+ return new HeapKeyedStateBackend<>(
+ kvStateRegistry,
+ keySerializerProvider,
+ userCodeClassLoader,
+ executionConfig,
+ ttlTimeProvider,
+ cancelStreamRegistryForBackend,
+ keyGroupCompressionDecorator,
+ registeredKVStates,
+ registeredPQStates,
+ localRecoveryConfig,
+ priorityQueueSetFactory,
+ snapshotStrategy,
+ keyContext);
}
private HeapSnapshotStrategy<K> initSnapshotStrategy(
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapRestoreOperation.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapRestoreOperation.java
index 35966ba..b2ddbd1 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapRestoreOperation.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapRestoreOperation.java
@@ -73,7 +73,7 @@ public class HeapRestoreOperation<K> implements RestoreOperation<Void> {
@Nonnegative
private final int numberOfKeyGroups;
private final HeapSnapshotStrategy<K> snapshotStrategy;
- private final HeapKeyedStateBackend<K> backend;
+ private final InternalKeyContext<K> keyContext;
HeapRestoreOperation(
@Nonnull Collection<KeyedStateHandle> restoreStateHandles,
@@ -86,7 +86,7 @@ public class HeapRestoreOperation<K> implements RestoreOperation<Void> {
@Nonnull KeyGroupRange keyGroupRange,
int numberOfKeyGroups,
HeapSnapshotStrategy<K> snapshotStrategy,
- HeapKeyedStateBackend<K> backend) {
+ InternalKeyContext<K> keyContext) {
this.restoreStateHandles = restoreStateHandles;
this.keySerializerProvider = keySerializerProvider;
this.userCodeClassLoader = userCodeClassLoader;
@@ -97,7 +97,7 @@ public class HeapRestoreOperation<K> implements RestoreOperation<Void> {
this.keyGroupRange = keyGroupRange;
this.numberOfKeyGroups = numberOfKeyGroups;
this.snapshotStrategy = snapshotStrategy;
- this.backend = backend;
+ this.keyContext = keyContext;
}
@Override
@@ -181,7 +181,10 @@ public class HeapRestoreOperation<K> implements RestoreOperation<Void> {
new RegisteredKeyValueStateBackendMetaInfo<>(metaInfoSnapshot);
registeredKVStates.put(
metaInfoSnapshot.getName(),
- snapshotStrategy.newStateTable(backend, registeredKeyedBackendStateMetaInfo));
+ snapshotStrategy.newStateTable(
+ keyContext,
+ registeredKeyedBackendStateMetaInfo,
+ keySerializerProvider.currentSchemaSerializer()));
}
break;
case PRIORITY_QUEUE:
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapSnapshotStrategy.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapSnapshotStrategy.java
index 8a7d280..2ca9f5e 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapSnapshotStrategy.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapSnapshotStrategy.java
@@ -238,8 +238,9 @@ class HeapSnapshotStrategy<K>
@Override
public <N, V> StateTable<K, N, V> newStateTable(
InternalKeyContext<K> keyContext,
- RegisteredKeyValueStateBackendMetaInfo<N, V> newMetaInfo) {
- return snapshotStrategySynchronicityTrait.newStateTable(keyContext, newMetaInfo);
+ RegisteredKeyValueStateBackendMetaInfo<N, V> newMetaInfo,
+ TypeSerializer<K> keySerializer) {
+ return snapshotStrategySynchronicityTrait.newStateTable(keyContext, newMetaInfo, keySerializer);
}
private void processSnapshotMetaInfoForAllStates(
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/InternalKeyContext.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/InternalKeyContext.java
index cb0582b..ffc6796 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/InternalKeyContext.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/InternalKeyContext.java
@@ -19,9 +19,10 @@
package org.apache.flink.runtime.state.heap;
import org.apache.flink.annotation.Internal;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.runtime.state.KeyGroupRange;
+import javax.annotation.Nonnull;
+
/**
* This interface is the current context of a keyed state. It provides information about the currently selected key in
* the context, the corresponding key-group, and other key and key-grouping related information.
@@ -53,8 +54,17 @@ public interface InternalKeyContext<K> {
KeyGroupRange getKeyGroupRange();
/**
- * {@link TypeSerializer} for the state backend key type.
+ * Set current key of the context.
+ *
+ * @param currentKey the current key to set to.
+ */
+ void setCurrentKey(@Nonnull K currentKey);
+
+ /**
+ * Set current key group index of the context.
+ *
+ * @param currentKeyGroupIndex the current key group index to set to.
*/
- TypeSerializer<K> getKeySerializer();
+ void setCurrentKeyGroupIndex(int currentKeyGroupIndex);
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/InternalKeyContextImpl.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/InternalKeyContextImpl.java
new file mode 100644
index 0000000..8b7510a
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/InternalKeyContextImpl.java
@@ -0,0 +1,77 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.state.heap;
+
+import org.apache.flink.runtime.state.KeyGroupRange;
+
+import javax.annotation.Nonnegative;
+import javax.annotation.Nonnull;
+
+/**
+ * The default {@link InternalKeyContext} implementation.
+ *
+ * @param <K> Type of the key.
+ */
+public class InternalKeyContextImpl<K> implements InternalKeyContext<K> {
+ /** Range of key-groups for which this backend is responsible. */
+ private final KeyGroupRange keyGroupRange;
+ /** The number of key-groups aka max parallelism. */
+ private final int numberOfKeyGroups;
+
+ /** The currently active key. */
+ private K currentKey;
+ /** The key group of the currently active key. */
+ private int currentKeyGroupIndex;
+
+ public InternalKeyContextImpl(@Nonnull KeyGroupRange keyGroupRange, @Nonnegative int numberOfKeyGroups) {
+ this.keyGroupRange = keyGroupRange;
+ this.numberOfKeyGroups = numberOfKeyGroups;
+ }
+
+ @Override
+ public K getCurrentKey() {
+ return currentKey;
+ }
+
+ @Override
+ public int getCurrentKeyGroupIndex() {
+ return currentKeyGroupIndex;
+ }
+
+ @Override
+ public int getNumberOfKeyGroups() {
+ return numberOfKeyGroups;
+ }
+
+ @Override
+ public KeyGroupRange getKeyGroupRange() {
+ return keyGroupRange;
+ }
+
+ @Override
+ public void setCurrentKey(@Nonnull K currentKey) {
+ this.currentKey = currentKey;
+ }
+
+ @Override
+ public void setCurrentKeyGroupIndex(int currentKeyGroupIndex) {
+ this.currentKeyGroupIndex = currentKeyGroupIndex;
+ }
+
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/NestedMapsStateTable.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/NestedMapsStateTable.java
index f6b5557..e5e4998 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/NestedMapsStateTable.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/NestedMapsStateTable.java
@@ -76,12 +76,15 @@ public class NestedMapsStateTable<K, N, S> extends StateTable<K, N, S> {
/**
* Creates a new {@link NestedMapsStateTable} for the given key context and meta info.
- *
- * @param keyContext the key context.
+ * @param keyContext the key context.
* @param metaInfo the meta information for this state table.
+ * @param keySerializer the serializer of the key.
*/
- public NestedMapsStateTable(InternalKeyContext<K> keyContext, RegisteredKeyValueStateBackendMetaInfo<N, S> metaInfo) {
- super(keyContext, metaInfo);
+ public NestedMapsStateTable(
+ InternalKeyContext<K> keyContext,
+ RegisteredKeyValueStateBackendMetaInfo<N, S> metaInfo,
+ TypeSerializer<K> keySerializer) {
+ super(keyContext, metaInfo, keySerializer);
this.keyGroupOffset = keyContext.getKeyGroupRange().getStartKeyGroup();
@SuppressWarnings("unchecked")
@@ -350,11 +353,12 @@ public class NestedMapsStateTable<K, N, S> extends StateTable<K, N, S> {
private final StateSnapshotTransformer<S> snapshotFilter;
NestedMapsStateTableSnapshot(
- NestedMapsStateTable<K, N, S> owningTable, StateSnapshotTransformFactory<S> snapshotTransformFactory) {
+ NestedMapsStateTable<K, N, S> owningTable,
+ StateSnapshotTransformFactory<S> snapshotTransformFactory) {
super(owningTable);
this.snapshotFilter = snapshotTransformFactory.createForDeserializedState().orElse(null);
- this.keySerializer = owningStateTable.keyContext.getKeySerializer();
+ this.keySerializer = owningStateTable.keySerializer;
this.namespaceSerializer = owningStateTable.metaInfo.getNamespaceSerializer();
this.stateSerializer = owningStateTable.metaInfo.getStateSerializer();
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/SnapshotStrategySynchronicityBehavior.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/SnapshotStrategySynchronicityBehavior.java
index 3e963ed..ee79873 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/SnapshotStrategySynchronicityBehavior.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/SnapshotStrategySynchronicityBehavior.java
@@ -18,6 +18,7 @@
package org.apache.flink.runtime.state.heap;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.runtime.state.RegisteredKeyValueStateBackendMetaInfo;
/**
@@ -35,5 +36,6 @@ interface SnapshotStrategySynchronicityBehavior<K> {
<N, V> StateTable<K, N, V> newStateTable(
InternalKeyContext<K> keyContext,
- RegisteredKeyValueStateBackendMetaInfo<N, V> newMetaInfo);
+ RegisteredKeyValueStateBackendMetaInfo<N, V> newMetaInfo,
+ TypeSerializer<K> keySerializer);
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/StateTable.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/StateTable.java
index 7f1b9c1..101829c 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/StateTable.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/StateTable.java
@@ -52,13 +52,22 @@ public abstract class StateTable<K, N, S> implements StateSnapshotRestore {
protected RegisteredKeyValueStateBackendMetaInfo<N, S> metaInfo;
/**
- *
- * @param keyContext the key context provides the key scope for all put/get/delete operations.
- * @param metaInfo the meta information, including the type serializer for state copy-on-write.
+ * The serializer of the key.
+ */
+ protected final TypeSerializer<K> keySerializer;
+
+ /**
+ * @param keyContext the key context provides the key scope for all put/get/delete operations.
+ * @param metaInfo the meta information, including the type serializer for state copy-on-write.
+ * @param keySerializer the serializer of the key.
*/
- public StateTable(InternalKeyContext<K> keyContext, RegisteredKeyValueStateBackendMetaInfo<N, S> metaInfo) {
+ public StateTable(
+ InternalKeyContext<K> keyContext,
+ RegisteredKeyValueStateBackendMetaInfo<N, S> metaInfo,
+ TypeSerializer<K> keySerializer) {
this.keyContext = Preconditions.checkNotNull(keyContext);
this.metaInfo = Preconditions.checkNotNull(metaInfo);
+ this.keySerializer = Preconditions.checkNotNull(keySerializer);
}
// Main interface methods of StateTable -------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/StateTableByKeyGroupReaders.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/StateTableByKeyGroupReaders.java
index 0e2d005..4891063 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/StateTableByKeyGroupReaders.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/StateTableByKeyGroupReaders.java
@@ -41,14 +41,16 @@ class StateTableByKeyGroupReaders {
* Creates a new StateTableByKeyGroupReader that inserts de-serialized mappings into the given table, using the
* de-serialization algorithm that matches the given version.
*
- * @param stateTable the {@link StateTable} into which de-serialized mappings are inserted.
- * @param version version for the de-serialization algorithm.
* @param <K> type of key.
* @param <N> type of namespace.
* @param <S> type of state.
+ * @param stateTable the {@link StateTable} into which de-serialized mappings are inserted.
+ * @param version version for the de-serialization algorithm.
* @return the appropriate reader.
*/
- static <K, N, S> StateSnapshotKeyGroupReader readerForVersion(StateTable<K, N, S> stateTable, int version) {
+ static <K, N, S> StateSnapshotKeyGroupReader readerForVersion(
+ StateTable<K, N, S> stateTable,
+ int version) {
switch (version) {
case 1:
return new StateTableByKeyGroupReaderV1<>(stateTable);
@@ -63,10 +65,11 @@ class StateTableByKeyGroupReaders {
}
}
- private static <K, N, S> StateSnapshotKeyGroupReader createV2PlusReader(StateTable<K, N, S> stateTable) {
- final TypeSerializer<K> keySerializer = stateTable.keyContext.getKeySerializer();
+ private static <K, N, S> StateSnapshotKeyGroupReader createV2PlusReader(
+ StateTable<K, N, S> stateTable) {
final TypeSerializer<N> namespaceSerializer = stateTable.getNamespaceSerializer();
final TypeSerializer<S> stateSerializer = stateTable.getStateSerializer();
+ final TypeSerializer<K> keySerializer = stateTable.keySerializer;
final Tuple3<N, K, S> buffer = new Tuple3<>();
return KeyGroupPartitioner.createKeyGroupPartitionReader((in) -> {
buffer.f0 = namespaceSerializer.deserialize(in);
@@ -79,9 +82,11 @@ class StateTableByKeyGroupReaders {
static final class StateTableByKeyGroupReaderV1<K, N, S> implements StateSnapshotKeyGroupReader {
protected final StateTable<K, N, S> stateTable;
+ protected final TypeSerializer<K> keySerializer;
StateTableByKeyGroupReaderV1(StateTable<K, N, S> stateTable) {
this.stateTable = stateTable;
+ this.keySerializer = stateTable.keySerializer;
}
@Override
@@ -91,7 +96,6 @@ class StateTableByKeyGroupReaders {
return;
}
- final TypeSerializer<K> keySerializer = stateTable.keyContext.getKeySerializer();
final TypeSerializer<N> namespaceSerializer = stateTable.getNamespaceSerializer();
final TypeSerializer<S> stateSerializer = stateTable.getStateSerializer();
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/SyncSnapshotStrategySynchronicityBehavior.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/SyncSnapshotStrategySynchronicityBehavior.java
index 2d553ab..2c9752c 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/SyncSnapshotStrategySynchronicityBehavior.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/SyncSnapshotStrategySynchronicityBehavior.java
@@ -18,6 +18,7 @@
package org.apache.flink.runtime.state.heap;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.runtime.state.RegisteredKeyValueStateBackendMetaInfo;
/**
@@ -41,7 +42,8 @@ class SyncSnapshotStrategySynchronicityBehavior<K> implements SnapshotStrategySy
@Override
public <N, V> StateTable<K, N, V> newStateTable(
InternalKeyContext<K> keyContext,
- RegisteredKeyValueStateBackendMetaInfo<N, V> newMetaInfo) {
- return new NestedMapsStateTable<>(keyContext, newMetaInfo);
+ RegisteredKeyValueStateBackendMetaInfo<N, V> newMetaInfo,
+ TypeSerializer<K> keySerializer) {
+ return new NestedMapsStateTable<>(keyContext, newMetaInfo, keySerializer);
}
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/CopyOnWriteStateTableTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/CopyOnWriteStateTableTest.java
index 3dd4fdb..089dff6 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/CopyOnWriteStateTableTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/CopyOnWriteStateTableTest.java
@@ -45,6 +45,7 @@ import java.util.Map;
import java.util.Random;
public class CopyOnWriteStateTableTest extends TestLogger {
+ private final TypeSerializer<Integer> keySerializer = IntSerializer.INSTANCE;
/**
* Testing the basic map operations.
@@ -58,10 +59,10 @@ public class CopyOnWriteStateTableTest extends TestLogger {
IntSerializer.INSTANCE,
new ArrayListSerializer<>(IntSerializer.INSTANCE)); // we use mutable state objects.
- final MockInternalKeyContext<Integer> keyContext = new MockInternalKeyContext<>(IntSerializer.INSTANCE);
+ final MockInternalKeyContext<Integer> keyContext = new MockInternalKeyContext<>();
final CopyOnWriteStateTable<Integer, Integer, ArrayList<Integer>> stateTable =
- new CopyOnWriteStateTable<>(keyContext, metaInfo);
+ new CopyOnWriteStateTable<>(keyContext, metaInfo, keySerializer);
ArrayList<Integer> state_1_1 = new ArrayList<>();
state_1_1.add(41);
@@ -130,10 +131,10 @@ public class CopyOnWriteStateTableTest extends TestLogger {
IntSerializer.INSTANCE,
new ArrayListSerializer<>(IntSerializer.INSTANCE)); // we use mutable state objects.
- final MockInternalKeyContext<Integer> keyContext = new MockInternalKeyContext<>(IntSerializer.INSTANCE);
+ final MockInternalKeyContext<Integer> keyContext = new MockInternalKeyContext<>();
final CopyOnWriteStateTable<Integer, Integer, ArrayList<Integer>> stateTable =
- new CopyOnWriteStateTable<>(keyContext, metaInfo);
+ new CopyOnWriteStateTable<>(keyContext, metaInfo, keySerializer);
int insert = 0;
int remove = 0;
@@ -175,10 +176,10 @@ public class CopyOnWriteStateTableTest extends TestLogger {
IntSerializer.INSTANCE,
new ArrayListSerializer<>(IntSerializer.INSTANCE)); // we use mutable state objects.
- final MockInternalKeyContext<Integer> keyContext = new MockInternalKeyContext<>(IntSerializer.INSTANCE);
+ final MockInternalKeyContext<Integer> keyContext = new MockInternalKeyContext<>();
final CopyOnWriteStateTable<Integer, Integer, ArrayList<Integer>> stateTable =
- new CopyOnWriteStateTable<>(keyContext, metaInfo);
+ new CopyOnWriteStateTable<>(keyContext, metaInfo, keySerializer);
final HashMap<Tuple2<Integer, Integer>, ArrayList<Integer>> referenceMap = new HashMap<>();
@@ -379,10 +380,10 @@ public class CopyOnWriteStateTableTest extends TestLogger {
IntSerializer.INSTANCE,
new ArrayListSerializer<>(IntSerializer.INSTANCE)); // we use mutable state objects.
- final MockInternalKeyContext<Integer> keyContext = new MockInternalKeyContext<>(IntSerializer.INSTANCE);
+ final MockInternalKeyContext<Integer> keyContext = new MockInternalKeyContext<>();
final CopyOnWriteStateTable<Integer, Integer, ArrayList<Integer>> stateTable =
- new CopyOnWriteStateTable<>(keyContext, metaInfo);
+ new CopyOnWriteStateTable<>(keyContext, metaInfo, keySerializer);
ArrayList<Integer> originalState1 = new ArrayList<>(1);
ArrayList<Integer> originalState2 = new ArrayList<>(1);
@@ -454,36 +455,9 @@ public class CopyOnWriteStateTableTest extends TestLogger {
namespaceSerializer,
stateSerializer);
- final KeyGroupRange keyGroupRange = new KeyGroupRange(0, 0);
- InternalKeyContext<Integer> mockKeyContext = new InternalKeyContext<Integer>() {
- @Override
- public Integer getCurrentKey() {
- return 0;
- }
-
- @Override
- public int getCurrentKeyGroupIndex() {
- return 0;
- }
-
- @Override
- public int getNumberOfKeyGroups() {
- return 1;
- }
-
- @Override
- public KeyGroupRange getKeyGroupRange() {
- return keyGroupRange;
- }
-
- @Override
- public TypeSerializer<Integer> getKeySerializer() {
- return keySerializer;
- }
- };
-
+ InternalKeyContext<Integer> mockKeyContext = new MockInternalKeyContext<>();
CopyOnWriteStateTable<Integer, Integer, Integer> table =
- new CopyOnWriteStateTable<>(mockKeyContext, metaInfo);
+ new CopyOnWriteStateTable<>(mockKeyContext, metaInfo, keySerializer);
table.put(0, 0, 0, 0);
table.put(1, 0, 0, 1);
@@ -570,44 +544,9 @@ public class CopyOnWriteStateTableTest extends TestLogger {
}
}
- static class MockInternalKeyContext<T> implements InternalKeyContext<T> {
-
- private T key;
- private final TypeSerializer<T> serializer;
- private final KeyGroupRange keyGroupRange;
-
- public MockInternalKeyContext(TypeSerializer<T> serializer) {
- this.serializer = serializer;
- this.keyGroupRange = new KeyGroupRange(0, 0);
- }
-
- public void setKey(T key) {
- this.key = key;
- }
-
- @Override
- public T getCurrentKey() {
- return key;
- }
-
- @Override
- public int getCurrentKeyGroupIndex() {
- return 0;
- }
-
- @Override
- public int getNumberOfKeyGroups() {
- return 1;
- }
-
- @Override
- public KeyGroupRange getKeyGroupRange() {
- return keyGroupRange;
- }
-
- @Override
- public TypeSerializer<T> getKeySerializer() {
- return serializer;
+ static class MockInternalKeyContext<T> extends InternalKeyContextImpl<T> {
+ MockInternalKeyContext() {
+ super(new KeyGroupRange(0,0),1);
}
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/StateTableSnapshotCompatibilityTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/StateTableSnapshotCompatibilityTest.java
index d7fd3b9..81d3e65 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/StateTableSnapshotCompatibilityTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/StateTableSnapshotCompatibilityTest.java
@@ -19,6 +19,7 @@
package org.apache.flink.runtime.state.heap;
import org.apache.flink.api.common.state.StateDescriptor;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.common.typeutils.base.IntSerializer;
import org.apache.flink.core.memory.ByteArrayInputStreamWithPos;
import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos;
@@ -40,6 +41,7 @@ import java.util.ArrayList;
import java.util.Random;
public class StateTableSnapshotCompatibilityTest {
+ private final TypeSerializer<Integer> keySerializer = IntSerializer.INSTANCE;
/**
* This test ensures that different implementations of {@link StateTable} are compatible in their serialization
@@ -56,10 +58,10 @@ public class StateTableSnapshotCompatibilityTest {
new ArrayListSerializer<>(IntSerializer.INSTANCE));
final CopyOnWriteStateTableTest.MockInternalKeyContext<Integer> keyContext =
- new CopyOnWriteStateTableTest.MockInternalKeyContext<>(IntSerializer.INSTANCE);
+ new CopyOnWriteStateTableTest.MockInternalKeyContext<>();
CopyOnWriteStateTable<Integer, Integer, ArrayList<Integer>> cowStateTable =
- new CopyOnWriteStateTable<>(keyContext, metaInfo);
+ new CopyOnWriteStateTable<>(keyContext, metaInfo, keySerializer);
for (int i = 0; i < 100; ++i) {
ArrayList<Integer> list = new ArrayList<>(5);
@@ -74,7 +76,7 @@ public class StateTableSnapshotCompatibilityTest {
StateSnapshot snapshot = cowStateTable.stateSnapshot();
final NestedMapsStateTable<Integer, Integer, ArrayList<Integer>> nestedMapsStateTable =
- new NestedMapsStateTable<>(keyContext, metaInfo);
+ new NestedMapsStateTable<>(keyContext, metaInfo, keySerializer);
restoreStateTableFromSnapshot(nestedMapsStateTable, snapshot, keyContext.getKeyGroupRange());
snapshot.release();
@@ -86,7 +88,7 @@ public class StateTableSnapshotCompatibilityTest {
}
snapshot = nestedMapsStateTable.stateSnapshot();
- cowStateTable = new CopyOnWriteStateTable<>(keyContext, metaInfo);
+ cowStateTable = new CopyOnWriteStateTable<>(keyContext, metaInfo, keySerializer);
restoreStateTableFromSnapshot(cowStateTable, snapshot, keyContext.getKeyGroupRange());
snapshot.release();
@@ -97,8 +99,8 @@ public class StateTableSnapshotCompatibilityTest {
}
}
- private static <K, N, S> void restoreStateTableFromSnapshot(
- StateTable<K, N, S> stateTable,
+ private void restoreStateTableFromSnapshot(
+ StateTable<Integer, Integer, ArrayList<Integer>> stateTable,
StateSnapshot snapshot,
KeyGroupRange keyGroupRange) throws IOException {
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/ttl/mock/MockKeyedStateBackend.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/ttl/mock/MockKeyedStateBackend.java
index 7cee4c8..5c82b80 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/ttl/mock/MockKeyedStateBackend.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/ttl/mock/MockKeyedStateBackend.java
@@ -48,6 +48,7 @@ import org.apache.flink.runtime.state.StateSnapshotTransformer.StateSnapshotTran
import org.apache.flink.runtime.state.StateSnapshotTransformers;
import org.apache.flink.runtime.state.heap.HeapPriorityQueueElement;
import org.apache.flink.runtime.state.heap.HeapPriorityQueueSet;
+import org.apache.flink.runtime.state.heap.InternalKeyContext;
import org.apache.flink.runtime.state.ttl.TtlStateFactory;
import org.apache.flink.runtime.state.ttl.TtlTimeProvider;
import org.apache.flink.util.FlinkRuntimeException;
@@ -92,15 +93,14 @@ public class MockKeyedStateBackend<K> extends AbstractKeyedStateBackend<K> {
TaskKvStateRegistry kvStateRegistry,
TypeSerializer<K> keySerializer,
ClassLoader userCodeClassLoader,
- int numberOfKeyGroups,
- KeyGroupRange keyGroupRange,
ExecutionConfig executionConfig,
TtlTimeProvider ttlTimeProvider,
Map<String, Map<K, Map<Object, Object>>> stateValues,
Map<String, StateSnapshotTransformer<Object>> stateSnapshotFilters,
- CloseableRegistry cancelStreamRegistry) {
+ CloseableRegistry cancelStreamRegistry,
+ InternalKeyContext<K> keyContext) {
super(kvStateRegistry, keySerializer, userCodeClassLoader,
- numberOfKeyGroups, keyGroupRange, executionConfig, ttlTimeProvider, cancelStreamRegistry);
+ executionConfig, ttlTimeProvider, cancelStreamRegistry, keyContext);
this.stateValues = stateValues;
this.stateSnapshotFilters = stateSnapshotFilters;
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/ttl/mock/MockKeyedStateBackendBuilder.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/ttl/mock/MockKeyedStateBackendBuilder.java
index efef923..3ffe183 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/ttl/mock/MockKeyedStateBackendBuilder.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/ttl/mock/MockKeyedStateBackendBuilder.java
@@ -27,6 +27,7 @@ import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.KeyedStateHandle;
import org.apache.flink.runtime.state.StateSnapshotTransformer;
import org.apache.flink.runtime.state.StreamCompressionDecorator;
+import org.apache.flink.runtime.state.heap.InternalKeyContextImpl;
import org.apache.flink.runtime.state.ttl.TtlTimeProvider;
import javax.annotation.Nonnull;
@@ -74,12 +75,14 @@ public class MockKeyedStateBackendBuilder<K> extends AbstractKeyedStateBackendBu
kvStateRegistry,
keySerializerProvider.currentSchemaSerializer(),
userCodeClassLoader,
- numberOfKeyGroups,
- keyGroupRange,
executionConfig,
ttlTimeProvider,
stateValues,
stateSnapshotFilters,
- cancelStreamRegistry);
+ cancelStreamRegistry,
+ new InternalKeyContextImpl<>(
+ keyGroupRange,
+ numberOfKeyGroups
+ ));
}
}
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
index 4a6bd3c..11f773f 100644
--- a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
+++ b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
@@ -41,7 +41,6 @@ import org.apache.flink.runtime.checkpoint.CheckpointType;
import org.apache.flink.runtime.query.TaskKvStateRegistry;
import org.apache.flink.runtime.state.AbstractKeyedStateBackend;
import org.apache.flink.runtime.state.CheckpointStreamFactory;
-import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.KeyGroupedInternalPriorityQueue;
import org.apache.flink.runtime.state.Keyed;
import org.apache.flink.runtime.state.KeyedStateHandle;
@@ -55,6 +54,7 @@ import org.apache.flink.runtime.state.StateSnapshotTransformer.StateSnapshotTran
import org.apache.flink.runtime.state.StreamCompressionDecorator;
import org.apache.flink.runtime.state.heap.HeapPriorityQueueElement;
import org.apache.flink.runtime.state.heap.HeapPriorityQueueSetFactory;
+import org.apache.flink.runtime.state.heap.InternalKeyContext;
import org.apache.flink.runtime.state.ttl.TtlTimeProvider;
import org.apache.flink.util.FileUtils;
import org.apache.flink.util.FlinkRuntimeException;
@@ -201,8 +201,6 @@ public class RocksDBKeyedStateBackend<K> extends AbstractKeyedStateBackend<K> {
Function<String, ColumnFamilyOptions> columnFamilyOptionsFactory,
TaskKvStateRegistry kvStateRegistry,
StateSerializerProvider<K> keySerializerProvider,
- int numberOfKeyGroups,
- KeyGroupRange keyGroupRange,
ExecutionConfig executionConfig,
TtlTimeProvider ttlTimeProvider,
RocksDB db,
@@ -218,10 +216,18 @@ public class RocksDBKeyedStateBackend<K> extends AbstractKeyedStateBackend<K> {
RocksDBNativeMetricMonitor nativeMetricMonitor,
RocksDBSerializedCompositeKeyBuilder<K> sharedRocksKeyBuilder,
PriorityQueueSetFactory priorityQueueFactory,
- RocksDbTtlCompactFiltersManager ttlCompactFiltersManager) {
-
- super(kvStateRegistry, keySerializerProvider, userCodeClassLoader, numberOfKeyGroups,
- keyGroupRange, executionConfig, ttlTimeProvider, cancelStreamRegistry, keyGroupCompressionDecorator);
+ RocksDbTtlCompactFiltersManager ttlCompactFiltersManager,
+ InternalKeyContext<K> keyContext) {
+
+ super(
+ kvStateRegistry,
+ keySerializerProvider,
+ userCodeClassLoader,
+ executionConfig,
+ ttlTimeProvider,
+ cancelStreamRegistry,
+ keyGroupCompressionDecorator,
+ keyContext);
this.ttlCompactFiltersManager = ttlCompactFiltersManager;
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackendBuilder.java b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackendBuilder.java
index 9b98e6e..b55a7a0 100644
--- a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackendBuilder.java
+++ b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackendBuilder.java
@@ -43,6 +43,8 @@ import org.apache.flink.runtime.state.PriorityQueueSetFactory;
import org.apache.flink.runtime.state.StateHandleID;
import org.apache.flink.runtime.state.StreamCompressionDecorator;
import org.apache.flink.runtime.state.heap.HeapPriorityQueueSetFactory;
+import org.apache.flink.runtime.state.heap.InternalKeyContext;
+import org.apache.flink.runtime.state.heap.InternalKeyContextImpl;
import org.apache.flink.runtime.state.ttl.TtlTimeProvider;
import org.apache.flink.util.FileUtils;
import org.apache.flink.util.IOUtils;
@@ -323,6 +325,10 @@ public class RocksDBKeyedStateBackendBuilder<K> extends AbstractKeyedStateBacken
throw new BackendBuildingException(errMsg, e);
}
}
+ InternalKeyContext<K> keyContext = new InternalKeyContextImpl<>(
+ keyGroupRange,
+ numberOfKeyGroups
+ );
return new RocksDBKeyedStateBackend<>(
this.userCodeClassLoader,
this.instanceBasePath,
@@ -330,8 +336,6 @@ public class RocksDBKeyedStateBackendBuilder<K> extends AbstractKeyedStateBacken
columnFamilyOptionsFactory,
this.kvStateRegistry,
this.keySerializerProvider,
- this.numberOfKeyGroups,
- this.keyGroupRange,
this.executionConfig,
this.ttlTimeProvider,
db,
@@ -347,8 +351,8 @@ public class RocksDBKeyedStateBackendBuilder<K> extends AbstractKeyedStateBacken
nativeMetricMonitor,
sharedRocksKeyBuilder,
priorityQueueFactory,
- ttlCompactFiltersManager
- );
+ ttlCompactFiltersManager,
+ keyContext);
}
private AbstractRocksDBRestoreOperation<K> getRocksDBRestoreOperation(