org.apache.flink.api.common.state.ListState Java Examples

@Override
public void processBroadcastElement(Integer value, Context ctx, Collector<String> out) throws Exception {
	// put an element in the broadcast state
	ctx.applyToKeyedState(
			listStateDesc,
			new KeyedStateFunction<String, ListState<String>>() {
				@Override
				public void process(String key, ListState<String> state) throws Exception {
					final Iterator<String> it = state.get().iterator();

					final List<String> list = new ArrayList<>();
					while (it.hasNext()) {
						list.add(it.next());
					}
					assertEquals(expectedKeyedStates.get(key), list);
				}
			});
}
 

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

/**
 * Restores a {@link MergingWindowSet} from the given state.
 */
public MergingWindowSet(MergingWindowAssigner<?, W> windowAssigner, ListState<Tuple2<W, W>> state) throws Exception {
	this.windowAssigner = windowAssigner;
	mapping = new HashMap<>();

	Iterable<Tuple2<W, W>> windowState = state.get();
	if (windowState != null) {
		for (Tuple2<W, W> window: windowState) {
			mapping.put(window.f0, window.f1);
		}
	}

	this.state = state;

	initialMapping = new HashMap<>();
	initialMapping.putAll(mapping);
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	ListState<String> unionListState = context.getOperatorStateStore().getListState(
			CheckpointingNonParallelSourceWithListState.STATE_DESCRIPTOR);

	if (context.isRestored()) {
		assertThat(unionListState.get(),
				containsInAnyOrder(
						CheckpointingNonParallelSourceWithListState.CHECKPOINTED_STRING,
						CheckpointingNonParallelSourceWithListState.CHECKPOINTED_STRING_1,
						CheckpointingNonParallelSourceWithListState.CHECKPOINTED_STRING_2,
						CheckpointingNonParallelSourceWithListState.CHECKPOINTED_STRING_3));

		getRuntimeContext().addAccumulator(SUCCESSFUL_RESTORE_CHECK_ACCUMULATOR, new IntCounter());
		getRuntimeContext().getAccumulator(SUCCESSFUL_RESTORE_CHECK_ACCUMULATOR).add(1);
	} else {
		throw new RuntimeException(
				"This source should always be restored because it's only used when restoring from a savepoint.");
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	ListState<String> unionListState = context.getOperatorStateStore().getUnionListState(
			CheckpointingNonParallelSourceWithListState.STATE_DESCRIPTOR);

	if (context.isRestored()) {
		assertThat(unionListState.get(),
				containsInAnyOrder(CheckpointingParallelSourceWithUnionListState.CHECKPOINTED_STRINGS));

		getRuntimeContext().addAccumulator(SUCCESSFUL_RESTORE_CHECK_ACCUMULATOR, new IntCounter());
		getRuntimeContext().getAccumulator(SUCCESSFUL_RESTORE_CHECK_ACCUMULATOR).add(1);
	} else {
		throw new RuntimeException(
				"This source should always be restored because it's only used when restoring from a savepoint.");
	}
}
 

/**
 * Verify that an empty {@code ListState} yields {@code null}.
 */
@Test
public void testListStateDefaultValue() throws Exception {
	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);

	ListStateDescriptor<String> kvId = new ListStateDescriptor<>("id", String.class);

	ListState<String> state = backend.getPartitionedState(
			VoidNamespace.INSTANCE,
			VoidNamespaceSerializer.INSTANCE, kvId);

	backend.setCurrentKey(1);
	assertNull(state.get());

	state.update(Arrays.asList("Ciao", "Bello"));
	assertThat(state.get(), containsInAnyOrder("Ciao", "Bello"));

	state.clear();
	assertNull(state.get());

	backend.dispose();
}
 

@Override
public void add(T input) {
  try {
    ListState<T> partitionedState =
        flinkStateBackend.getPartitionedState(
            namespace.stringKey(), StringSerializer.INSTANCE, flinkStateDescriptor);
    if (storesVoidValues) {
      Preconditions.checkState(input == null, "Expected to a null value but was: %s", input);
      // Flink does not allow storing null values
      // If we have null values, we use the structural null value
      input = (T) VoidCoder.of().structuralValue((Void) input);
    }
    partitionedState.add(input);
  } catch (Exception e) {
    throw new RuntimeException("Error adding to bag state.", e);
  }
}
 

@Test
public void testPersistOnlyIfHaveUpdates() throws Exception {
	@SuppressWarnings("unchecked")
	ListState<Tuple2<TimeWindow, TimeWindow>> mockState = mock(ListState.class);
	when(mockState.get()).thenReturn(Lists.newArrayList(
			new Tuple2<>(new TimeWindow(17, 42), new TimeWindow(42, 17)),
			new Tuple2<>(new TimeWindow(1, 2), new TimeWindow(3, 4))
	));

	MergingWindowSet<TimeWindow> windowSet = new MergingWindowSet<>(EventTimeSessionWindows.withGap(Time.milliseconds(3)), mockState);

	assertEquals(new TimeWindow(42, 17), windowSet.getStateWindow(new TimeWindow(17, 42)));
	assertEquals(new TimeWindow(3, 4), windowSet.getStateWindow(new TimeWindow(1, 2)));

	windowSet.persist();

	verify(mockState, times(0)).add(Matchers.<Tuple2<TimeWindow, TimeWindow>>anyObject());

}
 

private static BufferingDoFnRunner createBufferingDoFnRunner(
    int concurrentCheckpoints,
    List<BufferingDoFnRunner.CheckpointIdentifier> notYetAcknowledgeCheckpoints)
    throws Exception {
  DoFnRunner doFnRunner = Mockito.mock(DoFnRunner.class);
  OperatorStateBackend operatorStateBackend = Mockito.mock(OperatorStateBackend.class);

  // Setup not yet acknowledged checkpoint union list state
  ListState unionListState = Mockito.mock(ListState.class);
  Mockito.when(operatorStateBackend.getUnionListState(Mockito.any())).thenReturn(unionListState);
  Mockito.when(unionListState.get()).thenReturn(notYetAcknowledgeCheckpoints);

  // Setup buffer list state
  Mockito.when(operatorStateBackend.getListState(Mockito.any()))
      .thenReturn(Mockito.mock(ListState.class));

  return BufferingDoFnRunner.create(
      doFnRunner,
      "stable-input",
      StringUtf8Coder.of(),
      WindowedValue.getFullCoder(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE),
      operatorStateBackend,
      null,
      concurrentCheckpoints);
}
 

private static Buckets<String, String> restoreBuckets(
		final Path basePath,
		final RollingPolicy<String, String> rollingPolicy,
		final BucketLifeCycleListener<String, String> bucketLifeCycleListener,
		final int subtaskIdx,
		final ListState<byte[]> bucketState,
		final ListState<Long> partCounterState,
		final OutputFileConfig outputFileConfig) throws Exception {
	final Buckets<String, String> restoredBuckets = createBuckets(
		basePath,
		rollingPolicy,
		bucketLifeCycleListener,
		subtaskIdx,
		outputFileConfig);
	restoredBuckets.initializeState(bucketState, partCounterState);
	return restoredBuckets;
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	ListState<String> unionListState = context.getOperatorStateStore().getUnionListState(
			CheckpointingNonParallelSourceWithListState.STATE_DESCRIPTOR);

	if (context.isRestored()) {
		assertThat(unionListState.get(),
				containsInAnyOrder(CheckpointingParallelSourceWithUnionListState.CHECKPOINTED_STRINGS));

		getRuntimeContext().addAccumulator(SUCCESSFUL_RESTORE_CHECK_ACCUMULATOR, new IntCounter());
		getRuntimeContext().getAccumulator(SUCCESSFUL_RESTORE_CHECK_ACCUMULATOR).add(1);
	} else {
		throw new RuntimeException(
				"This source should always be restored because it's only used when restoring from a savepoint.");
	}
}
 

/**
 * This test verifies that all ListState implementations are consistent in not allowing
 * {@link ListState#addAll(List)} to be called with {@code null}.
 */
@Test
public void testListStateUpdateNull() throws Exception {
	AbstractKeyedStateBackend<String> keyedBackend = createKeyedBackend(StringSerializer.INSTANCE);

	final ListStateDescriptor<Long> stateDescr = new ListStateDescriptor<>("my-state", Long.class);

	try {
		ListState<Long> state =
			keyedBackend.getPartitionedState(
				VoidNamespace.INSTANCE,
				VoidNamespaceSerializer.INSTANCE,
				stateDescr);

		keyedBackend.setCurrentKey("abc");
		assertNull(state.get());

		expectedException.expect(NullPointerException.class);
		state.update(null);
	} finally {
		keyedBackend.close();
		keyedBackend.dispose();
	}
}
 

@Test
public void testApplyToAllKeysLambdaFunction() throws Exception {
	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);

	try {
		ListStateDescriptor<String> listStateDescriptor =
			new ListStateDescriptor<>("foo", StringSerializer.INSTANCE);

		ListState<String> listState =
			backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, listStateDescriptor);

		for (int i = 0; i < 100; ++i) {
			backend.setCurrentKey(i);
			listState.add("Hello" + i);
		}

		// valid state value via applyToAllKeys().
		backend.applyToAllKeys(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, listStateDescriptor,
			(Integer key, ListState<String> state) -> assertEquals("Hello" + key, state.get().iterator().next())
		);
	}
	finally {
		IOUtils.closeQuietly(backend);
		backend.dispose();
	}
}
 

/**
 * This test verifies that all ListState implementations are consistent in not allowing
 * adding {@code null}.
 */
@Test
public void testListStateAddNull() throws Exception {
	AbstractKeyedStateBackend<String> keyedBackend = createKeyedBackend(StringSerializer.INSTANCE);

	final ListStateDescriptor<Long> stateDescr = new ListStateDescriptor<>("my-state", Long.class);

	try {
		ListState<Long> state =
			keyedBackend.getPartitionedState(
				VoidNamespace.INSTANCE,
				VoidNamespaceSerializer.INSTANCE,
				stateDescr);

		keyedBackend.setCurrentKey("abc");
		assertNull(state.get());

		expectedException.expect(NullPointerException.class);
		state.add(null);
	} finally {
		keyedBackend.close();
		keyedBackend.dispose();
	}
}
 

/**
 * Restores a {@link MergingWindowSet} from the given state.
 */
public MergingWindowSet(MergingWindowAssigner<?, W> windowAssigner, ListState<Tuple2<W, W>> state) throws Exception {
	this.windowAssigner = windowAssigner;
	mapping = new HashMap<>();

	Iterable<Tuple2<W, W>> windowState = state.get();
	if (windowState != null) {
		for (Tuple2<W, W> window: windowState) {
			mapping.put(window.f0, window.f1);
		}
	}

	this.state = state;

	initialMapping = new HashMap<>();
	initialMapping.putAll(mapping);
}
 

/**
 * Test merging of a large new window that covers one existing windows.
 */
@Test
public void testMergeLargeWindowCoveringSingleWindow() throws Exception {
	@SuppressWarnings("unchecked")
	ListState<Tuple2<TimeWindow, TimeWindow>> mockState = mock(ListState.class);

	MergingWindowSet<TimeWindow> windowSet = new MergingWindowSet<>(EventTimeSessionWindows.withGap(Time.milliseconds(3)), mockState);

	TestingMergeFunction mergeFunction = new TestingMergeFunction();

	// add an initial small window

	mergeFunction.reset();
	assertEquals(new TimeWindow(1, 2), windowSet.addWindow(new TimeWindow(1, 2), mergeFunction));
	assertFalse(mergeFunction.hasMerged());
	assertEquals(new TimeWindow(1, 2), windowSet.getStateWindow(new TimeWindow(1, 2)));

	// add a new window that completely covers the existing window

	mergeFunction.reset();
	assertEquals(new TimeWindow(0, 3), windowSet.addWindow(new TimeWindow(0, 3), mergeFunction));
	assertTrue(mergeFunction.hasMerged());
	assertEquals(new TimeWindow(1, 2), windowSet.getStateWindow(new TimeWindow(0, 3)));
}
 

@Override
public void snapshotState(StateSnapshotContext context) throws Exception {
	super.snapshotState(context);

	ListState<StreamElement> partitionableState =
		getOperatorStateBackend().getListState(new ListStateDescriptor<>(STATE_NAME, inStreamElementSerializer));
	partitionableState.clear();

	try {
		partitionableState.addAll(queue.values());
	} catch (Exception e) {
		partitionableState.clear();

		throw new Exception("Could not add stream element queue entries to operator state " +
			"backend of operator " + getOperatorName() + '.', e);
	}
}
 

public EvictingWindowOperator(WindowAssigner<? super IN, W> windowAssigner,
		TypeSerializer<W> windowSerializer,
		KeySelector<IN, K> keySelector,
		TypeSerializer<K> keySerializer,
		StateDescriptor<? extends ListState<StreamRecord<IN>>, ?> windowStateDescriptor,
		InternalWindowFunction<Iterable<IN>, OUT, K, W> windowFunction,
		Trigger<? super IN, ? super W> trigger,
		Evictor<? super IN, ? super W> evictor,
		long allowedLateness,
		OutputTag<IN> lateDataOutputTag) {

	super(windowAssigner, windowSerializer, keySelector,
		keySerializer, null, windowFunction, trigger, allowedLateness, lateDataOutputTag);

	this.evictor = checkNotNull(evictor);
	this.evictingWindowStateDescriptor = checkNotNull(windowStateDescriptor);
}
 

void snapshotState(
		final long checkpointId,
		final ListState<byte[]> bucketStatesContainer,
		final ListState<Long> partCounterStateContainer) throws Exception {

	Preconditions.checkState(
			fsWriter != null && bucketStateSerializer != null,
			"sink has not been initialized");

	LOG.info("Subtask {} checkpointing for checkpoint with id={} (max part counter={}).",
			subtaskIndex, checkpointId, maxPartCounter);

	bucketStatesContainer.clear();
	partCounterStateContainer.clear();

	snapshotActiveBuckets(checkpointId, bucketStatesContainer);
	partCounterStateContainer.add(maxPartCounter);
}
 

@Test
public void testApplyToAllKeysLambdaFunction() throws Exception {
	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);

	try {
		ListStateDescriptor<String> listStateDescriptor =
			new ListStateDescriptor<>("foo", StringSerializer.INSTANCE);

		ListState<String> listState =
			backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, listStateDescriptor);

		for (int i = 0; i < 100; ++i) {
			backend.setCurrentKey(i);
			listState.add("Hello" + i);
		}

		// valid state value via applyToAllKeys().
		backend.applyToAllKeys(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, listStateDescriptor,
			(Integer key, ListState<String> state) -> assertEquals("Hello" + key, state.get().iterator().next())
		);
	}
	finally {
		IOUtils.closeQuietly(backend);
		backend.dispose();
	}
}
 

void snapshotState(
		final long checkpointId,
		final ListState<byte[]> bucketStatesContainer,
		final ListState<Long> partCounterStateContainer) throws Exception {

	Preconditions.checkState(
			fsWriter != null && bucketStateSerializer != null,
			"sink has not been initialized");

	LOG.info("Subtask {} checkpointing for checkpoint with id={} (max part counter={}).",
			subtaskIndex, checkpointId, maxPartCounter);

	bucketStatesContainer.clear();
	partCounterStateContainer.clear();

	snapshotActiveBuckets(checkpointId, bucketStatesContainer);
	partCounterStateContainer.add(maxPartCounter);
}
 

/**
 * Test adding a new window that is identical to an existing window. This should not cause
 * a merge.
 */
@Test
public void testAddingIdenticalWindows() throws Exception {
	@SuppressWarnings("unchecked")
	ListState<Tuple2<TimeWindow, TimeWindow>> mockState = mock(ListState.class);

	MergingWindowSet<TimeWindow> windowSet = new MergingWindowSet<>(EventTimeSessionWindows.withGap(Time.milliseconds(3)), mockState);

	TestingMergeFunction mergeFunction = new TestingMergeFunction();

	mergeFunction.reset();
	assertEquals(new TimeWindow(1, 2), windowSet.addWindow(new TimeWindow(1, 2), mergeFunction));
	assertFalse(mergeFunction.hasMerged());
	assertEquals(new TimeWindow(1, 2), windowSet.getStateWindow(new TimeWindow(1, 2)));

	mergeFunction.reset();
	assertEquals(new TimeWindow(1, 2), windowSet.addWindow(new TimeWindow(1, 2), mergeFunction));
	assertFalse(mergeFunction.hasMerged());
	assertEquals(new TimeWindow(1, 2), windowSet.getStateWindow(new TimeWindow(1, 2)));
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {

	readyToFail = false;

	if (context.isRestored()) {
		isRunning = false;
	} else {
		isRunning = true;

		OperatorStateStore operatorStateStore = context.getOperatorStateStore();
		for (int i = 0; i < numListStates; ++i) {

			ListStateDescriptor<String> listStateDescriptor =
				new ListStateDescriptor<>("test-list-state-" + i, String.class);

			ListState<String> unionListState =
				operatorStateStore.getUnionListState(listStateDescriptor);

			for (int j = 0; j < numPartitionsPerListState; ++j) {
				unionListState.add(String.valueOf(j));
			}
		}
	}
}
 

/**
 * Test adding a new window that is identical to an existing window. This should not cause
 * a merge.
 */
@Test
public void testAddingIdenticalWindows() throws Exception {
	@SuppressWarnings("unchecked")
	ListState<Tuple2<TimeWindow, TimeWindow>> mockState = mock(ListState.class);

	MergingWindowSet<TimeWindow> windowSet = new MergingWindowSet<>(EventTimeSessionWindows.withGap(Time.milliseconds(3)), mockState);

	TestingMergeFunction mergeFunction = new TestingMergeFunction();

	mergeFunction.reset();
	assertEquals(new TimeWindow(1, 2), windowSet.addWindow(new TimeWindow(1, 2), mergeFunction));
	assertFalse(mergeFunction.hasMerged());
	assertEquals(new TimeWindow(1, 2), windowSet.getStateWindow(new TimeWindow(1, 2)));

	mergeFunction.reset();
	assertEquals(new TimeWindow(1, 2), windowSet.addWindow(new TimeWindow(1, 2), mergeFunction));
	assertFalse(mergeFunction.hasMerged());
	assertEquals(new TimeWindow(1, 2), windowSet.getStateWindow(new TimeWindow(1, 2)));
}
 

private static Buckets<String, String> restoreBuckets(
		final Path basePath,
		final RollingPolicy<String, String> rollingPolicy,
		final int subtaskIdx,
		final ListState<byte[]> bucketState,
		final ListState<Long> partCounterState
) throws Exception {

	final Buckets<String, String> restoredBuckets = createBuckets(basePath, rollingPolicy, subtaskIdx);
	restoredBuckets.initializeState(bucketState, partCounterState);
	return restoredBuckets;
}
 

@Override
public <N, V> StateListView<N, V> getStateListView(String stateName, ListViewTypeInfo<V> listViewTypeInfo) throws Exception {
	ListStateDescriptor<V> listStateDesc = new ListStateDescriptor<>(
		stateName,
		listViewTypeInfo.getElementType());

	ListState<V> listState = keyedStateBackend.getOrCreateKeyedState(windowSerializer, listStateDesc);
	// explict cast to internal state
	InternalListState<?, N, V> internalListState = (InternalListState<?, N, V>) listState;

	return new StateListView.NamespacedStateListView<>(internalListState);
}
 

private void testOperatorUnionListStateUpgrade(
		ListStateDescriptor<TestType> initialAccessDescriptor,
		ListStateDescriptor<TestType> newAccessDescriptorAfterRestore) throws Exception {

	CheckpointStreamFactory streamFactory = createStreamFactory();

	OperatorStateBackend backend = createOperatorStateBackend();

	try {
		ListState<TestType> state = backend.getUnionListState(initialAccessDescriptor);

		state.add(new TestType("foo", 13));
		state.add(new TestType("bar", 278));

		OperatorStateHandle snapshot = runSnapshot(
			backend.snapshot(1L, 2L, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation()));
		backend.dispose();

		backend = restoreOperatorStateBackend(snapshot);

		state = backend.getUnionListState(newAccessDescriptorAfterRestore);

		// the state backend should have decided whether or not it needs to perform state migration;
		// make sure that reading and writing each state partition works with the new serializer
		Iterator<TestType> iterator = state.get().iterator();
		assertEquals(new TestType("foo", 13), iterator.next());
		assertEquals(new TestType("bar", 278), iterator.next());
		Assert.assertFalse(iterator.hasNext());
		state.add(new TestType("new-entry", 777));
	} finally {
		backend.dispose();
	}
}
 

@Override
public <T> ListState<T> getListState(ListStateDescriptor<T> stateProperties) {
	if (!stateRegistrationAllowed) {
		throw new RuntimeException(REGISTRATION_EXCEPTION_MSG);
	}

	registeredDescriptors.add(stateProperties);
	return keyedStateStore.getListState(stateProperties);
}
 

@Test
public void testRestoreFromState() throws Exception {
	@SuppressWarnings("unchecked")
	ListState<Tuple2<TimeWindow, TimeWindow>> mockState = mock(ListState.class);
	when(mockState.get()).thenReturn(Lists.newArrayList(
			new Tuple2<>(new TimeWindow(17, 42), new TimeWindow(42, 17)),
			new Tuple2<>(new TimeWindow(1, 2), new TimeWindow(3, 4))
	));

	MergingWindowSet<TimeWindow> windowSet = new MergingWindowSet<>(EventTimeSessionWindows.withGap(Time.milliseconds(3)), mockState);

	assertEquals(new TimeWindow(42, 17), windowSet.getStateWindow(new TimeWindow(17, 42)));
	assertEquals(new TimeWindow(3, 4), windowSet.getStateWindow(new TimeWindow(1, 2)));
}
 

private static boolean tryRestoreFunction(
		StateInitializationContext context,
		Function userFunction) throws Exception {

	if (userFunction instanceof CheckpointedFunction) {
		((CheckpointedFunction) userFunction).initializeState(context);

		return true;
	}

	if (context.isRestored() && userFunction instanceof ListCheckpointed) {
		@SuppressWarnings("unchecked")
		ListCheckpointed<Serializable> listCheckpointedFun = (ListCheckpointed<Serializable>) userFunction;

		ListState<Serializable> listState = context.getOperatorStateStore().
				getSerializableListState(DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME);

		List<Serializable> list = new ArrayList<>();

		for (Serializable serializable : listState.get()) {
			list.add(serializable);
		}

		try {
			listCheckpointedFun.restoreState(list);
		} catch (Exception e) {

			throw new Exception("Failed to restore state to function: " + e.getMessage(), e);
		}

		return true;
	}

	return false;
}