Polymorphic Stream Collector In Java

Let’s answer this Stackoverflow question by implementing a polymorphic stream collector, mixing F.P. and O.O.P. concepts.

Overview

In this short article, we’ll play lambda expressions and various ways of collecting streams in Java. We’ll start with this StackOverflow question and we’ll solve it by combing functional features such as passing functions as the first-class citizens and O.OP. concepts such as polymorphism.

In my opinion, there is no such thing as “the” silver bullet and all paradigms have interesting things to offer. I believe this combination of F.P. and O.O.P. can lead to an elegant way of solving the problem asked in the question.

The Question

The original question from stack overflow has an example that is perhaps harder to read or understand. This is probably because the author tried to hide the original code due to internal policy reasons. Unfortunately, class names such as A and B, or variable names such as AvsB will make the task at hand even harder, therefore I decided to rephrase a bit the question.

Let’s assume we have a FieldValue class that wraps a long number:

record FieldValue(long value){}

These values, come as input as part of a Map where the key is an Aggregator. An aggregator is an object that has an id and a type:

record Aggregator(long id, Type type) 
    public enum Type {
        SUM, AVG, MAX, MIN;
    }
}

Let’s assume this is the Map<Aggregator, FieldValue> input of our program:

Map<Aggregator, FieldValue> values = Map.of(
    new Aggregator(1L, Type.AVG), new FieldValue(10L),
    new Aggregator(2L, Type.AVG), new FieldValue(20L),
    new Aggregator(3L, Type.SUM), new FieldValue(30L),
    new Aggregator(4L, Type.SUM), new FieldValue(40L),
    new Aggregator(5L, Type.MAX), new FieldValue(50L),
    new Aggregator(6L, Type.MAX), new FieldValue(60L),
    new Aggregator(7L, Type.MIN), new FieldValue(70L),
    new Aggregator(8L, Type.MIN), new FieldValue(80L)
);

Our goal is to group the data by the Aggregator.Type and aggregate their values accordingly. Therefore we would like our function to return a Map<Aggregator.Type, Double> with the key-value pairs:

• Type.AVG: 15 // avg(10, 20)
• Type.SUM: 70 // sum(30, 40)
• Type.MAX: 60 // max(50, 60)
• Type.MIN: 70 // min(70, 80)

In fact, we can write the unit-test already:

@Test
void test() {
    //given
    Map<Aggregator, FieldValue> values = Map.of(
        new Aggregator(1L, Type.AVG), new FieldValue(10L),
        new Aggregator(2L, Type.AVG), new FieldValue(20L),

        new Aggregator(3L, Type.SUM), new FieldValue(30L),
        new Aggregator(4L, Type.SUM), new FieldValue(40L),

        new Aggregator(5L, Type.MAX), new FieldValue(50L),
        new Aggregator(6L, Type.MAX), new FieldValue(60L),

        new Aggregator(7L, Type.MIN), new FieldValue(70L),
        new Aggregator(8L, Type.MIN), new FieldValue(80L)
    );

    // when
    Map<Aggregator.Type, Double> result = aggregate(values);

    //then
    assertThat(result).isEqualTo(Map.of(
        Type.AVG, 15.0,
        Type.SUM, 70.0,
        Type.MAX, 60.0,
        Type.MIN, 70.0
    ));
}

public Map<Aggregator.Type, Double> aggregate(Map<Aggregator, FieldValue> fields) {
    // TODO: implement me!
    return new HashMap<>();
}

The Challenges

Even though it looks like a simple task, there are a few challenges we’ll encounter. First of all, we can try streaming the entry set and then collect it using Collectors.toMap or Collectors.groupBy. Unfortunately, it is impossible to group elements using different rules or different collectors (without creating your own custom collector).

Secondly, the various out-of-the-box collectors we might be tempted to use have different return types:

• Collectors.maxBy() and minBy() — will return Optional<Value>
• Collectors.summingDouble() and averagingDouble() — will return Double

Even if we will be using a LongStream or DoubleStream and attempt to use max(), min(), average(), and sum(), the return types will be different.

The Implementation

Firstly, let’s unwrap the data and discard the useless pieces of information, such as the aggregator id:

Map<Aggregator.Type, List<Long>> valuesByType = fields.entrySet()
    .stream()
    .collect(Collectors.groupingBy(
        entry -> entry.getKey().type(),
        Collectors.mapping(
            entry -> entry.getValue().value(),
            Collectors.toList())
    ));

After this step, the new map should look something like this:

{
  Type.AVG: [10, 20],
  Type.SUM: [30, 40],
  Type.MAX: [50, 60],
  Type.MIN: [70, 80]
}

At this point, we can re-iterate through the key-value pairs and try to apply a specific collector to each set of values:

return valuesByType.entrySet()
    .stream()
    .collect(Collectors.toMap(
        entry.getKey(),
        entry.getKey().aggregate(entry.getValue())
    ));

As we can see, we’ll keep the Aggregator.Type as key of our map, and we’ll call a collectValues that will perform the specific collector for each type.

PS: I decided to pass the key and the value as separate arguments instead of passing the whole entry not to couple this function with the internal implementation of java.util.Map

private double collectValues(Type aggregator, List<Long> values) {
    DoubleStream stream = values.stream().mapToDouble(x -> x + 0.0d);
    return switch (aggregator) {
        case AVG -> stream.average().orElseThrow();
        case SUM -> stream.sum();
        case MAX -> stream.max().orElseThrow();
        case MIN -> stream.min().orElseThrow();
        default -> throw new IllegalArgumentException();
    };
}

This new enhanced switch will allow us to call the appropriate aggregator in a declarative manner. In this example, the type is an Enum, but, Starting with Java20, we’ll be able to do it on sealed classes. This pattern is often called “ad-hoc polymorphism”.

In my opinion, this is not true polymorphism because this method still depends on the various aggregator types. Let’s do a bit of refactoring and add a bit of logic inside the Aggregator.Type:

@RequiredArgsConstructor
public enum Type {
    SUM(values -> doubleStream(values).sum()),
    AVG(values -> doubleStream(values).average().orElseThrow()),
    MAX(values -> doubleStream(values).max().orElseThrow()),
    MIN(values -> doubleStream(values).min().orElseThrow());

    private final Function<List<Long>, Double> aggregator;

    public Double aggregate(List<Long> values) {
        return aggregator.apply(values);
    }

    private static DoubleStream doubleStream(List<Long> values) {
        return  values.stream().mapToDouble(v -> v + 0.0d);
    }
}

Now, we can use the polymorphic aggregate method and remove the switch statement. Let's read through the whole code again and run the tests:

@Test
void test() {
    //given
    Map<Aggregator, FieldValue> values = // Map.of( ... );

    // when
    Map<Aggregator.Type, Double> result = aggregate(values);

    //then
    assertThat(result).isEqualTo(Map.of(
        Type.AVG, 15d,
        Type.SUM, 70d,
        Type.MAX, 60d,
        Type.MIN, 70d
    ));
}

public Map<Aggregator.Type, Double> aggregate(Map<Aggregator, FieldValue> fields) {
    Map<Aggregator.Type, List<Long>> valuesByType = fields.entrySet()
        .stream()
        .collect(Collectors.groupingBy(
            entry -> entry.getKey().type(),
            Collectors.mapping(
                entry -> entry.getValue().value(),
                Collectors.toList())
        ));
    
    return valuesByType.entrySet()
        .stream()
        .collect(Collectors.toMap(
            entry.getKey(),
            entry.getKey().aggregate(entry.getValue())
        ));
}

Code Smells?

The test is passing! We followed the question and solved the problem, great!

I just want to mention two potential code smells here. Firstly, look at the aggregate method: we collect the data into a local variable, and only after that, we use it to stream again. I believe this is the proper way of doing it. Even though it might be tempting to continue the pipeline right away, this can lead to very large and cluttered functional pipelines. So, in my opinion, we should stay away from this anti-pattern: .collect(…).stream()

Secondly, the Maps<> we are using here might not be the correct data structure in our case. We can enrich the project with more meaningful classes and types. Lines such as this one: entry.getKey().aggregate(entry.getValue()) are a clear manifestation of the “primitive obsession” anti-pattern.

If you are interested in reading more about this anti-pattern, I have an article dedicated to this topic. Feel free to dive deeper into the subject and let me know what you think:

Thank You!

Thanks for reading the article and please let me know what you think! Any feedback is welcome.

If you want to read more about clean code, design, unit testing, object-oriented programming, functional programming, and many others, make sure to check out my other articles. Do you like the content? Consider following or subscribing to the email list.

Finally, if you consider becoming a Medium member and supporting my blog, here’s my referral.

Happy Coding!