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

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

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

/**
 * 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();
}
 

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

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

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

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

@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 <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() 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 <UK, UV> MapState<UK, UV> getMapState(MapStateDescriptor<UK, UV> stateProperties) {
	if (!stateRegistrationAllowed) {
		throw new RuntimeException(REGISTRATION_EXCEPTION_MSG);
	}

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

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

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

	output = new JoinedRowData();

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

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

	ValueStateDescriptor<Long> cleanupTsStateDescriptor = new ValueStateDescriptor<>(
		"cleanupTsState",
		Types.LONG
	);
	this.cleanupTsState = getRuntimeContext().getState(cleanupTsStateDescriptor);
}
 

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

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

private void testBroadcastStateKeyUpgrade(
		MapStateDescriptor<TestType, Integer> initialAccessDescriptor,
		MapStateDescriptor<TestType, Integer> newAccessDescriptorAfterRestore) throws Exception {

	CheckpointStreamFactory streamFactory = createStreamFactory();

	OperatorStateBackend backend = createOperatorStateBackend();

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

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

		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
		assertEquals((Integer) 3, state.get(new TestType("foo", 13)));
		assertEquals((Integer) 5, state.get(new TestType("bar", 278)));
		state.put(new TestType("new-entry", 777), 17);
	} finally {
		backend.dispose();
	}
}
 

@Override
public <K, V> BroadcastState<K, V> getBroadcastState(MapStateDescriptor<K, V> descriptor) {
	try {
		return getOperatorStateBackend().getBroadcastState(descriptor);
	} catch (Exception e) {
		throw new FlinkRuntimeException(e);
	}
}
 

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

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

ReadWriteContextImpl(
		final ExecutionConfig executionConfig,
		final BroadcastProcessFunction<IN1, IN2, OUT> function,
		final Map<MapStateDescriptor<?, ?>, BroadcastState<?, ?>> broadcastStates,
		final ProcessingTimeService timerService) {

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

@Test
public void testBroadcastStateValueMigration() throws Exception {
	final String stateName = "broadcast-state";

	testBroadcastStateValueUpgrade(
		new MapStateDescriptor<>(
			stateName,
			IntSerializer.INSTANCE,
			new TestType.V1TestTypeSerializer()),
		new MapStateDescriptor<>(
			stateName,
			IntSerializer.INSTANCE,
			// new value serializer is a V2 serializer with a different schema
			new TestType.V2TestTypeSerializer()));
}
 

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

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

protected BroadcastConnectedStream(
		final StreamExecutionEnvironment env,
		final DataStream<IN1> input1,
		final BroadcastStream<IN2> input2,
		final List<MapStateDescriptor<?, ?>> broadcastStateDescriptors) {
	this.environment = requireNonNull(env);
	this.inputStream1 = requireNonNull(input1);
	this.inputStream2 = requireNonNull(input2);
	this.broadcastStateDescriptors = requireNonNull(broadcastStateDescriptors);
}
 

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

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

private TestMapState() throws Exception {
	this.state = backend.createInternalState(
		VoidNamespaceSerializer.INSTANCE,
		new MapStateDescriptor<>("TestMapState", StringSerializer.INSTANCE, StringSerializer.INSTANCE),
		snapshotTransformFactory);
	state.setCurrentNamespace(VoidNamespace.INSTANCE);
}
 

@Test
public void testBroadcastStateValueMigration() throws Exception {
	final String stateName = "broadcast-state";

	testBroadcastStateValueUpgrade(
		new MapStateDescriptor<>(
			stateName,
			IntSerializer.INSTANCE,
			new TestType.V1TestTypeSerializer()),
		new MapStateDescriptor<>(
			stateName,
			IntSerializer.INSTANCE,
			// new value serializer is a V2 serializer with a different schema
			new TestType.V2TestTypeSerializer()));
}
 

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

	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.V2TestTypeSerializer()));
}
 

private InputSideHasNoUniqueKey(
		RuntimeContext ctx,
		String stateName,
		BaseRowTypeInfo recordType,
		StateTtlConfig ttlConfig) {
	MapStateDescriptor<BaseRow, Integer> recordStateDesc = new MapStateDescriptor<>(
		stateName,
		recordType,
		Types.INT);
	if (!ttlConfig.equals(StateTtlConfig.DISABLED)) {
		recordStateDesc.enableTimeToLive(ttlConfig);
	}
	this.recordState = ctx.getMapState(recordStateDesc);
}