Flink
KeyBy
Multiple KeyBy
Data Streaming
Programming Tips

How to support multiple KeyBy in Flink

System Design practice on Codemia

Work through 120+ system design problems with detailed solutions, from rate limiters to multi-region storage.

Practice system design

Introduction

In Flink, a stream is keyed by one key definition at a time. If you need logic based on more than one field, the right solution is usually a composite key, a re-keying step after some transformation, or branching the stream into separate keyed pipelines.

Understand What keyBy Really Does

keyBy partitions a stream so that all records with the same key go to the same logical keyed partition. Stateful operators after that point use that key for state access and grouping.

If you write:

java
stream.keyBy(event -> event.userId())

you now have a stream keyed by user. Flink does not keep a second independent keyed state space for product, country, or some other field at the same time in that same operator chain.

That is why "multiple keyBy" usually means one of three different needs.

Use a Composite Key When the Grouping Is Combined

If the real grouping key is the combination of two fields, create one composite key and key by that.

java
1DataStream<Event> stream = ...;
2
3KeyedStream<Event, Tuple2<String, String>> keyed =
4    stream.keyBy(event -> Tuple2.of(event.country(), event.city()));

You can then aggregate on the pair:

java
keyed
    .sum("amount");

This is the right approach when the question is "group by country and city together".

Re-Key After a Transformation When the Stages Differ

Sometimes you need one stage keyed by one field and a later stage keyed by another. That is allowed, but it is sequential, not simultaneous.

java
1DataStream<CountrySummary> countrySummaries =
2    stream
3        .keyBy(Event::country)
4        .process(new CountryProcessFunction());
5
6DataStream<CitySummary> citySummaries =
7    countrySummaries
8        .keyBy(CountrySummary::city)
9        .process(new CityProcessFunction());

This works because the second keyBy repartitions the stream again. The important tradeoff is that keyed state from the first stage does not magically carry over as keyed state for the second key.

Branch the Stream If You Need Two Independent Keyed Views

If the same raw stream must drive separate user-keyed and product-keyed computations, branch it into multiple pipelines.

java
1DataStream<Event> stream = ...;
2
3DataStream<UserResult> byUser =
4    stream
5        .keyBy(Event::userId)
6        .process(new UserProcessFunction());
7
8DataStream<ProductResult> byProduct =
9    stream
10        .keyBy(Event::productId)
11        .process(new ProductProcessFunction());

That is the clean design when the two keyed analyses are logically separate. Each branch gets its own keyed state and scaling behavior.

Common Pitfalls

The biggest mistake is expecting one operator to hold keyed state for several unrelated keyBy dimensions at once. A keyed operator works with one current key definition.

Another common issue is using nested keyBy calls and assuming they create hierarchical state automatically. In reality, each keyBy repartitions the stream for downstream operators only.

People also force composite keys into cases where they really need separate views. If you need both per-user state and per-product state independently, branching the stream is usually clearer than trying to encode everything into one huge key.

Finally, remember that every keyBy can trigger a shuffle. Re-keying too often has a performance cost, so choose the structure that matches the actual computation.

When in doubt, sketch the state ownership per operator before coding. That makes it much easier to see whether you need one combined key or several separate keyed branches.

Summary

  • A Flink stream is keyed by one key definition at a time in a given keyed operator chain.
  • Use a composite key when the grouping is the combination of several fields.
  • Re-key the stream when later stages need a different grouping key.
  • Branch the stream when you need independent keyed computations over the same events.
  • Do not expect several unrelated keyed state spaces to exist implicitly in one operator.

Related reading
Course
Beginner
27 lessons
10 hours
System Design Fundamentals

Build a strong foundation in designing scalable, reliable distributed systems.

View the course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

System Design practice on Codemia

Work through 120+ system design problems with detailed solutions, from rate limiters to multi-region storage.

Practice system design

All Rights Reserved.