Equivalent of DataSet groupBy/withPartitioner for DataStream

user_231762 0
Viewed 68

Previously with a DataSet I could do a .groupBy(...) followed by a .withPartitioner(...) to create groups such that one group (known to be much, much bigger than all the others) would be assigned to its own slot, and the other groups would be distributed among the remaining slots.

In switching to a DataStream, I don't see any straightforward way to do the same thing. If I dig into .keyBy(...), I see it using a PartitionTransformation with a KeyGroupStreamPartitioner, which is promising - but PartitionTransformation is an internal-use only class (or so the annotation says).

What's the recommended approach with a DataStream for achieving the same result?

flink-streaming
2 Answers

With DataStream it's not as straightforward. You can implement a custom Partitioner that you use with partitionCustom, but then you do not have a KeyedStream, and so can not use keyed state or timers.

Another solution is to do a two-step, local/global aggregation, e.g.,

.keyBy(randomizedKey).process(local).keyBy(key).process(global)

And in some cases, the first level of random keying isn't necessary (if the keys are already well distributed among the source partitions).

In principle, given a priori knowledge of the hot keys, you should be able to somehow implement a KeySelector that does a good job of balancing the load among the task slots. I believe one or two people have actually done this (by brute force searching for a suitable mapping from original keys to actual keys), but I don't have a reference implementation at hand.

user_20008230

As David noted, you can sometimes do the double-keyBy trick (initially using a random key) to reduce the impact of key skew. In my case that wasn't viable, as I'm processing records in each group using a large deep learning network with significant memory requirements, which means having all models loaded at the same time for the first grouping.

I re-used a technique I'd gotten to work with an older version of Flink, where you decide which sub-task (operator index) should get each record, and then calculate a key that Flink will assign to the target sub-task. The code, which calculates an Integer key, looks something like:

    public static Integer makeKeyForOperatorIndex(int maxParallelism, int parallelism,
            int operatorIndex) {

        for (int i = 0; i < maxParallelism * 2; i++) {
            Integer key = new Integer(i);
            int index = KeyGroupRangeAssignment.assignKeyToParallelOperator(
                    i, maxParallelism, parallelism);
            
            if (index == operatorIndex) {
                return key;
            }
        }

        throw new RuntimeException(String.format(
                "Unable to find key for target operator index %d (max parallelism = %d, parallelism = %d",
                operatorIndex, maxParallelism, parallelism));
    }

But note this is very fragile, as it depends on internal Flink implementation details.

user_2317620
Related