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

@Override
public void open(Configuration parameters) {
	MapStateDescriptor<EmailId, EmailInformation> stateDescriptor =
			new MapStateDescriptor<>(
					QsConstants.STATE_NAME,
					TypeInformation.of(new TypeHint<EmailId>() {

					}),
					TypeInformation.of(new TypeHint<EmailInformation>() {

					})
			);
	stateDescriptor.setQueryable(QsConstants.QUERY_NAME);
	state = getRuntimeContext().getMapState(stateDescriptor);
	count = -1;
}
 

@Test
public void testMapStateInstantiation() throws Exception {

	final ExecutionConfig config = new ExecutionConfig();
	config.registerKryoType(Path.class);

	final AtomicReference<Object> descriptorCapture = new AtomicReference<>();

	StreamingRuntimeContext context = new StreamingRuntimeContext(
			createDescriptorCapturingMockOp(descriptorCapture, config),
			createMockEnvironment(),
			Collections.<String, Accumulator<?, ?>>emptyMap());

	MapStateDescriptor<String, TaskInfo> descr =
			new MapStateDescriptor<>("name", String.class, TaskInfo.class);

	context.getMapState(descr);

	MapStateDescriptor<?, ?> descrIntercepted = (MapStateDescriptor<?, ?>) descriptorCapture.get();
	TypeSerializer<?> valueSerializer = descrIntercepted.getValueSerializer();

	// check that the Path class is really registered, i.e., the execution config was applied
	assertTrue(valueSerializer instanceof KryoSerializer);
	assertTrue(((KryoSerializer<?>) valueSerializer).getKryo().getRegistration(Path.class).getId() > 0);
}
 

@Override
public void open() throws Exception {
	super.open();

	LOG.info("Opening RowTimeSortOperator");
	if (gComparator != null) {
		comparator = gComparator.newInstance(getContainingTask().getUserCodeClassLoader());
		gComparator = null;
	}

	BasicTypeInfo<Long> keyTypeInfo = BasicTypeInfo.LONG_TYPE_INFO;
	ListTypeInfo<RowData> valueTypeInfo = new ListTypeInfo<>(inputRowType);
	MapStateDescriptor<Long, List<RowData>> mapStateDescriptor = new MapStateDescriptor<>(
			"dataState", keyTypeInfo, valueTypeInfo);
	dataState = getRuntimeContext().getMapState(mapStateDescriptor);

	ValueStateDescriptor<Long> lastTriggeringTsDescriptor = new ValueStateDescriptor<>("lastTriggeringTsState",
			Long.class);
	lastTriggeringTsState = getRuntimeContext().getState(lastTriggeringTsDescriptor);
}
 

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

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

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

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

	state.put("Ciao", "Hello");
	state.put("Bello", "Nice");

	assertNotNull(state.entries());
	assertEquals(state.get("Ciao"), "Hello");
	assertEquals(state.get("Bello"), "Nice");

	state.clear();
	assertNull(state.entries());

	backend.dispose();
}
 

@Override
public void initializeState(StateInitializationContext context) throws Exception {
	super.initializeState(context);

	// initializeState through the provided context
	computationStates = context.getKeyedStateStore().getState(
		new ValueStateDescriptor<>(
			NFA_STATE_NAME,
			new NFAStateSerializer()));

	partialMatches = new SharedBuffer<>(context.getKeyedStateStore(), inputSerializer);

	elementQueueState = context.getKeyedStateStore().getMapState(
			new MapStateDescriptor<>(
					EVENT_QUEUE_STATE_NAME,
					LongSerializer.INSTANCE,
					new ListSerializer<>(inputSerializer)));

	migrateOldState();
}
 

public SharedBuffer(KeyedStateStore stateStore, TypeSerializer<V> valueSerializer) {
	this.eventsBuffer = stateStore.getMapState(
		new MapStateDescriptor<>(
			eventsStateName,
			EventId.EventIdSerializer.INSTANCE,
			new Lockable.LockableTypeSerializer<>(valueSerializer)));

	this.entries = stateStore.getMapState(
		new MapStateDescriptor<>(
			entriesStateName,
			new NodeId.NodeIdSerializer(),
			new Lockable.LockableTypeSerializer<>(new SharedBufferNode.SharedBufferNodeSerializer())));

	this.eventsCount = stateStore.getMapState(
		new MapStateDescriptor<>(
			eventsCountStateName,
			LongSerializer.INSTANCE,
			IntSerializer.INSTANCE));
}
 

@Override
public <KeyT, ValueT> org.apache.beam.sdk.state.MapState<KeyT, ValueT> bindMap(
    String id,
    StateSpec<org.apache.beam.sdk.state.MapState<KeyT, ValueT>> spec,
    Coder<KeyT> mapKeyCoder,
    Coder<ValueT> mapValueCoder) {
  try {
    keyedStateBackend.getOrCreateKeyedState(
        StringSerializer.INSTANCE,
        new MapStateDescriptor<>(
            id,
            new CoderTypeSerializer<>(mapKeyCoder),
            new CoderTypeSerializer<>(mapValueCoder)));
  } catch (Exception e) {
    throw new RuntimeException(e);
  }
  return null;
}
 

@Override
public void open(Configuration parameters) throws Exception {
	function = genAggsHandler.newInstance(getRuntimeContext().getUserCodeClassLoader());
	function.open(new PerKeyStateDataViewStore(getRuntimeContext()));

	output = new JoinedRowData();

	sortedTimestamps = new LinkedList<Long>();

	// initialize accumulator state
	RowDataTypeInfo accTypeInfo = new RowDataTypeInfo(accTypes);
	ValueStateDescriptor<RowData> accStateDesc =
		new ValueStateDescriptor<RowData>("accState", accTypeInfo);
	accState = getRuntimeContext().getState(accStateDesc);

	// input element are all binary row as they are came from network
	RowDataTypeInfo inputType = new RowDataTypeInfo(inputFieldTypes);
	ListTypeInfo<RowData> rowListTypeInfo = new ListTypeInfo<RowData>(inputType);
	MapStateDescriptor<Long, List<RowData>> inputStateDesc = new MapStateDescriptor<Long, List<RowData>>(
		"inputState",
		Types.LONG,
		rowListTypeInfo);
	inputState = getRuntimeContext().getMapState(inputStateDesc);

	initCleanupTimeState("RowTimeUnboundedOverCleanupTime");
}
 

@Test
public void testKeyedMapStateRegistrationFailsIfNewStateSerializerIsIncompatible() {
	final String stateName = "test-name";

	try {
		testKeyedMapStateUpgrade(
			new MapStateDescriptor<>(
			stateName,
			IntSerializer.INSTANCE,
			new TestType.V1TestTypeSerializer()),
			new MapStateDescriptor<>(
			stateName,
			IntSerializer.INSTANCE,
			// restore with a V2 serializer that has a different schema
			new TestType.IncompatibleTestTypeSerializer()));
		fail("should have failed");
	} catch (Exception expected) {
		Assert.assertTrue(ExceptionUtils.findThrowable(expected, StateMigrationException.class).isPresent());
	}
}
 

@Test
public void testKeyedMapStateRegistrationFailsIfNewStateSerializerIsIncompatible() {
	final String stateName = "test-name";

	try {
		testKeyedMapStateUpgrade(
			new MapStateDescriptor<>(
			stateName,
			IntSerializer.INSTANCE,
			new TestType.V1TestTypeSerializer()),
			new MapStateDescriptor<>(
			stateName,
			IntSerializer.INSTANCE,
			// restore with a V2 serializer that has a different schema
			new TestType.IncompatibleTestTypeSerializer()));
		fail("should have failed");
	} catch (Exception expected) {
		Assert.assertTrue(ExceptionUtils.findThrowable(expected, StateMigrationException.class).isPresent());
	}
}
 

@SuppressWarnings("unchecked")
private <SV, SEV> StateSnapshotTransformer<SV> getStateSnapshotTransformer(
	StateDescriptor<?, SV> stateDesc,
	StateSnapshotTransformFactory<SEV> snapshotTransformFactory) {
	Optional<StateSnapshotTransformer<SEV>> original = snapshotTransformFactory.createForDeserializedState();
	if (original.isPresent()) {
		if (stateDesc instanceof ListStateDescriptor) {
			return (StateSnapshotTransformer<SV>) new StateSnapshotTransformers.ListStateSnapshotTransformer<>(original.get());
		} else if (stateDesc instanceof MapStateDescriptor) {
			return (StateSnapshotTransformer<SV>) new StateSnapshotTransformers.MapStateSnapshotTransformer<>(original.get());
		} else {
			return (StateSnapshotTransformer<SV>) original.get();
		}
	} else {
		return null;
	}
}
 

@Override
public void open() throws Exception {
	joinCondition = generatedJoinCondition.newInstance(getRuntimeContext().getUserCodeClassLoader());
	joinCondition.setRuntimeContext(getRuntimeContext());
	joinCondition.open(new Configuration());

	nextLeftIndex = getRuntimeContext().getState(
		new ValueStateDescriptor<>(NEXT_LEFT_INDEX_STATE_NAME, Types.LONG));
	leftState = getRuntimeContext().getMapState(
		new MapStateDescriptor<>(LEFT_STATE_NAME, Types.LONG, leftType));
	rightState = getRuntimeContext().getMapState(
		new MapStateDescriptor<>(RIGHT_STATE_NAME, Types.LONG, rightType));
	registeredTimer = getRuntimeContext().getState(
		new ValueStateDescriptor<>(REGISTERED_TIMER_STATE_NAME, Types.LONG));

	timerService = getInternalTimerService(
		TIMERS_STATE_NAME, VoidNamespaceSerializer.INSTANCE, this);
	collector = new TimestampedCollector<>(output);
	outRow = new JoinedRow();
	outRow.setHeader(BaseRowUtil.ACCUMULATE_MSG);
}
 

@Override
public void open(Configuration parameters) throws Exception {
	super.open(parameters);

	descriptor = new MapStateDescriptor<>(
		"broadcast-state", BasicTypeInfo.LONG_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO
	);
}
 

@Override
public void open(Context<K, W> ctx) throws Exception {
	super.open(ctx);
	MapStateDescriptor<W, W> mappingStateDescriptor = new MapStateDescriptor<>(
		"session-window-mapping",
		windowSerializer,
		windowSerializer);
	MapState<W, W> windowMapping = ctx.getPartitionedState(mappingStateDescriptor);
	this.mergingWindows = new MergingWindowSet<>(windowAssigner, windowMapping);
	this.mergingFunction = new MergingFunctionImpl();
}
 

@Override
public <K, V> ReadOnlyBroadcastState<K, V> getBroadcastState(MapStateDescriptor<K, V> stateDescriptor) {
	Preconditions.checkNotNull(stateDescriptor);

	stateDescriptor.initializeSerializerUnlessSet(config);
	ReadOnlyBroadcastState<K, V> state = (ReadOnlyBroadcastState<K, V>) states.get(stateDescriptor);
	if (state == null) {
		throw new IllegalArgumentException("The requested state does not exist. " +
				"Check for typos in your state descriptor, or specify the state descriptor " +
				"in the datastream.broadcast(...) call if you forgot to register it.");
	}
	return state;
}
 

@SuppressWarnings("unchecked")
public static <K, V, T, S extends State> S createState(
	StateDescriptor<S, T> stateDescriptor,
	byte[] serializedState) throws IOException {
	MapStateDescriptor<K, V> mapStateDescriptor = (MapStateDescriptor<K, V>) stateDescriptor;
	final Map<K, V> state = KvStateSerializer.deserializeMap(
		serializedState,
		mapStateDescriptor.getKeySerializer(),
		mapStateDescriptor.getValueSerializer());
	return (S) new ImmutableMapState<>(state);
}
 

/**
 * Read operator {@code BroadcastState} from a {@code Savepoint}.
 * @param uid The uid of the operator.
 * @param name The (unique) name for the state.
 * @param keyTypeInfo The type information for the keys in the state.
 * @param valueTypeInfo The type information for the values in the state.
 * @param <K> The type of keys in state.
 * @param <V> The type of values in state.
 * @return A {@code DataSet} of key-value pairs from state.
 * @throws IOException If the savepoint does not contain the specified uid.
 */
public <K, V> DataSet<Tuple2<K, V>> readBroadcastState(
	String uid,
	String name,
	TypeInformation<K> keyTypeInfo,
	TypeInformation<V> valueTypeInfo) throws IOException {

	OperatorState operatorState = metadata.getOperatorState(uid);
	MapStateDescriptor<K, V> descriptor = new MapStateDescriptor<>(name, keyTypeInfo, valueTypeInfo);
	BroadcastStateInputFormat<K, V> inputFormat = new BroadcastStateInputFormat<>(operatorState, descriptor);
	return env.createInput(inputFormat, new TupleTypeInfo<>(keyTypeInfo, valueTypeInfo));
}
 

OnTimerContextImpl(
		final ExecutionConfig executionConfig,
		final KeyedBroadcastProcessFunction<KS, IN1, IN2, OUT> function,
		final Map<MapStateDescriptor<?, ?>, BroadcastState<?, ?>> broadcastStates,
		final TimerService timerService) {

	function.super();
	this.config = Preconditions.checkNotNull(executionConfig);
	this.states = Preconditions.checkNotNull(broadcastStates);
	this.timerService = Preconditions.checkNotNull(timerService);
}
 

@Override
public <K, V> BroadcastState<K, V> getBroadcastState(MapStateDescriptor<K, V> stateDescriptor) {
	Preconditions.checkNotNull(stateDescriptor);

	stateDescriptor.initializeSerializerUnlessSet(config);
	BroadcastState<K, V> state = (BroadcastState<K, V>) states.get(stateDescriptor);
	if (state == null) {
		throw new IllegalArgumentException("The requested state does not exist. " +
				"Check for typos in your state descriptor, or specify the state descriptor " +
				"in the datastream.broadcast(...) call if you forgot to register it.");
	}
	return state;
}
 

private InputSideHasNoUniqueKey(
		RuntimeContext ctx,
		String stateName,
		RowDataTypeInfo recordType,
		StateTtlConfig ttlConfig) {
	MapStateDescriptor<RowData, Integer> recordStateDesc = new MapStateDescriptor<>(
		stateName,
		recordType,
		Types.INT);
	if (ttlConfig.isEnabled()) {
		recordStateDesc.enableTimeToLive(ttlConfig);
	}
	this.recordState = ctx.getMapState(recordStateDesc);
}
 

ReadOnlyContextImpl(
		final ExecutionConfig executionConfig,
		final KeyedBroadcastProcessFunction<KS, IN1, IN2, OUT> function,
		final Map<MapStateDescriptor<?, ?>, BroadcastState<?, ?>> broadcastStates,
		final TimerService timerService) {

	function.super();
	this.config = Preconditions.checkNotNull(executionConfig);
	this.states = Preconditions.checkNotNull(broadcastStates);
	this.timerService = Preconditions.checkNotNull(timerService);
}
 

@Override
public <K, V> BroadcastState<K, V> getBroadcastState(MapStateDescriptor<K, V> stateDescriptor) {
	Preconditions.checkNotNull(stateDescriptor);

	stateDescriptor.initializeSerializerUnlessSet(config);
	BroadcastState<K, V> state = (BroadcastState<K, V>) states.get(stateDescriptor);
	if (state == null) {
		throw new IllegalArgumentException("The requested state does not exist. " +
				"Check for typos in your state descriptor, or specify the state descriptor " +
				"in the datastream.broadcast(...) call if you forgot to register it.");
	}
	return state;
}
 

@Override
public <K, V> ReadOnlyBroadcastState<K, V> getBroadcastState(MapStateDescriptor<K, V> stateDescriptor) {
	Preconditions.checkNotNull(stateDescriptor);

	stateDescriptor.initializeSerializerUnlessSet(config);
	ReadOnlyBroadcastState<K, V> state = (ReadOnlyBroadcastState<K, V>) states.get(stateDescriptor);
	if (state == null) {
		throw new IllegalArgumentException("The requested state does not exist. " +
				"Check for typos in your state descriptor, or specify the state descriptor " +
				"in the datastream.broadcast(...) call if you forgot to register it.");
	}
	return state;
}
 

private void testBroadcastStateValueUpgrade(
		MapStateDescriptor<Integer, TestType> initialAccessDescriptor,
		MapStateDescriptor<Integer, TestType> newAccessDescriptorAfterRestore) throws Exception {
	CheckpointStreamFactory streamFactory = createStreamFactory();

	OperatorStateBackend backend = createOperatorStateBackend();

	try {
		BroadcastState<Integer, TestType> state = backend.getBroadcastState(initialAccessDescriptor);

		state.put(3, new TestType("foo", 13));
		state.put(5, new TestType("bar", 278));

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

		backend = restoreOperatorStateBackend(snapshot);

		state = backend.getBroadcastState(newAccessDescriptorAfterRestore);

		// the state backend should have decided whether or not it needs to perform state migration;
		// make sure that reading and writing each broadcast entry works with the new serializer
		Assert.assertEquals(new TestType("foo", 13), state.get(3));
		Assert.assertEquals(new TestType("bar", 278), state.get(5));
		state.put(17, new TestType("new-entry", 777));
	} finally {
		backend.dispose();
	}
}
 

private static <IN1, IN2, OUT> TwoInputStreamOperatorTestHarness<IN1, IN2, OUT> getInitializedTestHarness(
		final BroadcastProcessFunction<IN1, IN2, OUT> function,
		final MapStateDescriptor<?, ?>... descriptors) throws Exception {

	return getInitializedTestHarness(
			function,
			1,
			1,
			0,
			descriptors);
}
 

@Override
public <UK, UV> MapState<UK, UV> getMapState(MapStateDescriptor<UK, UV> stateProperties) {
	requireNonNull(stateProperties, "The state properties must not be null");
	try {
		stateProperties.initializeSerializerUnlessSet(executionConfig);
		MapState<UK, UV> originalState = getPartitionedState(stateProperties);
		return new UserFacingMapState<>(originalState);
	} catch (Exception e) {
		throw new RuntimeException("Error while getting state", e);
	}
}
 

@Override
public void open(Configuration parameters) throws Exception {
	super.open(parameters);

	firstStateDesc = new MapStateDescriptor<>(
			"broadcast-state-1", BasicTypeInfo.LONG_TYPE_INFO, BasicTypeInfo.LONG_TYPE_INFO
	);

	secondStateDesc = new MapStateDescriptor<>(
			"broadcast-state-2", BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.LONG_TYPE_INFO
	);
}
 

private static <IN1, IN2, OUT> TwoInputStreamOperatorTestHarness<IN1, IN2, OUT> getInitializedTestHarness(
		final BroadcastProcessFunction<IN1, IN2, OUT> function,
		final MapStateDescriptor<?, ?>... descriptors) throws Exception {

	return getInitializedTestHarness(
			function,
			1,
			1,
			0,
			descriptors);
}
 

@Test
public void testMapStateReturnsEmptyMapByDefault() throws Exception {

	StreamingRuntimeContext context = new StreamingRuntimeContext(
			createMapPlainMockOp(),
			createMockEnvironment(),
			Collections.<String, Accumulator<?, ?>>emptyMap());

	MapStateDescriptor<Integer, String> descr = new MapStateDescriptor<>("name", Integer.class, String.class);
	MapState<Integer, String> state = context.getMapState(descr);

	Iterable<Map.Entry<Integer, String>> value = state.entries();
	assertNotNull(value);
	assertFalse(value.iterator().hasNext());
}
 

/**
 * Tests to cover case when user misuse optimizeForPointLookup with iterator interfaces on map state.
 *
 * <p>The option {@link ColumnFamilyOptions#optimizeForPointLookup(long)} would lead to iterator.seek with prefix bytes invalid.
 */
@Test
public void testMisuseOptimizePointLookupWithMapState() throws Exception {
	RocksDBStateBackend rocksDBStateBackend = createStateBackendWithOptimizePointLookup();
	RocksDBKeyedStateBackend<Integer> keyedStateBackend =
		createKeyedStateBackend(rocksDBStateBackend, new MockEnvironmentBuilder().build(), IntSerializer.INSTANCE);
	try {
		MapStateDescriptor<Integer, Long> stateDescriptor = new MapStateDescriptor<>("map", IntSerializer.INSTANCE, LongSerializer.INSTANCE);
		MapState<Integer, Long> mapState = keyedStateBackend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, stateDescriptor);

		keyedStateBackend.setCurrentKey(1);
		Map<Integer, Long> expectedResult = new HashMap<>();
		for (int i = 0; i < 100; i++) {
			long uv = ThreadLocalRandom.current().nextLong();
			mapState.put(i, uv);
			expectedResult.put(i, uv);
		}

		Iterator<Map.Entry<Integer, Long>> iterator = mapState.entries().iterator();
		while (iterator.hasNext()) {
			Map.Entry<Integer, Long> entry = iterator.next();
			assertEquals(entry.getValue(), expectedResult.remove(entry.getKey()));
			iterator.remove();
		}
		assertTrue(expectedResult.isEmpty());
		assertTrue(mapState.isEmpty());
	} finally {
		keyedStateBackend.dispose();
	}
}