Java Code Examples for org.apache.flink.runtime.state.KeyedStateBackend#setCurrentKey()

@Test
public void testApplyToAllKeys() 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);
	}
	applyToAllKeys(
		backend,
		listStateDescriptor,
		(k, state) -> {
			backend.setCurrentKey(k);
			state.clear();
		});
	for (long i = 0; i < 10; i++) {
		backend.setCurrentKey(i);
		Assert.assertNull(listState.get());
	}
	cleanUp(backend);
}
 

/**
 * Processes all pending processing timers. This is intended for use during shutdown. From Flink
 * 1.10 on, processing timer execution is stopped when the operator is closed. This leads to
 * problems for applications which assume all pending timers will be completed. Although Flink
 * does drain the remaining timers after close(), this is not sufficient because no new timers
 * are allowed to be scheduled anymore. This breaks Beam pipelines which rely on all processing
 * timers to be scheduled and executed.
 */
void processPendingProcessingTimeTimers() {
  final KeyedStateBackend<Object> keyedStateBackend = getKeyedStateBackend();
  final InternalPriorityQueue<InternalTimer<Object, TimerData>> processingTimeTimersQueue =
      Workarounds.retrieveInternalProcessingTimerQueue(timerService);

  InternalTimer<Object, TimerData> internalTimer;
  while ((internalTimer = processingTimeTimersQueue.poll()) != null) {
    keyedStateBackend.setCurrentKey(internalTimer.getKey());
    TimerData timer = internalTimer.getNamespace();
    checkInvokeStartBundle();
    fireTimer(timer);
  }
}
 

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);
    }
  }
}
 

@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);
}
 

@Test
public void testSetTtlTimeProvider() throws Exception {
	AbstractStreamOperator<Integer> operator = new AbstractStreamOperator<Integer>() {};
	try (AbstractStreamOperatorTestHarness<Integer> result = new AbstractStreamOperatorTestHarness<>(
			operator,
			1,
			1,
			0)) {

		result.config.setStateKeySerializer(IntSerializer.INSTANCE);

		Time timeToLive = Time.hours(1);
		result.initializeState(new OperatorSubtaskState());
		result.open();

		ValueStateDescriptor<Integer> stateDescriptor = new ValueStateDescriptor<>("test", IntSerializer.INSTANCE);
		stateDescriptor.enableTimeToLive(StateTtlConfig.newBuilder(timeToLive).build());
		KeyedStateBackend<Integer> keyedStateBackend = operator.getKeyedStateBackend();
		ValueState<Integer> state = keyedStateBackend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, stateDescriptor);

		int expectedValue = 42;
		keyedStateBackend.setCurrentKey(1);
		result.setStateTtlProcessingTime(0L);
		state.update(expectedValue);
		Assert.assertEquals(expectedValue, (int) state.value());
		result.setStateTtlProcessingTime(timeToLive.toMilliseconds() + 1);
		Assert.assertNull(state.value());
	}
}
 

@Test
public void testWatermarkHoldsPersistence() throws Exception {
  KeyedStateBackend<ByteBuffer> keyedStateBackend = createStateBackend();
  FlinkStateInternals stateInternals =
      new FlinkStateInternals<>(keyedStateBackend, StringUtf8Coder.of());

  StateTag<WatermarkHoldState> stateTag =
      StateTags.watermarkStateInternal("hold", TimestampCombiner.EARLIEST);
  WatermarkHoldState globalWindow = stateInternals.state(StateNamespaces.global(), stateTag);
  WatermarkHoldState fixedWindow =
      stateInternals.state(
          StateNamespaces.window(
              IntervalWindow.getCoder(), new IntervalWindow(new Instant(0), new Instant(10))),
          stateTag);

  Instant noHold = new Instant(Long.MAX_VALUE);
  assertThat(stateInternals.minWatermarkHoldMs(), is(noHold.getMillis()));

  Instant high = new Instant(10);
  globalWindow.add(high);
  assertThat(stateInternals.minWatermarkHoldMs(), is(high.getMillis()));

  Instant middle = new Instant(5);
  fixedWindow.add(middle);
  assertThat(stateInternals.minWatermarkHoldMs(), is(middle.getMillis()));

  Instant low = new Instant(1);
  globalWindow.add(low);
  assertThat(stateInternals.minWatermarkHoldMs(), is(low.getMillis()));

  // Try to overwrite with later hold (should not succeed)
  globalWindow.add(high);
  assertThat(stateInternals.minWatermarkHoldMs(), is(low.getMillis()));
  fixedWindow.add(high);
  assertThat(stateInternals.minWatermarkHoldMs(), is(low.getMillis()));

  // Watermark hold should be computed across all keys
  ByteBuffer firstKey = keyedStateBackend.getCurrentKey();
  changeKey(keyedStateBackend);
  ByteBuffer secondKey = keyedStateBackend.getCurrentKey();
  assertThat(firstKey, is(Matchers.not(secondKey)));
  assertThat(stateInternals.minWatermarkHoldMs(), is(low.getMillis()));
  // ..but be tracked per key / window
  assertThat(globalWindow.read(), is(Matchers.nullValue()));
  assertThat(fixedWindow.read(), is(Matchers.nullValue()));
  globalWindow.add(middle);
  fixedWindow.add(high);
  assertThat(globalWindow.read(), is(middle));
  assertThat(fixedWindow.read(), is(high));
  // Old key should give previous results
  keyedStateBackend.setCurrentKey(firstKey);
  assertThat(globalWindow.read(), is(low));
  assertThat(fixedWindow.read(), is(middle));

  // Discard watermark view and recover it
  stateInternals = new FlinkStateInternals<>(keyedStateBackend, StringUtf8Coder.of());
  globalWindow = stateInternals.state(StateNamespaces.global(), stateTag);
  fixedWindow =
      stateInternals.state(
          StateNamespaces.window(
              IntervalWindow.getCoder(), new IntervalWindow(new Instant(0), new Instant(10))),
          stateTag);

  // Watermark hold across all keys should be unchanged
  assertThat(stateInternals.minWatermarkHoldMs(), is(low.getMillis()));

  // Check the holds for the second key and clear them
  keyedStateBackend.setCurrentKey(secondKey);
  assertThat(globalWindow.read(), is(middle));
  assertThat(fixedWindow.read(), is(high));
  globalWindow.clear();
  fixedWindow.clear();

  // Check the holds for the first key and clear them
  keyedStateBackend.setCurrentKey(firstKey);
  assertThat(globalWindow.read(), is(low));
  assertThat(fixedWindow.read(), is(middle));

  fixedWindow.clear();
  assertThat(stateInternals.minWatermarkHoldMs(), is(low.getMillis()));

  globalWindow.clear();
  assertThat(stateInternals.minWatermarkHoldMs(), is(noHold.getMillis()));
}
 

private static void changeKey(KeyedStateBackend<ByteBuffer> keyedStateBackend)
    throws CoderException {
  keyedStateBackend.setCurrentKey(
      ByteBuffer.wrap(
          CoderUtils.encodeToByteArray(StringUtf8Coder.of(), UUID.randomUUID().toString())));
}