Parallel rayon version of itertools ZipLongest iterator?

Hello. I'm looking for a rayon ParallelIterator which is analogous to the zip_longest iterator from itertools.

In summary ZipLongest iterates over two iterators simultaneously but only stops when both of the iterators have been exhausted.

As far as I can tell, there is no combination of existing iterators in rayon that can be used to accomplish the same task. I tried to create my own parallel iterator but rayon's iter::plumbing module still eludes me.

You could do the following:

  1. Make both iterators wrap their items in Some and have them emit None items forever once you reach the end. I.e. for each iterator, iter.map(Some).chain(repeat(None)).
  2. Zip the two iterators together to obtain something with item type (Option<T>, Option<U>).
  3. Map the iterator to have item type Option<EitherOrBoth<T, U>>.
  4. Make it stop once both iterators reach the end using while_some.
2 Likes

Thanks. I did go with that approach initially but the problem I ran into was that repeat returns a regular ParallelIterator while zip requires both arguments to be IndexedParallelIterators. We can use repeatn instead.

use rayon::iter::repeatn;
use rayon::prelude::*;

pub fn main() {
    let a = [1, 2, 3];
    let b = [4, 5, 6, 7, 8];

    zip_longest(a.into_par_iter(), b.into_par_iter()).for_each(|x| println!("{x:?}"));
}

fn zip_longest<A, B>(a: A, b: B) -> impl ParallelIterator<Item = EitherOrBoth<A::Item, B::Item>>
where
    A: IndexedParallelIterator,
    A::Item: Clone,
    B: IndexedParallelIterator,
    B::Item: Clone,
{
    let max = a.len().max(b.len());
    let a_extra = max - a.len();
    let b_extra = max - b.len();

    a.map(Some)
        .chain(repeatn(None, a_extra))
        .zip(b.map(Some).chain(repeatn(None, b_extra)))
        .map(|(a, b)| match (a, b) {
            (Some(a), Some(b)) => EitherOrBoth::Both(a, b),
            (Some(a), None) => EitherOrBoth::Left(a),
            (None, Some(b)) => EitherOrBoth::Right(b),
            (None, None) => unreachable!(),
        })
}

#[derive(Debug)]
enum EitherOrBoth<A, B> {
    Both(A, B),
    Left(A),
    Right(B),
}