com.dua3.utility.lang.BatchCollector Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of utility Show documentation
Java utilities (core)
There is a newer version: 15.0.2
package com.dua3.utility.lang;

import com.dua3.cabe.annotations.Nullable;
import com.dua3.utility.data.Pair;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Deque;
import java.util.EnumSet;
import java.util.List;
import java.util.Objects;
import java.util.Set;
import java.util.function.BiConsumer;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collector;

/**
 * A collector that puts subsequent items into batches. A new batch is started by generating a new batch key.
 *
 * @param  the item type
 * @param  the key type
 */
public class BatchCollector implements Collector>>, List>>> {
    private final Function keyMapper;
    private final K defaultKey;

    /**
     * Constructor.
     *
     * @param keyMapper mapping that maps each item to the grouping key
     */
    public BatchCollector(Function keyMapper) {
        this(keyMapper, null);
    }

    /**
     * Create a new BatchCollector.
     * 
     * For each item in the stream, a key is determined by applying the keyMapper. If the generated key is null, or
     * equals the last item's key, the item is added to the current batch. If not, a new batch is created and the
     * item added.
     *
     * @param keyMapper  the key mapper
     * @param defaultKey the default key
     */
    public BatchCollector(Function keyMapper, @Nullable K defaultKey) {
        this.keyMapper = keyMapper;
        this.defaultKey = defaultKey;
    }

    @Override
    public Supplier>>> supplier() {
        return ArrayDeque::new;
    }

    @Override
    public BiConsumer>>, T> accumulator() {
        return (accu, item) -> {
            K key = keyMapper.apply(item);

            List bucket;
            if (accu.isEmpty() || (!Objects.equals(key, accu.peekLast().first()))) {
                bucket = new ArrayList<>();
                accu.addLast(Pair.of(key == null ? defaultKey : key, bucket));
            } else {
                bucket = accu.peekLast().second();
            }

            bucket.add(item);
        };
    }

    @Override
    public BinaryOperator>>> combiner() {
        return (left, right) -> {
            left.addAll(right);
            return left;
        };
    }

    @Override
    public Function>>, List>>> finisher() {
        return ArrayList::new;
    }

    @Override
    public Set characteristics() {
        return EnumSet.noneOf(Characteristics.class);
    }
}