一尘不染

对Java 8流进行分区

java

如何在Java 8 Stream上实现“分区”操作?划分是指将流分成给定大小的子流。它在某种程度上与Guava Iterators.partition()方法相同,只是希望分区是延迟评估的Streams,而不是List的Streams。


阅读 471

收藏
2020-03-19

共1个答案

一尘不染

将任意源流划分为固定大小的批次是不可能的,因为这会加重并行处理。并行处理时,你可能不知道拆分后的第一个子任务中有多少个元素,因此你无法为下一个子任务创建分区,直到完全处理第一个子任务。

但是,可以从随机访问中创建分区流List。此类功能在例如我的StreamEx图书馆中可用:

List<Type> input = Arrays.asList(...);

Stream<List<Type>> stream = StreamEx.ofSubLists(input, partitionSize);

或者,如果你真的想要流:

Stream<Stream<Type>> stream = StreamEx.ofSubLists(input, partitionSize).map(List::stream);

如果你不想依赖第三方库,则可以ofSubLists手动实现这种方法:

public static <T> Stream<List<T>> ofSubLists(List<T> source, int length) {
    if (length <= 0)
        throw new IllegalArgumentException("length = " + length);
    int size = source.size();
    if (size <= 0)
        return Stream.empty();
    int fullChunks = (size - 1) / length;
    return IntStream.range(0, fullChunks + 1).mapToObj(
        n -> source.subList(n * length, n == fullChunks ? size : (n + 1) * length));
}

该实现看起来有些长,但是它考虑了一些极端情况,例如接近MAX_VALUE列表大小。

如果你想为无序流提供并行友好的解决方案(因此你不必担心将哪些流元素合并为单个批处理),则可以使用这样的收集器(感谢@sibnick的启发):

public static <T, A, R> Collector<T, ?, R> unorderedBatches(int batchSize, 
                   Collector<List<T>, A, R> downstream) {
    class Acc {
        List<T> cur = new ArrayList<>();
        A acc = downstream.supplier().get();
    }
    BiConsumer<Acc, T> accumulator = (acc, t) -> {
        acc.cur.add(t);
        if(acc.cur.size() == batchSize) {
            downstream.accumulator().accept(acc.acc, acc.cur);
            acc.cur = new ArrayList<>();
        }
    };
    return Collector.of(Acc::new, accumulator,
            (acc1, acc2) -> {
                acc1.acc = downstream.combiner().apply(acc1.acc, acc2.acc);
                for(T t : acc2.cur) accumulator.accept(acc1, t);
                return acc1;
            }, acc -> {
                if(!acc.cur.isEmpty())
                    downstream.accumulator().accept(acc.acc, acc.cur);
                return downstream.finisher().apply(acc.acc);
            }, Collector.Characteristics.UNORDERED);
}

用法示例:

List<List<Integer>> list = IntStream.range(0,20)
                                    .boxed().parallel()
                                    .collect(unorderedBatches(3, Collectors.toList()));

结果:

[[2, 3, 4], [7, 8, 9], [0, 1, 5], [12, 13, 14], [17, 18, 19], [10, 11, 15], [6, 16]]

这样的收集器是完全线程安全的,并为顺序流生成有序批。

如果要对每个批次应用中间转换,则可以使用以下版本:

public static <T, AA, A, B, R> Collector<T, ?, R> unorderedBatches(int batchSize,
        Collector<T, AA, B> batchCollector,
        Collector<B, A, R> downstream) {
    return unorderedBatches(batchSize, 
            Collectors.mapping(list -> list.stream().collect(batchCollector), downstream));
}

例如,通过这种方式,你可以即时对每个批次中的数字求和:

List<Integer> list = IntStream.range(0,20)
        .boxed().parallel()
        .collect(unorderedBatches(3, Collectors.summingInt(Integer::intValue), 
            Collectors.toList()));
2020-03-19