Java Code Examples for org.apache.flink.runtime.state.KeyGroupRangeAssignment#assignToKeyGroup()

@Override
public byte[] getSerializedValue(
		final byte[] serializedKeyAndNamespace,
		final TypeSerializer<K> safeKeySerializer,
		final TypeSerializer<N> safeNamespaceSerializer,
		final TypeSerializer<V> safeValueSerializer) throws Exception {

	//TODO make KvStateSerializer key-group aware to save this round trip and key-group computation
	Tuple2<K, N> keyAndNamespace = KvStateSerializer.deserializeKeyAndNamespace(
			serializedKeyAndNamespace, safeKeySerializer, safeNamespaceSerializer);

	int keyGroup = KeyGroupRangeAssignment.assignToKeyGroup(keyAndNamespace.f0, backend.getNumberOfKeyGroups());

	RocksDBSerializedCompositeKeyBuilder<K> keyBuilder =
					new RocksDBSerializedCompositeKeyBuilder<>(
						safeKeySerializer,
						backend.getKeyGroupPrefixBytes(),
						32
					);
	keyBuilder.setKeyAndKeyGroup(keyAndNamespace.f0, keyGroup);
	byte[] key = keyBuilder.buildCompositeKeyNamespace(keyAndNamespace.f1, namespaceSerializer);
	return backend.db.get(columnFamily, key);
}
 

@Override
public byte[] getSerializedValue(
		final byte[] serializedKeyAndNamespace,
		final TypeSerializer<K> safeKeySerializer,
		final TypeSerializer<N> safeNamespaceSerializer,
		final TypeSerializer<V> safeValueSerializer) throws Exception {

	//TODO make KvStateSerializer key-group aware to save this round trip and key-group computation
	Tuple2<K, N> keyAndNamespace = KvStateSerializer.deserializeKeyAndNamespace(
			serializedKeyAndNamespace, safeKeySerializer, safeNamespaceSerializer);

	int keyGroup = KeyGroupRangeAssignment.assignToKeyGroup(keyAndNamespace.f0, backend.getNumberOfKeyGroups());

	RocksDBSerializedCompositeKeyBuilder<K> keyBuilder =
					new RocksDBSerializedCompositeKeyBuilder<>(
						safeKeySerializer,
						backend.getKeyGroupPrefixBytes(),
						32
					);
	keyBuilder.setKeyAndKeyGroup(keyAndNamespace.f0, keyGroup);
	byte[] key = keyBuilder.buildCompositeKeyNamespace(keyAndNamespace.f1, namespaceSerializer);
	return backend.db.get(columnFamily, key);
}
 

/** Tries to append next entry to {@code partitioningSource} array snapshot and returns next index.*/
int tryAddToSource(int currentIndex, CopyOnWriteStateTable.StateTableEntry<K, N, S> entry) {
	final int keyGroup = KeyGroupRangeAssignment.assignToKeyGroup(entry.key, totalKeyGroups);
	reportKeyGroupOfElementAtIndex(currentIndex, keyGroup);
	partitioningSource[currentIndex] = entry;
	return currentIndex + 1;
}
 

@Override
public KeyedStateRow map(Tuple2<K, V> t) throws Exception {
	int keyGroup = KeyGroupRangeAssignment.assignToKeyGroup(t.f0, maxParallelism);
	ByteArrayOutputStreamWithPos os = new ByteArrayOutputStreamWithPos();
	DataOutputViewStreamWrapper ov = new DataOutputViewStreamWrapper(os);

	RocksDBUtils.writeKeyGroup(keyGroup, keygroupPrefixBytes, ov);
	RocksDBUtils.writeKey(t.f0, keySerializer, os, ov, false);
	RocksDBUtils.writeNameSpace(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, os,
			ov, false);

	os.close();
	return new KeyedStateRow(stateName, os.toByteArray(),
			InstantiationUtil.serializeToByteArray(valueSerializer, t.f1));
}
 

private static int getKeyInKeyGroup(int keyGroup, int maxParallelism) {
	Random rand = new Random(System.currentTimeMillis());
	int result = rand.nextInt();
	while (KeyGroupRangeAssignment.assignToKeyGroup(result, maxParallelism) != keyGroup) {
		result = rand.nextInt();
	}
	return result;
}
 

private static int getKeyInKeyGroup(int keyGroup, int maxParallelism) {
	Random rand = new Random(System.currentTimeMillis());
	int result = rand.nextInt();
	while (KeyGroupRangeAssignment.assignToKeyGroup(result, maxParallelism) != keyGroup) {
		result = rand.nextInt();
	}
	return result;
}
 

public static int getOperatorIndexForKey(String key, int maxParallelism, int parallelism) {
    int keyGroup = KeyGroupRangeAssignment.assignToKeyGroup(key, maxParallelism);
    return KeyGroupRangeAssignment.computeOperatorIndexForKeyGroup(maxParallelism,
            parallelism, keyGroup);
}
 

@Override
public byte[] getSerializedValue(
		final byte[] serializedKeyAndNamespace,
		final TypeSerializer<K> safeKeySerializer,
		final TypeSerializer<N> safeNamespaceSerializer,
		final TypeSerializer<Map<UK, UV>> safeValueSerializer) throws Exception {

	Preconditions.checkNotNull(serializedKeyAndNamespace);
	Preconditions.checkNotNull(safeKeySerializer);
	Preconditions.checkNotNull(safeNamespaceSerializer);
	Preconditions.checkNotNull(safeValueSerializer);

	//TODO make KvStateSerializer key-group aware to save this round trip and key-group computation
	Tuple2<K, N> keyAndNamespace = KvStateSerializer.deserializeKeyAndNamespace(
			serializedKeyAndNamespace, safeKeySerializer, safeNamespaceSerializer);

	int keyGroup = KeyGroupRangeAssignment.assignToKeyGroup(keyAndNamespace.f0, backend.getNumberOfKeyGroups());

	RocksDBSerializedCompositeKeyBuilder<K> keyBuilder =
		new RocksDBSerializedCompositeKeyBuilder<>(
			safeKeySerializer,
			backend.getKeyGroupPrefixBytes(),
			32);

	keyBuilder.setKeyAndKeyGroup(keyAndNamespace.f0, keyGroup);

	final byte[] keyPrefixBytes = keyBuilder.buildCompositeKeyNamespace(keyAndNamespace.f1, namespaceSerializer);

	final MapSerializer<UK, UV> serializer = (MapSerializer<UK, UV>) safeValueSerializer;

	final TypeSerializer<UK> dupUserKeySerializer = serializer.getKeySerializer();
	final TypeSerializer<UV> dupUserValueSerializer = serializer.getValueSerializer();
	final DataInputDeserializer inputView = new DataInputDeserializer();

	final Iterator<Map.Entry<UK, UV>> iterator = new RocksDBMapIterator<Map.Entry<UK, UV>>(
			backend.db,
			keyPrefixBytes,
			dupUserKeySerializer,
			dupUserValueSerializer,
			inputView
		) {

		@Override
		public Map.Entry<UK, UV> next() {
			return nextEntry();
		}
	};

	// Return null to make the behavior consistent with other backends
	if (!iterator.hasNext()) {
		return null;
	}

	return KvStateSerializer.serializeMap(() -> iterator, dupUserKeySerializer, dupUserValueSerializer);
}
 

public void setCurrentKeyAndKeyGroup(K key) {
	super.setCurrentKey(key);
	super.setCurrentKeyGroupIndex(KeyGroupRangeAssignment.assignToKeyGroup(key, getNumberOfKeyGroups()));
}
 

@Override
public int applyAsInt(T value) {
  Object key = keySelector.apply(value);
  return KeyGroupRangeAssignment.assignToKeyGroup(key, maxParallelism);
}
 

@Override
public byte[] getSerializedValue(
		final byte[] serializedKeyAndNamespace,
		final TypeSerializer<K> safeKeySerializer,
		final TypeSerializer<N> safeNamespaceSerializer,
		final TypeSerializer<Map<UK, UV>> safeValueSerializer) throws Exception {

	Preconditions.checkNotNull(serializedKeyAndNamespace);
	Preconditions.checkNotNull(safeKeySerializer);
	Preconditions.checkNotNull(safeNamespaceSerializer);
	Preconditions.checkNotNull(safeValueSerializer);

	//TODO make KvStateSerializer key-group aware to save this round trip and key-group computation
	Tuple2<K, N> keyAndNamespace = KvStateSerializer.deserializeKeyAndNamespace(
			serializedKeyAndNamespace, safeKeySerializer, safeNamespaceSerializer);

	int keyGroup = KeyGroupRangeAssignment.assignToKeyGroup(keyAndNamespace.f0, backend.getNumberOfKeyGroups());

	RocksDBSerializedCompositeKeyBuilder<K> keyBuilder =
		new RocksDBSerializedCompositeKeyBuilder<>(
			safeKeySerializer,
			backend.getKeyGroupPrefixBytes(),
			32);

	keyBuilder.setKeyAndKeyGroup(keyAndNamespace.f0, keyGroup);

	final byte[] keyPrefixBytes = keyBuilder.buildCompositeKeyNamespace(keyAndNamespace.f1, namespaceSerializer);

	final MapSerializer<UK, UV> serializer = (MapSerializer<UK, UV>) safeValueSerializer;

	final TypeSerializer<UK> dupUserKeySerializer = serializer.getKeySerializer();
	final TypeSerializer<UV> dupUserValueSerializer = serializer.getValueSerializer();
	final DataInputDeserializer inputView = new DataInputDeserializer();

	final Iterator<Map.Entry<UK, UV>> iterator = new RocksDBMapIterator<Map.Entry<UK, UV>>(
			backend.db,
			keyPrefixBytes,
			dupUserKeySerializer,
			dupUserValueSerializer,
			inputView
		) {

		@Override
		public Map.Entry<UK, UV> next() {
			return nextEntry();
		}
	};

	// Return null to make the behavior consistent with other backends
	if (!iterator.hasNext()) {
		return null;
	}

	return KvStateSerializer.serializeMap(() -> iterator, dupUserKeySerializer, dupUserValueSerializer);
}
 

@Test
public void testTimerAssignmentToKeyGroups() {
	int totalNoOfTimers = 100;

	int totalNoOfKeyGroups = 100;
	int startKeyGroupIdx = 0;
	int endKeyGroupIdx = totalNoOfKeyGroups - 1; // we have 0 to 99

	@SuppressWarnings("unchecked")
	Set<TimerHeapInternalTimer<Integer, String>>[] expectedNonEmptyTimerSets = new HashSet[totalNoOfKeyGroups];
	TestKeyContext keyContext = new TestKeyContext();

	final KeyGroupRange keyGroupRange = new KeyGroupRange(startKeyGroupIdx, endKeyGroupIdx);

	final PriorityQueueSetFactory priorityQueueSetFactory =
		createQueueFactory(keyGroupRange, totalNoOfKeyGroups);

	InternalTimerServiceImpl<Integer, String> timerService = createInternalTimerService(
		keyGroupRange,
		keyContext,
		new TestProcessingTimeService(),
		IntSerializer.INSTANCE,
		StringSerializer.INSTANCE,
		priorityQueueSetFactory);

	timerService.startTimerService(IntSerializer.INSTANCE, StringSerializer.INSTANCE, mock(Triggerable.class));

	for (int i = 0; i < totalNoOfTimers; i++) {

		// create the timer to be registered
		TimerHeapInternalTimer<Integer, String> timer = new TimerHeapInternalTimer<>(10 + i, i, "hello_world_" + i);
		int keyGroupIdx =  KeyGroupRangeAssignment.assignToKeyGroup(timer.getKey(), totalNoOfKeyGroups);

		// add it in the adequate expected set of timers per keygroup
		Set<TimerHeapInternalTimer<Integer, String>> timerSet = expectedNonEmptyTimerSets[keyGroupIdx];
		if (timerSet == null) {
			timerSet = new HashSet<>();
			expectedNonEmptyTimerSets[keyGroupIdx] = timerSet;
		}
		timerSet.add(timer);

		// register the timer as both processing and event time one
		keyContext.setCurrentKey(timer.getKey());
		timerService.registerEventTimeTimer(timer.getNamespace(), timer.getTimestamp());
		timerService.registerProcessingTimeTimer(timer.getNamespace(), timer.getTimestamp());
	}

	List<Set<TimerHeapInternalTimer<Integer, String>>> eventTimeTimers =
		timerService.getEventTimeTimersPerKeyGroup();
	List<Set<TimerHeapInternalTimer<Integer, String>>> processingTimeTimers =
		timerService.getProcessingTimeTimersPerKeyGroup();

	// finally verify that the actual timers per key group sets are the expected ones.
	for (int i = 0; i < expectedNonEmptyTimerSets.length; i++) {
		Set<TimerHeapInternalTimer<Integer, String>> expected = expectedNonEmptyTimerSets[i];
		Set<TimerHeapInternalTimer<Integer, String>> actualEvent = eventTimeTimers.get(i);
		Set<TimerHeapInternalTimer<Integer, String>> actualProcessing = processingTimeTimers.get(i);

		if (expected == null) {
			Assert.assertTrue(actualEvent.isEmpty());
			Assert.assertTrue(actualProcessing.isEmpty());
		} else {
			Assert.assertEquals(expected, actualEvent);
			Assert.assertEquals(expected, actualProcessing);
		}
	}
}
 

@Test
public void testTimerAssignmentToKeyGroups() {
	int totalNoOfTimers = 100;

	int totalNoOfKeyGroups = 100;
	int startKeyGroupIdx = 0;
	int endKeyGroupIdx = totalNoOfKeyGroups - 1; // we have 0 to 99

	@SuppressWarnings("unchecked")
	Set<TimerHeapInternalTimer<Integer, String>>[] expectedNonEmptyTimerSets = new HashSet[totalNoOfKeyGroups];
	TestKeyContext keyContext = new TestKeyContext();

	final KeyGroupRange keyGroupRange = new KeyGroupRange(startKeyGroupIdx, endKeyGroupIdx);

	final PriorityQueueSetFactory priorityQueueSetFactory =
		createQueueFactory(keyGroupRange, totalNoOfKeyGroups);

	InternalTimerServiceImpl<Integer, String> timerService = createInternalTimerService(
		keyGroupRange,
		keyContext,
		new TestProcessingTimeService(),
		IntSerializer.INSTANCE,
		StringSerializer.INSTANCE,
		priorityQueueSetFactory);

	timerService.startTimerService(IntSerializer.INSTANCE, StringSerializer.INSTANCE, mock(Triggerable.class));

	for (int i = 0; i < totalNoOfTimers; i++) {

		// create the timer to be registered
		TimerHeapInternalTimer<Integer, String> timer = new TimerHeapInternalTimer<>(10 + i, i, "hello_world_" + i);
		int keyGroupIdx =  KeyGroupRangeAssignment.assignToKeyGroup(timer.getKey(), totalNoOfKeyGroups);

		// add it in the adequate expected set of timers per keygroup
		Set<TimerHeapInternalTimer<Integer, String>> timerSet = expectedNonEmptyTimerSets[keyGroupIdx];
		if (timerSet == null) {
			timerSet = new HashSet<>();
			expectedNonEmptyTimerSets[keyGroupIdx] = timerSet;
		}
		timerSet.add(timer);

		// register the timer as both processing and event time one
		keyContext.setCurrentKey(timer.getKey());
		timerService.registerEventTimeTimer(timer.getNamespace(), timer.getTimestamp());
		timerService.registerProcessingTimeTimer(timer.getNamespace(), timer.getTimestamp());
	}

	List<Set<TimerHeapInternalTimer<Integer, String>>> eventTimeTimers =
		timerService.getEventTimeTimersPerKeyGroup();
	List<Set<TimerHeapInternalTimer<Integer, String>>> processingTimeTimers =
		timerService.getProcessingTimeTimersPerKeyGroup();

	// finally verify that the actual timers per key group sets are the expected ones.
	for (int i = 0; i < expectedNonEmptyTimerSets.length; i++) {
		Set<TimerHeapInternalTimer<Integer, String>> expected = expectedNonEmptyTimerSets[i];
		Set<TimerHeapInternalTimer<Integer, String>> actualEvent = eventTimeTimers.get(i);
		Set<TimerHeapInternalTimer<Integer, String>> actualProcessing = processingTimeTimers.get(i);

		if (expected == null) {
			Assert.assertTrue(actualEvent.isEmpty());
			Assert.assertTrue(actualProcessing.isEmpty());
		} else {
			Assert.assertEquals(expected, actualEvent);
			Assert.assertEquals(expected, actualProcessing);
		}
	}
}
 

private int computeKeyGroupIndex(T element) {
	final Object extractKeyFromElement = keyExtractor.extractKeyFromElement(element);
	final int keyGroupId = KeyGroupRangeAssignment.assignToKeyGroup(extractKeyFromElement, totalKeyGroups);
	return globalKeyGroupToLocalIndex(keyGroupId);
}
 

private HashMap<T, T> getDedupMapForElement(T element) {
	int keyGroup = KeyGroupRangeAssignment.assignToKeyGroup(
		keyExtractor.extractKeyFromElement(element),
		totalNumberOfKeyGroups);
	return getDedupMapForKeyGroup(keyGroup);
}
 

@Override
public byte[] getSerializedValue(
		final byte[] serializedKeyAndNamespace,
		final TypeSerializer<K> safeKeySerializer,
		final TypeSerializer<N> safeNamespaceSerializer,
		final TypeSerializer<Map<UK, UV>> safeValueSerializer) throws Exception {

	Preconditions.checkNotNull(serializedKeyAndNamespace);
	Preconditions.checkNotNull(safeKeySerializer);
	Preconditions.checkNotNull(safeNamespaceSerializer);
	Preconditions.checkNotNull(safeValueSerializer);

	//TODO make KvStateSerializer key-group aware to save this round trip and key-group computation
	Tuple2<K, N> keyAndNamespace = KvStateSerializer.deserializeKeyAndNamespace(
			serializedKeyAndNamespace, safeKeySerializer, safeNamespaceSerializer);

	int keyGroup = KeyGroupRangeAssignment.assignToKeyGroup(keyAndNamespace.f0, backend.getNumberOfKeyGroups());

	RocksDBSerializedCompositeKeyBuilder<K> keyBuilder =
		new RocksDBSerializedCompositeKeyBuilder<>(
			safeKeySerializer,
			backend.getKeyGroupPrefixBytes(),
			32);

	keyBuilder.setKeyAndKeyGroup(keyAndNamespace.f0, keyGroup);

	final byte[] keyPrefixBytes = keyBuilder.buildCompositeKeyNamespace(keyAndNamespace.f1, namespaceSerializer);

	final MapSerializer<UK, UV> serializer = (MapSerializer<UK, UV>) safeValueSerializer;

	final TypeSerializer<UK> dupUserKeySerializer = serializer.getKeySerializer();
	final TypeSerializer<UV> dupUserValueSerializer = serializer.getValueSerializer();
	final DataInputDeserializer inputView = new DataInputDeserializer();

	final Iterator<Map.Entry<UK, UV>> iterator = new RocksDBMapIterator<Map.Entry<UK, UV>>(
			backend.db,
			keyPrefixBytes,
			dupUserKeySerializer,
			dupUserValueSerializer,
			inputView
		) {

		@Override
		public Map.Entry<UK, UV> next() {
			return nextEntry();
		}
	};

	// Return null to make the behavior consistent with other backends
	if (!iterator.hasNext()) {
		return null;
	}

	return KvStateSerializer.serializeMap(() -> iterator, dupUserKeySerializer, dupUserValueSerializer);
}
 

@Test
public void testTimerAssignmentToKeyGroups() {
	int totalNoOfTimers = 100;

	int totalNoOfKeyGroups = 100;
	int startKeyGroupIdx = 0;
	int endKeyGroupIdx = totalNoOfKeyGroups - 1; // we have 0 to 99

	@SuppressWarnings("unchecked")
	Set<TimerHeapInternalTimer<Integer, String>>[] expectedNonEmptyTimerSets = new HashSet[totalNoOfKeyGroups];
	TestKeyContext keyContext = new TestKeyContext();

	final KeyGroupRange keyGroupRange = new KeyGroupRange(startKeyGroupIdx, endKeyGroupIdx);

	final PriorityQueueSetFactory priorityQueueSetFactory =
		createQueueFactory(keyGroupRange, totalNoOfKeyGroups);

	InternalTimerServiceImpl<Integer, String> timerService = createInternalTimerService(
		keyGroupRange,
		keyContext,
		new TestProcessingTimeService(),
		IntSerializer.INSTANCE,
		StringSerializer.INSTANCE,
		priorityQueueSetFactory);

	timerService.startTimerService(IntSerializer.INSTANCE, StringSerializer.INSTANCE, mock(Triggerable.class));

	for (int i = 0; i < totalNoOfTimers; i++) {

		// create the timer to be registered
		TimerHeapInternalTimer<Integer, String> timer = new TimerHeapInternalTimer<>(10 + i, i, "hello_world_" + i);
		int keyGroupIdx =  KeyGroupRangeAssignment.assignToKeyGroup(timer.getKey(), totalNoOfKeyGroups);

		// add it in the adequate expected set of timers per keygroup
		Set<TimerHeapInternalTimer<Integer, String>> timerSet = expectedNonEmptyTimerSets[keyGroupIdx];
		if (timerSet == null) {
			timerSet = new HashSet<>();
			expectedNonEmptyTimerSets[keyGroupIdx] = timerSet;
		}
		timerSet.add(timer);

		// register the timer as both processing and event time one
		keyContext.setCurrentKey(timer.getKey());
		timerService.registerEventTimeTimer(timer.getNamespace(), timer.getTimestamp());
		timerService.registerProcessingTimeTimer(timer.getNamespace(), timer.getTimestamp());
	}

	List<Set<TimerHeapInternalTimer<Integer, String>>> eventTimeTimers =
		timerService.getEventTimeTimersPerKeyGroup();
	List<Set<TimerHeapInternalTimer<Integer, String>>> processingTimeTimers =
		timerService.getProcessingTimeTimersPerKeyGroup();

	// finally verify that the actual timers per key group sets are the expected ones.
	for (int i = 0; i < expectedNonEmptyTimerSets.length; i++) {
		Set<TimerHeapInternalTimer<Integer, String>> expected = expectedNonEmptyTimerSets[i];
		Set<TimerHeapInternalTimer<Integer, String>> actualEvent = eventTimeTimers.get(i);
		Set<TimerHeapInternalTimer<Integer, String>> actualProcessing = processingTimeTimers.get(i);

		if (expected == null) {
			Assert.assertTrue(actualEvent.isEmpty());
			Assert.assertTrue(actualProcessing.isEmpty());
		} else {
			Assert.assertEquals(expected, actualEvent);
			Assert.assertEquals(expected, actualProcessing);
		}
	}
}
 

@Override
public S get(K key, N namespace) {
	int keyGroup = KeyGroupRangeAssignment.assignToKeyGroup(key, keyContext.getNumberOfKeyGroups());
	return get(key, keyGroup, namespace);
}
 

private int computeKeyGroupIndex(T element) {
	final Object extractKeyFromElement = keyExtractor.extractKeyFromElement(element);
	final int keyGroupId = KeyGroupRangeAssignment.assignToKeyGroup(extractKeyFromElement, totalKeyGroups);
	return globalKeyGroupToLocalIndex(keyGroupId);
}
 

/**
 * Returns the state for the composite of active key and given namespace. This is typically used by
 * queryable state.
 *
 * @param key       the key. Not null.
 * @param namespace the namespace. Not null.
 * @return the state of the mapping with the specified key/namespace composite key, or {@code null}
 * if no mapping for the specified key is found.
 */
public S get(K key, N namespace) {
	int keyGroup = KeyGroupRangeAssignment.assignToKeyGroup(key, keyContext.getNumberOfKeyGroups());
	return get(key, keyGroup, namespace);
}