Writing a custom collector that takes/skips a given number of elements – Functional style programming – extending API

198. Writing a custom collector that takes/skips a given number of elements

In Problem 195, we have written a hand of custom collectors grouped in the MyCollectors class. Now, let’s continue our journey, and let’s try to add here two more custom collectors for taking/keeping a given number of elements from the current stream.Let’s assume the following model and data:

public class Car {
  private final String brand;
  private final String fuel;
  private final int horsepower;
  …
}
List<Car> cars = List.of(
  new Car(“Chevrolet”, “diesel”, 350),
  … // 10 more
  new Car(“Lexus”, “diesel”, 300)
);

The Stream API provides an intermediate operation named limit(long n) which can be used to truncate the stream to n elements. So, if this is exactly what we want then we can use it out of the box. For instance, we can limit the resulting stream to the first 5 cars as follows:

List<Car> first5CarsLimit = cars.stream()
  .limit(5)
  .collect(Collectors.toList());

Moreover, the Stream API provides an intermediate operation named skip(long n) which can be used to skip the first n elements in the stream pipeline. For instance, we can skip the first 5 cars as follows:

List<Car> last5CarsSkip = cars.stream()
  .skip(5)
  .collect(Collectors.toList());

However, there are cases when we need to compute different things and collect only the first/last 5 results. In such cases, a custom collector is welcome.Relying on Collector.of() method (details in Problem 195), we can write a custom collector that keeps/collects the first n elements as follows (just for fun, let’s collect these n elements in an unmodifiable list):

public static <T> Collector<T, List<T>, List<T>>  
    toUnmodifiableListKeep(int max) {
  return Collector.of(ArrayList::new,
    (list, value) -> {
       if (list.size() < max) {
         list.add(value);
       }
    },
    (left, right) -> {
       left.addAll(right);
       return left;
    },
    Collections::unmodifiableList);
}

So, the supplier is ArrayList::new, the accumulator is List#add(), the combiner is List#addAll(), and the finalizer is Collections::unmodifiableList. Basically, the accumulator’s job is to accumulate elements only until the given max is reached. From that point forward, nothing gets accumulated. This way, we can keep only the first 5 cars as follows:

List<Car> first5Cars = cars.stream()
  .collect(MyCollectors.toUnmodifiableListKeep(5));

On the other hand, if we want to skip the first n elements and collect the rest then we can try to accumulate null elements until we reach the given index. From this point forward, we accumulate the real elements. In the end, the finalizer removes the part of the list containing null values (from 0 to the given index) and returns an unmodifiable list from the remaining elements (from the given index to the end):

public static <T> Collector<T, List<T>, List<T>>
    toUnmodifiableListSkip(int index) {
  return Collector.of(ArrayList::new,
    (list, value) -> {
       if (list.size() >= index) {
         list.add(value);
       } else {
         list.add(null);
       }
    },
    (left, right) -> {
       left.addAll(right);
 
       return left;
    },
    list -> Collections.unmodifiableList(
      list.subList(index, list.size())));
}

Alternatively, we can optimize this approach by using a supplier class that contains the resulting list and a counter. While the given index is not reached, we simply increase the counter. Once the given index was reached, we start to accumulate elements:

public static <T> Collector<T, ?, List<T>>
    toUnmodifiableListSkip(int index) {
  class Sublist {
    int index;
    List<T> list = new ArrayList<>();          
  }
  return Collector.of(Sublist::new,
    (sublist, value) -> {
       if (sublist.index >= index) {
         sublist.list.add(value);
       } else {
         sublist.index++;
       }
     },
     (left, right) -> {
        left.list.addAll(right.list);
        left.index = left.index + right.index;
       return left;
     },
     sublist -> Collections.unmodifiableList(sublist.list));
}

Both of these approaches can be used as in the following example:

List<Car> last5Cars = cars.stream()
  .collect(MyCollectors.toUnmodifiableListSkip(5));

Challenge yourself to implement a custom collector that collects in a given range.