partition_map with a fallible closure - how to shortcircuit?

[SuperFluffy]

I am trying to do something like the code below which works fine if I choose to ignore the Err coming out of fallible_calculation. However, I do not know how to short-circuit the calculation and turn (Vec<T>, Vec<U>) into Result<(Vec<T>, Vec<U>), E>.

let res: (Vec<T>, Vec<U>) = us
    .par_iter()
    .partition_map(|u match fallible_calculation(u) {
        Ok(t) if cond(t) => Either::Right(u.clone()),
        Ok(t) => Either::Left(t),
        Err(SpecificError) => Either::Right(u.clone()),
        Err(e) => {},
    });

How would I go about properly handling the error case? In the case of the map adapter the issue is pretty straight forward, but I don't see how can be done in this case, other than maybe collecting to a (Vec<T>, Result<Vec<U>, E>>) instead, which however doesn't seem to short circuit? That would look like so:

let res: (Vec<T>, Result<Vec<U>, E>) = us
    .par_iter()
    .partition_map(|u match fallible_calculation(u) {
        Ok(t) if cond(t) => Either::Right(Ok(u.clone())),
        Ok(t) => Either::Left(t),
        Err(SpecificError) => Either::Right(Ok(u.clone())),
        Err(e) => Either::Right(Err(e)),
    });

Are there better ways?

EDIT: I didn't realize that Result<T, E> doesn't implement Default and ParallelExtend, which means that this tactic unfortunately doesn't work either. :-(

EDIT: The fix that I could come up with to make this workable at all (albeit without short-circuiting) is this:

let (a, b, c): (Vec<T>, Vec<U>, Vec<E>) = us
    .par_iter()
    .partition_map(|u match fallible_calculation(u) {
        Ok(t) if cond(t) => Either::Right(Ok(u.clone())),
        Ok(t) => Either::Left(t),
        Err(SpecificError) => Either::Right(Ok(u.clone())),
        Err(e) => Either::Right(Err(e)),
    });

if c.len() > 0 {
    Err(c.swap_remove(0))?
}

It's not particularly elegant, but works.

[cuviper]

I do not know how to short-circuit the calculation and turn (Vec<T>, Vec<U>) into Result<(Vec<T>, Vec<U>), E>.

I think you could do this today by mimicking FromParallelIterator<Result<T, E>> for Result<C, E>:

rayon/src/result.rs

Lines 88 to 132 in 4a26ac4

	/// Collect an arbitrary `Result`-wrapped collection.
	///
	/// If any item is `Err`, then all previous `Ok` items collected are
	/// discarded, and it returns that error. If there are multiple errors, the
	/// one returned is not deterministic.
	impl<C, T, E> FromParallelIterator<Result<T, E>> for Result<C, E>
	where
	C: FromParallelIterator<T>,
	T: Send,
	E: Send,
	{
	fn from_par_iter<I>(par_iter: I) -> Self
	where
	I: IntoParallelIterator<Item = Result<T, E>>,
	{
	fn ok<T, E>(saved: &Mutex<Option<E>>) -> impl Fn(Result<T, E>) -> Option<T> + '_ {
	move |item| match item {
	Ok(item) => Some(item),
	Err(error) => {
	// We don't need a blocking `lock()`, as anybody
	// else holding the lock will also be writing
	// `Some(error)`, and then ours is irrelevant.
	if let Ok(mut guard) = saved.try_lock() {
	if guard.is_none() {
	*guard = Some(error);
	}
	}
	None
	}
	}
	}
	let saved_error = Mutex::new(None);
	let collection = par_iter
	.into_par_iter()
	.map(ok(&saved_error))
	.while_some()
	.collect();
	match saved_error.into_inner().unwrap() {
	Some(error) => Err(error),
	None => Ok(collection),
	}
	}
	}

Your map in this case would create Option<Either<L, R>>, saving the error aside. Then .while_some() to handle short-circuiting, and .partition_map(identity) to split your left/right collections.

maybe collecting to a (Vec<T>, Result<Vec<U>, E>>) instead, which however doesn't seem to short circuit?

That's controlled by full, which only stops when both sides say so:

rayon/src/iter/unzip.rs

Lines 388 to 391 in 4a26ac4

	fn full(&self) -> bool {
	// don't stop until everyone is full
	self.left.full() && self.right.full()
	}

edit: but you're right, Default and ParallelExtend are also missing for Result, and in particular we can't do anything about Default from here.

---

A while ago, I was also experimenting with various FromParallelIterator for tuples as a followup to #604. That included FromParallelIterator<Either<L, R>> for (A, B) as a partition_map. I think with that in place, you could map your fallible calculation to Result<Either<T, U>, E>, then just collect to Result<(Vec<T>, Vec<U>), E> as you wanted. Now that you've presented a use case, I'll see if I can dig up that branch and make a pull request.

#721 is somewhat related too.

[cuviper]

A while ago, I was also experimenting with various FromParallelIterator for tuples as a followup to #604. That included FromParallelIterator<Either<L, R>> for (A, B) as a partition_map. I think with that in place, you could map your fallible calculation to Result<Either<T, U>, E>, then just collect to Result<(Vec<T>, Vec<U>), E> as you wanted. Now that you've presented a use case, I'll see if I can dig up that branch and make a pull request.

See #802, if you'd like to try it out...