Using groupingBy into a nested Map, but collecting to a different type of object

Question

So I have this code that "works" (replacing some names for simplicity):

 Map<String, Map<String, ImmutableList<SomeClassA>>> someMap =
      someListOfClassA.stream()
      .filter(...)
      .collect(Collectors.groupingBy(SomeClassA::someCriteriaA,
            Collectors.groupingBy(SomeClassA::someCriteriaB, GuavaCollectors.toImmutableList()
            )
      ));

However, I want to change this code so that the inner collection is of SomeClassB after grouping by SomeClassA fields. For example, if the classes look like this:

assuming they both have all args constructors

class SomeClassA { 
    String someCriteriaA;
    String someCriteriaB;
    T someData;
    String someId;
}

class SomeClassB {
    T someData;
    String someId; 
}

And there's a method somewhere:

public static Collection<SomeClassB> getSomeClassBsFromSomeClassA(SomeClassA someA) {
    List<Some List of Class B> listOfB = someMethod(someA);
    return listOfB; // calls something using someClassA, gets a list of SomeClassB 
}

I want to flatten the the resulting lists of SomeClass Bs into

Map<String, Map<String, ImmutableList<SomeClassB>>> someMap = 
    someListOfClassA.stream()
    .filter(...)
    . // not sure how to group by SomeClassA fields but result in a list of SomeClassB since one SomeClassA can result in multiple SomeClassB

I'm not sure how this would fit into the code above. How can I collect a bunch of lists based on SomeClassB into a single list for all the values of SomeClassA? If a single ClassA maps to a single ClassB I know how to get it to work using Collectors.mapping but since each ClassA results in multiple ClassBs I'm not sure how to get it to work.

Any ideas would be appreciated. Thanks!

Answer 1

With a custom collector like so:

private static Collector<Collection<SomeClassB>, ?, ImmutableList<SomeClassB>>
        flatMapToImmutableList() {
        return Collectors.collectingAndThen(Collectors.toList(),
                listOfCollectionsOfB ->
                        listOfCollectionsOfB.stream()
                                .flatMap(Collection::stream)
                                .collect(GuavaCollectors.toImmutableList()));
    }

you can achieve what you're after:

Map<String, Map<String, List<SomeClassB>>> someMap =
                someListOfClassA.stream()
                        .filter(...)
                        .collect(Collectors.groupingBy(SomeClassA::getSomeCriteriaA,
                                Collectors.groupingBy(SomeClassA::getSomeCriteriaB,
                                        Collectors.mapping(a -> getSomeClassBsFromSomeClassA(a),
                                                flatMapToImmutableList()))));

Answer 2

While we all wait for Java 9's Collectors.flatMapping (thanks to @shmosel for the link), you can write your own collector to achieve what you want:

public static <T, D, R> Collector<T, ?, R> flatMapping(
        Function<? super T, ? extends Stream<? extends D>> streamMapper,
        Collector<? super D, ?, R> downstream) {

    class Acc {
        Stream.Builder<Stream<? extends D>> builder = Stream.builder();

        void add(T t) {
            builder.accept(streamMapper.apply(t));
        }

        Acc combine(Acc another) {
            another.builder.build().forEach(builder);
            return this;
        }

        R finish() {
            return builder.build()
                    .flatMap(Function.identity()) // Here!
                    .collect(downstream);
        }
    }
    return Collector.of(Acc::new, Acc::add, Acc::combine, Acc::finish);
}

This helper method uses Collector.of and a local class Acc to accumulate the streams returned by the provided streamMapper function, which takes an element of the original stream as an argument. These streams are accumulated in a Stream.Builder that will be built when the collector's finisher function is applied.

Immediately after this stream of streams is built, it is flat-mapped with the identity function, since we only want to concatenate the streams. (I could have used a list of streams instead of a stream of streams, but I think that Stream.Builder is both very efficient and highly underused).

Acc also implements a combiner method that will merge the given Acc's stream of streams into this Acc's stream builder. This functionality will be used only if the original stream is parallel.

Here's how to use this method with your example:

Map<String, Map<String, ImmutableList<SomeClassB>>> map = someListOfClassA.stream()
    .filter(...)
    .collect(
        Collectors.groupingBy(SomeClassA::getSomeCriteriaA,
            Collectors.groupingBy(SomeClassA::getSomeCriteriaB,
                flatMapping(
                    a -> getSomeClassBsFromSomeClassA(a).stream(),
                    ImmutableList.toImmutableList()))));

---

EDIT: As @Holger indicates in the comments below, there's no need for buffering data into a stream builder when accumulating. Instead, a flat-mapping collector can be implemented by performing the flattening right in the accumulator function. Here is @Holger's own implementation of such collector, which I'm copying here verbatim with his consent:

public static <T, U, A, R> Collector<T, ?, R> flatMapping(
        Function<? super T, ? extends Stream<? extends U>> mapper,
        Collector<? super U, A, R> downstream) {

    BiConsumer<A, ? super U> acc = downstream.accumulator();
    return Collector.of(downstream.supplier(),
            (a, t) -> {
                try (Stream<? extends U> s = mapper.apply(t)) {
                    if (s != null) s.forEachOrdered(u -> acc.accept(a, u));
                }
            },
            downstream.combiner(), downstream.finisher(),
            downstream.characteristics().toArray(new Collector.Characteristics[0]));
}