org.apache.flink.runtime.state.KeyedStateBackend Java Examples

public static KeyedStateBackend<ByteBuffer> createStateBackend() throws Exception {
  MemoryStateBackend backend = new MemoryStateBackend();

  AbstractKeyedStateBackend<ByteBuffer> keyedStateBackend =
      backend.createKeyedStateBackend(
          new DummyEnvironment("test", 1, 0),
          new JobID(),
          "test_op",
          new GenericTypeInfo<>(ByteBuffer.class).createSerializer(new ExecutionConfig()),
          2,
          new KeyGroupRange(0, 1),
          new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()),
          TtlTimeProvider.DEFAULT,
          null,
          Collections.emptyList(),
          new CloseableRegistry());

  changeKey(keyedStateBackend);

  return keyedStateBackend;
}
 

@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createDescriptorCapturingMockOp(
		final AtomicReference<Object> ref, final ExecutionConfig config) throws Exception {

	AbstractStreamOperator<?> operatorMock = mock(AbstractStreamOperator.class);

	KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);

	DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);

	when(operatorMock.getExecutionConfig()).thenReturn(config);

	doAnswer(new Answer<Object>() {

		@Override
		public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
			ref.set(invocationOnMock.getArguments()[2]);
			return null;
		}
	}).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(StateDescriptor.class));

	when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore);
	when(operatorMock.getOperatorID()).thenReturn(new OperatorID());

	return operatorMock;
}
 

private Context(KeyedStateBackend<K> keyedStateBackend, InternalTimerService<VoidNamespace> timerService) throws Exception {
	eventTimers = keyedStateBackend.getPartitionedState(
		USER_TIMERS_NAME,
		StringSerializer.INSTANCE,
		new ListStateDescriptor<>(EVENT_TIMER_STATE, Types.LONG));

	timerService.forEachEventTimeTimer((namespace, timer) -> {
		if (namespace.equals(VoidNamespace.INSTANCE)) {
			eventTimers.add(timer);
		}
	});

	procTimers = keyedStateBackend.getPartitionedState(
		USER_TIMERS_NAME,
		StringSerializer.INSTANCE,
		new ListStateDescriptor<>(PROC_TIMER_STATE, Types.LONG));

	timerService.forEachProcessingTimeTimer((namespace, timer) -> {
		if (namespace.equals(VoidNamespace.INSTANCE)) {
			procTimers.add(timer);
		}
	});
}
 

FlinkCombiningStateWithContext(
    KeyedStateBackend<ByteBuffer> flinkStateBackend,
    String stateId,
    CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn,
    StateNamespace namespace,
    Coder<AccumT> accumCoder,
    CombineWithContext.Context context) {

  this.namespace = namespace;
  this.stateId = stateId;
  this.combineFn = combineFn;
  this.flinkStateBackend = flinkStateBackend;
  this.context = context;

  flinkStateDescriptor =
      new ValueStateDescriptor<>(stateId, new CoderTypeSerializer<>(accumCoder));
}
 

@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createDescriptorCapturingMockOp(
		final AtomicReference<Object> ref, final ExecutionConfig config) throws Exception {

	AbstractStreamOperator<?> operatorMock = mock(AbstractStreamOperator.class);

	KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);

	DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);

	when(operatorMock.getExecutionConfig()).thenReturn(config);

	doAnswer(new Answer<Object>() {

		@Override
		public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
			ref.set(invocationOnMock.getArguments()[2]);
			return null;
		}
	}).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(StateDescriptor.class));

	when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore);
	when(operatorMock.getOperatorID()).thenReturn(new OperatorID());

	return operatorMock;
}
 

static void fireExpiredAsyncOperations(
    MapStateDescriptor<Long, Message> asyncOperationStateDescriptor,
    Reductions reductions,
    MapState<Long, Message> asyncOperationState,
    KeyedStateBackend<String> keyedStateBackend)
    throws Exception {

  AsyncOperationFailureNotifier asyncOperationFailureNotifier =
      new AsyncOperationFailureNotifier(reductions, asyncOperationState);

  keyedStateBackend.applyToAllKeys(
      VoidNamespace.get(),
      VoidNamespaceSerializer.INSTANCE,
      asyncOperationStateDescriptor,
      asyncOperationFailureNotifier);

  if (asyncOperationFailureNotifier.enqueued()) {
    reductions.processEnvelopes();
  }
}
 

public static <InputT, OutputT> BufferingDoFnRunner<InputT, OutputT> create(
    DoFnRunner<InputT, OutputT> doFnRunner,
    String stateName,
    org.apache.beam.sdk.coders.Coder windowedInputCoder,
    org.apache.beam.sdk.coders.Coder windowCoder,
    OperatorStateBackend operatorStateBackend,
    @Nullable KeyedStateBackend<Object> keyedStateBackend,
    int maxConcurrentCheckpoints)
    throws Exception {
  return new BufferingDoFnRunner<>(
      doFnRunner,
      stateName,
      windowedInputCoder,
      windowCoder,
      operatorStateBackend,
      keyedStateBackend,
      maxConcurrentCheckpoints);
}
 

public TestOnMergeContext(
		K key,
		W targetWindow,
		Collection<W> mergedWindows,
		InternalTimerService<W> timerService,
		KeyedStateBackend<Integer> stateBackend,
		TypeSerializer<W> windowSerializer) {
	super(key, targetWindow, timerService, stateBackend, windowSerializer);

	this.mergedWindows = mergedWindows;
}
 

public final void setup(
	ExecutionConfig executionConfig,
	KeyedStateBackend<KEY> keyKeyedStateBackend,
	InternalTimeServiceManager<KEY> timerServiceManager,
	SavepointRuntimeContext ctx) {

	this.executionConfig = executionConfig;
	this.keyedStateBackend = keyKeyedStateBackend;
	this.timerServiceManager = timerServiceManager;
	this.keySerializer = keyType.createSerializer(executionConfig);

	FunctionUtils.setFunctionRuntimeContext(function, ctx);
}
 

public static KeyedStateBackend<Long> createKeyedStateBackend(StateBackendType backendType) throws IOException {
	switch (backendType) {
		case HEAP:
			rootDir = prepareDirectory(rootDirName, null);
			return createHeapKeyedStateBackend(rootDir);
		case ROCKSDB:
			rootDir = prepareDirectory(rootDirName, null);
			return createRocksDBKeyedStateBackend(rootDir);
		default:
			throw new IllegalArgumentException("Unknown backend type: " + backendType);
	}
}
 

public PerWindowStateDataViewStore(
		KeyedStateBackend<?> keyedStateBackend,
		TypeSerializer<?> windowSerializer,
		RuntimeContext runtimeContext) {
	this.keyedStateBackend = keyedStateBackend;
	this.windowSerializer = windowSerializer;
	this.ctx = runtimeContext;
}
 

@Override
protected StateInternals createStateInternals() {
  try {
    KeyedStateBackend<ByteBuffer> keyedStateBackend = createStateBackend();
    return new FlinkStateInternals<>(keyedStateBackend, StringUtf8Coder.of());
  } catch (Exception e) {
    throw new RuntimeException(e);
  }
}
 

@Test
public void testCompactState() throws Exception {
	KeyedStateBackend<Long> backend = createKeyedStateBackend(backendType);
	ListState<Long> listState = getListState(backend, listStateDescriptor);
	for (long i = 0; i < 10; i++) {
		backend.setCurrentKey(i);
		listState.add(i);
	}
	if (backend instanceof RocksDBKeyedStateBackend) {
		RocksDBKeyedStateBackend<Long> rocksDBKeyedStateBackend = (RocksDBKeyedStateBackend<Long>) backend;
		compactState(rocksDBKeyedStateBackend, listStateDescriptor);
	}
	cleanUp(backend);
}
 

private void populateOutputTimestampQueue() {
  Preconditions.checkState(
      outputTimestampQueue.isEmpty(),
      "Output timestamp queue should be empty when recomputing the minimum output timestamp across all timers.");
  final KeyedStateBackend<Object> keyedStateBackend = getKeyedStateBackend();
  final Object currentKey = keyedStateBackend.getCurrentKey();
  try (Stream<Object> keys =
      keyedStateBackend.getKeys(PENDING_TIMERS_STATE_NAME, VoidNamespace.INSTANCE)) {
    keys.forEach(
        key -> {
          keyedStateBackend.setCurrentKey(key);
          try {
            for (TimerData timerData : pendingTimersById.values()) {
              if (timerData.getDomain() == TimeDomain.EVENT_TIME) {
                long outputTimeStampMs = timerData.getOutputTimestamp().getMillis();
                if (timerUsesOutputTimestamp(timerData)) {
                  outputTimestampQueue.add(outputTimeStampMs);
                }
              }
            }
          } catch (Exception e) {
            throw new RuntimeException(
                "Exception while reading set of timers for key: " + key, e);
          }
        });
  } finally {
    if (currentKey != null) {
      keyedStateBackend.setCurrentKey(currentKey);
    }
  }
}
 

FlinkMapState(
    KeyedStateBackend<ByteBuffer> flinkStateBackend,
    String stateId,
    StateNamespace namespace,
    Coder<KeyT> mapKeyCoder,
    Coder<ValueT> mapValueCoder) {
  this.namespace = namespace;
  this.stateId = stateId;
  this.flinkStateBackend = flinkStateBackend;
  this.flinkStateDescriptor =
      new MapStateDescriptor<>(
          stateId,
          new CoderTypeSerializer<>(mapKeyCoder),
          new CoderTypeSerializer<>(mapValueCoder));
}
 

FlinkBagState(
    KeyedStateBackend<ByteBuffer> flinkStateBackend,
    String stateId,
    StateNamespace namespace,
    Coder<T> coder) {

  this.namespace = namespace;
  this.stateId = stateId;
  this.flinkStateBackend = flinkStateBackend;
  this.storesVoidValues = coder instanceof VoidCoder;
  this.flinkStateDescriptor =
      new ListStateDescriptor<>(stateId, new CoderTypeSerializer<>(coder));
}
 

FlinkCombiningState(
    KeyedStateBackend<ByteBuffer> flinkStateBackend,
    String stateId,
    Combine.CombineFn<InputT, AccumT, OutputT> combineFn,
    StateNamespace namespace,
    Coder<AccumT> accumCoder) {

  this.namespace = namespace;
  this.stateId = stateId;
  this.combineFn = combineFn;
  this.flinkStateBackend = flinkStateBackend;

  flinkStateDescriptor =
      new ValueStateDescriptor<>(stateId, new CoderTypeSerializer<>(accumCoder));
}
 

FlinkValueState(
    KeyedStateBackend<ByteBuffer> flinkStateBackend,
    String stateId,
    StateNamespace namespace,
    Coder<T> coder) {

  this.namespace = namespace;
  this.stateId = stateId;
  this.flinkStateBackend = flinkStateBackend;

  flinkStateDescriptor = new ValueStateDescriptor<>(stateId, new CoderTypeSerializer<>(coder));
}
 

public int numKeyedStateEntries() {
	AbstractStreamOperator<?> abstractStreamOperator = (AbstractStreamOperator<?>) operator;
	KeyedStateBackend<Object> keyedStateBackend = abstractStreamOperator.getKeyedStateBackend();
	if (keyedStateBackend instanceof HeapKeyedStateBackend) {
		return ((HeapKeyedStateBackend) keyedStateBackend).numKeyValueStateEntries();
	} else {
		throw new UnsupportedOperationException();
	}
}
 

private BufferingDoFnRunner(
    DoFnRunner<InputT, OutputT> underlying,
    String stateName,
    org.apache.beam.sdk.coders.Coder inputCoder,
    org.apache.beam.sdk.coders.Coder windowCoder,
    OperatorStateBackend operatorStateBackend,
    @Nullable KeyedStateBackend keyedStateBackend,
    int maxConcurrentCheckpoints)
    throws Exception {
  Preconditions.checkArgument(
      maxConcurrentCheckpoints > 0 && maxConcurrentCheckpoints < Short.MAX_VALUE,
      "Maximum number of concurrent checkpoints not within the bounds of 0 and %s",
      Short.MAX_VALUE);

  this.underlying = underlying;
  this.notYetAcknowledgedSnapshots =
      operatorStateBackend.getUnionListState(
          new ListStateDescriptor<>("notYetAcknowledgedSnapshots", CheckpointIdentifier.class));
  this.bufferingElementsHandlerFactory =
      (stateId) -> {
        ListStateDescriptor<BufferedElement> stateDescriptor =
            new ListStateDescriptor<>(
                stateName + stateId,
                new CoderTypeSerializer<>(
                    new BufferedElements.Coder(inputCoder, windowCoder, null)));
        if (keyedStateBackend != null) {
          return KeyedBufferingElementsHandler.create(keyedStateBackend, stateDescriptor);
        } else {
          return NonKeyedBufferingElementsHandler.create(
              operatorStateBackend.getListState(stateDescriptor));
        }
      };
  this.numCheckpointBuffers = initializeState(maxConcurrentCheckpoints);
  this.currentBufferingElementsHandler =
      bufferingElementsHandlerFactory.get(rotateAndGetStateIndex());
}
 

public static <K, S extends State, T> void applyToAllKeys(
	KeyedStateBackend<K> backend,
	final StateDescriptor<S, T> stateDescriptor,
	final KeyedStateFunction<K, S> function) throws Exception {
	backend.applyToAllKeys(
		VoidNamespace.INSTANCE,
		VoidNamespaceSerializer.INSTANCE,
		stateDescriptor,
		function);
}
 

public <N> int numKeyedStateEntries(N namespace) {
	AbstractStreamOperator<?> abstractStreamOperator = (AbstractStreamOperator<?>) operator;
	KeyedStateBackend<Object> keyedStateBackend = abstractStreamOperator.getKeyedStateBackend();
	if (keyedStateBackend instanceof HeapKeyedStateBackend) {
		return ((HeapKeyedStateBackend) keyedStateBackend).numKeyValueStateEntries(namespace);
	} else {
		throw new UnsupportedOperationException();
	}
}
 

public int numKeyedStateEntries() {
	AbstractStreamOperator<?> abstractStreamOperator = (AbstractStreamOperator<?>) operator;
	KeyedStateBackend<Object> keyedStateBackend = abstractStreamOperator.getKeyedStateBackend();
	if (keyedStateBackend instanceof HeapKeyedStateBackend) {
		return ((HeapKeyedStateBackend) keyedStateBackend).numKeyValueStateEntries();
	} else {
		throw new UnsupportedOperationException();
	}
}
 

@Inject
public FlinkState(
    @Label("runtime-context") RuntimeContext runtimeContext,
    @Label("keyed-state-backend") KeyedStateBackend<Object> keyedStateBackend,
    DynamicallyRegisteredTypes types) {

  this.runtimeContext = Objects.requireNonNull(runtimeContext);
  this.keyedStateBackend = Objects.requireNonNull(keyedStateBackend);
  this.types = Objects.requireNonNull(types);
}
 

private InternalListState<String, Long, Message> delayedMessagesBufferState(
    ListStateDescriptor<Message> delayedMessageStateDescriptor) {
  try {
    KeyedStateBackend<String> keyedStateBackend = getKeyedStateBackend();
    return (InternalListState<String, Long, Message>)
        keyedStateBackend.getOrCreateKeyedState(
            LongSerializer.INSTANCE, delayedMessageStateDescriptor);
  } catch (Exception e) {
    throw new RuntimeException("Error registered Flink state for delayed messages buffer.", e);
  }
}
 

@Inject
public FlinkState(
    @Label("runtime-context") RuntimeContext runtimeContext,
    @Label("keyed-state-backend") KeyedStateBackend<Object> keyedStateBackend,
    DynamicallyRegisteredTypes types) {

  this.runtimeContext = Objects.requireNonNull(runtimeContext);
  this.keyedStateBackend = Objects.requireNonNull(keyedStateBackend);
  this.types = Objects.requireNonNull(types);
}
 

@Inject
public MultiplexedState(
    @Label("runtime-context") RuntimeContext runtimeContext,
    @Label("keyed-state-backend") KeyedStateBackend<Object> keyedStateBackend,
    DynamicallyRegisteredTypes types) {

  this.keyedStateBackend = Objects.requireNonNull(keyedStateBackend);
  this.types = Objects.requireNonNull(types);
  this.sharedMapStateHandle = createSharedMapState(runtimeContext);
  this.executionConfiguration = Objects.requireNonNull(runtimeContext.getExecutionConfig());
}
 

public <N> int numKeyedStateEntries(N namespace) {
	AbstractStreamOperator<?> abstractStreamOperator = (AbstractStreamOperator<?>) operator;
	KeyedStateBackend<Object> keyedStateBackend = abstractStreamOperator.getKeyedStateBackend();
	if (keyedStateBackend instanceof HeapKeyedStateBackend) {
		return ((HeapKeyedStateBackend) keyedStateBackend).numKeyValueStateEntries(namespace);
	} else {
		throw new UnsupportedOperationException();
	}
}
 

public int numKeyedStateEntries() {
	AbstractStreamOperator<?> abstractStreamOperator = (AbstractStreamOperator<?>) operator;
	KeyedStateBackend<Object> keyedStateBackend = abstractStreamOperator.getKeyedStateBackend();
	if (keyedStateBackend instanceof HeapKeyedStateBackend) {
		return ((HeapKeyedStateBackend) keyedStateBackend).numKeyValueStateEntries();
	} else {
		throw new UnsupportedOperationException();
	}
}
 

private TtlStateFactory(
	@Nonnull TypeSerializer<N> namespaceSerializer,
	@Nonnull StateDescriptor<S, SV> stateDesc,
	@Nonnull KeyedStateBackend<K> stateBackend,
	@Nonnull TtlTimeProvider timeProvider) {
	this.namespaceSerializer = namespaceSerializer;
	this.stateDesc = stateDesc;
	this.stateBackend = stateBackend;
	this.ttlConfig = stateDesc.getTtlConfig();
	this.timeProvider = timeProvider;
	this.ttl = ttlConfig.getTtl().toMilliseconds();
	this.stateFactories = createStateFactories();
	this.incrementalCleanup = getTtlIncrementalCleanup();
}