Can I use a custom partitioner with group by?

Question

Let's say that I know that my dataset is unbalanced and I know the distribution of the keys. I'd like leverage this to write a custom partitioner to get the most out of the operator instances.

I know about DataStream#partitionCustom. However, if my stream is keyed, will it still work properly? My job would look something like:

KeyedDataStream afterCustomPartition = keyedStream.partitionCustom(new MyPartitioner(), MyPartitionKeySelector())

DataStreamUtils.reinterpretAsKeyedStream(afterCustomPartition, new MyGroupByKeySelector<>()).sum()

What I'm trying to achieve is:

• Having a stream keyBy according to some key so that the reduce function will only be called with elements from that key.
• The group by split the work across nodes based on some custom partitioning.
• The custom partitioning returning a number based on the number of parallel operator instances (which will be fixed and not subject to rescaling).
• The custom partioning returning different values from the keyBy. However, keyBy(x) = keyBy(y) => partition(x) = partition(y).
• Having pre-aggregation to minimize network traffic before partitioning.

Example of the use-case:

• Dataset: [(0, A), (0, B), (0, C), (1, D), (2, E)]
• Number of parallel operator instances: 2
• Group by function: returns the 1st element of the pair
• Partition function: returns 0 for key 0 and 1 for keys 1 and 2. Advantage: deal with the data skew that could be sending keys 0 and 1 to the same operator instance which would mean that one operator instance would receive 80% of the dataset.

Answer 1

For those who still are looking for solution, I made a dirty hack to workaround this limitation. Now, this is bad workaround, so I would recommend to use only if there absolutely no other way to solve your issue. What I did is created an inverse function for MurMurHash function that is used inside Flink code to partition the data. Now, you can provide a tampered key for KeySelector like that:

stream.keyBy(x -> inverseMurmurHash(x.getSomeIntField()))

And that will be mapped to a int value you provided from record.

With this function you can basically implement your own physical partitioner that returns KeyedStream. Of course you need to be mindful and careful, because it will work as long as Flink doesn't change hash algorithm.

Answer 2

That is unfortunately not possible. DataStreamUtils.reinterpretAsKeyedStream() requires that the data is identically partitioned as if you would have called keyBy().

The reason for this limitation are key groups and how keys are mapped to key groups. A key group is Flink's unit of how keyed state is distributed. The number of key groups determines the maximum parallelism of an operator and is configured with setMaxParallelism(). Keys are assigned to key groups with an internal hash function. By changing the partitioning of keys, keys for the same key group will be distributed across multiple machines which will not work.

In order to tweak the assignment of key to machines, you would need to change the assignment of keys to key groups. However, there is no a public or accessible interface to do that. Therefore, custom key distributions are not supported in Flink 1.6.