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

The following examples show how to use org.apache.flink.runtime.state.DefaultKeyedStateStore. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.

Example #1

Source File: StreamingRuntimeContextTest.java From Flink-CEPplus with Apache License 2.0

6 votes

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

Example #2

Source File: StreamingRuntimeContextTest.java From flink with Apache License 2.0

6 votes

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

Example #3

Source File: StreamOperatorStateHandler.java From flink with Apache License 2.0

6 votes

public StreamOperatorStateHandler(
		StreamOperatorStateContext context,
		ExecutionConfig executionConfig,
		CloseableRegistry closeableRegistry) {
	this.context = context;
	operatorStateBackend = context.operatorStateBackend();
	keyedStateBackend = context.keyedStateBackend();
	this.closeableRegistry = closeableRegistry;

	if (keyedStateBackend != null) {
		keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, executionConfig);
	}
	else {
		keyedStateStore = null;
	}
}

Example #4

Source File: StreamingRuntimeContextTest.java From Flink-CEPplus with Apache License 2.0

5 votes

@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createListPlainMockOp() throws Exception {

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

	KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);

	DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);

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

	doAnswer(new Answer<ListState<String>>() {

		@Override
		public ListState<String> answer(InvocationOnMock invocationOnMock) throws Throwable {
			ListStateDescriptor<String> descr =
					(ListStateDescriptor<String>) invocationOnMock.getArguments()[2];

			AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend(
				new DummyEnvironment("test_task", 1, 0),
				new JobID(),
				"test_op",
				IntSerializer.INSTANCE,
				1,
				new KeyGroupRange(0, 0),
				new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()),
				TtlTimeProvider.DEFAULT,
				new UnregisteredMetricsGroup(),
				Collections.emptyList(),
				new CloseableRegistry());
			backend.setCurrentKey(0);
			return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr);
		}
	}).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(ListStateDescriptor.class));

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

Example #5

Source File: StreamingRuntimeContextTest.java From Flink-CEPplus with Apache License 2.0

5 votes

@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createMapPlainMockOp() throws Exception {

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

	KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);

	DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);

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

	doAnswer(new Answer<MapState<Integer, String>>() {

		@Override
		public MapState<Integer, String> answer(InvocationOnMock invocationOnMock) throws Throwable {
			MapStateDescriptor<Integer, String> descr =
					(MapStateDescriptor<Integer, String>) invocationOnMock.getArguments()[2];

			AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend(
				new DummyEnvironment("test_task", 1, 0),
				new JobID(),
				"test_op",
				IntSerializer.INSTANCE,
				1,
				new KeyGroupRange(0, 0),
				new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()),
				TtlTimeProvider.DEFAULT,
				new UnregisteredMetricsGroup(),
				Collections.emptyList(),
				new CloseableRegistry());
			backend.setCurrentKey(0);
			return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr);
		}
	}).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(MapStateDescriptor.class));

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

Example #6

Source File: KeyedStateInputFormat.java From flink with Apache License 2.0

5 votes

@Override
@SuppressWarnings("unchecked")
public void open(KeyGroupRangeInputSplit split) throws IOException {
	registry = new CloseableRegistry();

	final Environment environment = new SavepointEnvironment
		.Builder(getRuntimeContext(), split.getNumKeyGroups())
		.setSubtaskIndex(split.getSplitNumber())
		.setPrioritizedOperatorSubtaskState(split.getPrioritizedOperatorSubtaskState())
		.build();

	final StreamOperatorStateContext context = getStreamOperatorStateContext(environment);

	keyedStateBackend = (AbstractKeyedStateBackend<K>) context.keyedStateBackend();

	final DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, getRuntimeContext().getExecutionConfig());
	SavepointRuntimeContext ctx = new SavepointRuntimeContext(getRuntimeContext(), keyedStateStore);
	FunctionUtils.setFunctionRuntimeContext(userFunction, ctx);

	keys = getKeyIterator(ctx);

	final InternalTimerService<VoidNamespace> timerService = restoreTimerService(context);
	try {
		this.ctx = new Context(keyedStateBackend, timerService);
	} catch (Exception e) {
		throw new IOException("Failed to restore timer state", e);
	}
}

Example #7

Source File: StreamingRuntimeContextTest.java From flink with Apache License 2.0

5 votes

@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createListPlainMockOp() throws Exception {

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

	KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);

	DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);

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

	doAnswer(new Answer<ListState<String>>() {

		@Override
		public ListState<String> answer(InvocationOnMock invocationOnMock) throws Throwable {
			ListStateDescriptor<String> descr =
					(ListStateDescriptor<String>) invocationOnMock.getArguments()[2];

			AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend(
				new DummyEnvironment("test_task", 1, 0),
				new JobID(),
				"test_op",
				IntSerializer.INSTANCE,
				1,
				new KeyGroupRange(0, 0),
				new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()),
				TtlTimeProvider.DEFAULT,
				new UnregisteredMetricsGroup(),
				Collections.emptyList(),
				new CloseableRegistry());
			backend.setCurrentKey(0);
			return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr);
		}
	}).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(ListStateDescriptor.class));

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

Example #8

Source File: StreamingRuntimeContextTest.java From flink with Apache License 2.0

5 votes

@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createMapPlainMockOp() throws Exception {

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

	KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);

	DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);

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

	doAnswer(new Answer<MapState<Integer, String>>() {

		@Override
		public MapState<Integer, String> answer(InvocationOnMock invocationOnMock) throws Throwable {
			MapStateDescriptor<Integer, String> descr =
					(MapStateDescriptor<Integer, String>) invocationOnMock.getArguments()[2];

			AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend(
				new DummyEnvironment("test_task", 1, 0),
				new JobID(),
				"test_op",
				IntSerializer.INSTANCE,
				1,
				new KeyGroupRange(0, 0),
				new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()),
				TtlTimeProvider.DEFAULT,
				new UnregisteredMetricsGroup(),
				Collections.emptyList(),
				new CloseableRegistry());
			backend.setCurrentKey(0);
			return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr);
		}
	}).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(MapStateDescriptor.class));

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

Example #9

Source File: KeyedStateInputFormat.java From flink with Apache License 2.0

5 votes

@Override
@SuppressWarnings("unchecked")
public void open(KeyGroupRangeInputSplit split) throws IOException {
	registry = new CloseableRegistry();

	final Environment environment = new SavepointEnvironment
		.Builder(getRuntimeContext(), split.getNumKeyGroups())
		.setConfiguration(configuration)
		.setSubtaskIndex(split.getSplitNumber())
		.setPrioritizedOperatorSubtaskState(split.getPrioritizedOperatorSubtaskState())
		.build();

	final StreamOperatorStateContext context = getStreamOperatorStateContext(environment);

	AbstractKeyedStateBackend<K> keyedStateBackend = (AbstractKeyedStateBackend<K>) context.keyedStateBackend();

	final DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, getRuntimeContext().getExecutionConfig());
	SavepointRuntimeContext ctx = new SavepointRuntimeContext(getRuntimeContext(), keyedStateStore);

	InternalTimeServiceManager<K> timeServiceManager = (InternalTimeServiceManager<K>) context.internalTimerServiceManager();
	try {
		operator.setup(getRuntimeContext().getExecutionConfig(), keyedStateBackend, timeServiceManager, ctx);
		operator.open();
		keysAndNamespaces = operator.getKeysAndNamespaces(ctx);
	} catch (Exception e) {
		throw new IOException("Failed to restore timer state", e);
	}
}

Example #10

Source File: StreamingRuntimeContextTest.java From flink with Apache License 2.0

5 votes

@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createListPlainMockOp() throws Exception {

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

	KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);

	DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);

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

	doAnswer(new Answer<ListState<String>>() {

		@Override
		public ListState<String> answer(InvocationOnMock invocationOnMock) throws Throwable {
			ListStateDescriptor<String> descr =
					(ListStateDescriptor<String>) invocationOnMock.getArguments()[2];

			AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend(
				new DummyEnvironment("test_task", 1, 0),
				new JobID(),
				"test_op",
				IntSerializer.INSTANCE,
				1,
				new KeyGroupRange(0, 0),
				new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()),
				TtlTimeProvider.DEFAULT,
				new UnregisteredMetricsGroup(),
				Collections.emptyList(),
				new CloseableRegistry());
			backend.setCurrentKey(0);
			return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr);
		}
	}).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(ListStateDescriptor.class));

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

Example #11

Source File: AbstractStreamOperator.java From Flink-CEPplus with Apache License 2.0

4 votes

@Override
public final void initializeState() throws Exception {

	final TypeSerializer<?> keySerializer = config.getStateKeySerializer(getUserCodeClassloader());

	final StreamTask<?, ?> containingTask =
		Preconditions.checkNotNull(getContainingTask());
	final CloseableRegistry streamTaskCloseableRegistry =
		Preconditions.checkNotNull(containingTask.getCancelables());
	final StreamTaskStateInitializer streamTaskStateManager =
		Preconditions.checkNotNull(containingTask.createStreamTaskStateInitializer());

	final StreamOperatorStateContext context =
		streamTaskStateManager.streamOperatorStateContext(
			getOperatorID(),
			getClass().getSimpleName(),
			this,
			keySerializer,
			streamTaskCloseableRegistry,
			metrics);

	this.operatorStateBackend = context.operatorStateBackend();
	this.keyedStateBackend = context.keyedStateBackend();

	if (keyedStateBackend != null) {
		this.keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, getExecutionConfig());
	}

	timeServiceManager = context.internalTimerServiceManager();

	CloseableIterable<KeyGroupStatePartitionStreamProvider> keyedStateInputs = context.rawKeyedStateInputs();
	CloseableIterable<StatePartitionStreamProvider> operatorStateInputs = context.rawOperatorStateInputs();

	try {
		StateInitializationContext initializationContext = new StateInitializationContextImpl(
			context.isRestored(), // information whether we restore or start for the first time
			operatorStateBackend, // access to operator state backend
			keyedStateStore, // access to keyed state backend
			keyedStateInputs, // access to keyed state stream
			operatorStateInputs); // access to operator state stream

		initializeState(initializationContext);
	} finally {
		closeFromRegistry(operatorStateInputs, streamTaskCloseableRegistry);
		closeFromRegistry(keyedStateInputs, streamTaskCloseableRegistry);
	}
}

Example #12

Source File: AbstractStreamOperator.java From flink with Apache License 2.0

4 votes

@Override
public final void initializeState() throws Exception {

	final TypeSerializer<?> keySerializer = config.getStateKeySerializer(getUserCodeClassloader());

	final StreamTask<?, ?> containingTask =
		Preconditions.checkNotNull(getContainingTask());
	final CloseableRegistry streamTaskCloseableRegistry =
		Preconditions.checkNotNull(containingTask.getCancelables());
	final StreamTaskStateInitializer streamTaskStateManager =
		Preconditions.checkNotNull(containingTask.createStreamTaskStateInitializer());

	final StreamOperatorStateContext context =
		streamTaskStateManager.streamOperatorStateContext(
			getOperatorID(),
			getClass().getSimpleName(),
			this,
			keySerializer,
			streamTaskCloseableRegistry,
			metrics);

	this.operatorStateBackend = context.operatorStateBackend();
	this.keyedStateBackend = context.keyedStateBackend();

	if (keyedStateBackend != null) {
		this.keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, getExecutionConfig());
	}

	timeServiceManager = context.internalTimerServiceManager();

	CloseableIterable<KeyGroupStatePartitionStreamProvider> keyedStateInputs = context.rawKeyedStateInputs();
	CloseableIterable<StatePartitionStreamProvider> operatorStateInputs = context.rawOperatorStateInputs();

	try {
		StateInitializationContext initializationContext = new StateInitializationContextImpl(
			context.isRestored(), // information whether we restore or start for the first time
			operatorStateBackend, // access to operator state backend
			keyedStateStore, // access to keyed state backend
			keyedStateInputs, // access to keyed state stream
			operatorStateInputs); // access to operator state stream

		initializeState(initializationContext);
	} finally {
		closeFromRegistry(operatorStateInputs, streamTaskCloseableRegistry);
		closeFromRegistry(keyedStateInputs, streamTaskCloseableRegistry);
	}
}

Example #13

Source File: StreamingRuntimeContextTest.java From flink with Apache License 2.0

4 votes

@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createMapPlainMockOp() throws Exception {

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

	KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);

	DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);

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

	doAnswer(new Answer<MapState<Integer, String>>() {

		@Override
		public MapState<Integer, String> answer(InvocationOnMock invocationOnMock) throws Throwable {
			MapStateDescriptor<Integer, String> descr =
					(MapStateDescriptor<Integer, String>) invocationOnMock.getArguments()[2];

			AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend(
				new DummyEnvironment("test_task", 1, 0),
				new JobID(),
				"test_op",
				IntSerializer.INSTANCE,
				1,
				new KeyGroupRange(0, 0),
				new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()),
				TtlTimeProvider.DEFAULT,
				new UnregisteredMetricsGroup(),
				Collections.emptyList(),
				new CloseableRegistry());
			backend.setCurrentKey(0);
			return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr);
		}
	}).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(MapStateDescriptor.class));

	when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore);
	when(operatorMock.getOperatorID()).thenReturn(new OperatorID());
	when(operatorMock.getProcessingTimeService()).thenReturn(new TestProcessingTimeService());
	return operatorMock;
}