public static < T > Collector < T , ArrayList < List < T > > , Stream < List < T > > > inNumBatches ( int numBatches ) { // Keep the position separate from the possible threads to enable parallel // collection. AtomicInteger nextPos = new AtomicInteger ( 0 ) ; AtomicInteger numTotal = new AtomicInteger ( 0 ) ; return Collector . of ( // instantiation ( ) -> { ArrayList < List < T >> batches = new ArrayList <> ( numBatches ) ; for ( int i = 0 ; i < numBatches ; ++ i ) { batches . add ( new ArrayList <> ( ) ) ; } return batches ; } , // accumulator ( batches , item ) -> { int pos = nextPos . getAndUpdate ( i -> ( ++ i ) % numBatches ) ; batches . get ( pos ) . add ( item ) ; numTotal . incrementAndGet ( ) ; } , // combiner ( a , b ) -> { // Merge the two lists so the batches matches the order // of the non-parallel inBatchesOf with (a1..an) + (b1..bn) // as the set of items. It's not extremely efficient, but // works fine as this is not optimized for parallel streams. for ( int i = 0 ; i < numBatches ; ++ i ) { List < T > al = a . get ( i ) ; List < T > bl = b . get ( i ) ; al . addAll ( bl ) ; } return a ; } , // finalizer batches -> { if ( numTotal . get ( ) < numBatches ) { return batches . subList ( 0 , numTotal . get ( ) ) . stream ( ) ; } return batches . stream ( ) ; } ) ; }
Collect into N batches of approximate equal size . Creates a stream of lists as response .