From d04a555338398f665f7cff9aef7ae42f59bc4b9b Mon Sep 17 00:00:00 2001 From: Nathan Scheinkman Date: Fri, 19 Jul 2024 22:59:32 -0700 Subject: [PATCH 1/3] Add cartesian_product to ParallelIterator --- src/iter/cartesian_product.rs | 48 +++++++++++++++++++++++++++++++ src/iter/mod.rs | 27 ++++++++++++++++++ tests/cartesian_product.rs | 53 +++++++++++++++++++++++++++++++++++ 3 files changed, 128 insertions(+) create mode 100644 src/iter/cartesian_product.rs create mode 100644 tests/cartesian_product.rs diff --git a/src/iter/cartesian_product.rs b/src/iter/cartesian_product.rs new file mode 100644 index 000000000..8b2b865da --- /dev/null +++ b/src/iter/cartesian_product.rs @@ -0,0 +1,48 @@ +use super::plumbing::*; +use super::repeat::*; +use super::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator}; + +/// `CartesianProduct` is an iterator that combines `i` and `j` into a single iterator that +/// iterates over the cartesian product of the elements of `i` and `j`. This struct is created by +/// the [`cartesian_product()`] method on [`ParallelIterator`] +#[must_use = "iterator adaptors are lazy and do nothing unless consumed"] +#[derive(Debug, Clone, Copy)] +pub struct CartesianProduct +where + I: ParallelIterator, + J: ParallelIterator, +{ + i: I, + j: J, +} + +impl CartesianProduct +where + I: ParallelIterator, + J: IndexedParallelIterator, +{ + /// Creates a new `CartesianProduct` iterator. + pub(super) fn new(i: I, j: J) -> Self { + CartesianProduct { i, j } + } +} + +impl ParallelIterator for CartesianProduct +where + I: ParallelIterator, + J: IndexedParallelIterator + Clone + Sync, + I::Item: Clone + Send, +{ + type Item = (I::Item, J::Item); + + fn drive_unindexed(self, consumer: C) -> C::Result + where + C: UnindexedConsumer, + { + self.i + .into_par_iter() + .map(|i_item| repeat(i_item.clone()).zip(self.j.clone())) + .flatten() + .drive_unindexed(consumer) + } +} diff --git a/src/iter/mod.rs b/src/iter/mod.rs index 963457229..02b365da1 100644 --- a/src/iter/mod.rs +++ b/src/iter/mod.rs @@ -104,6 +104,7 @@ mod test; // can be readily distinguished. mod blocks; +mod cartesian_product; mod chain; mod chunks; mod cloned; @@ -163,6 +164,7 @@ mod zip_eq; pub use self::{ blocks::{ExponentialBlocks, UniformBlocks}, + cartesian_product::CartesianProduct, chain::Chain, chunks::Chunks, cloned::Cloned, @@ -606,6 +608,31 @@ pub trait ParallelIterator: Sized + Send { Map::new(self, map_op) } + /// Produces a new iterator that iterates over the cartesian product of the element sets of + /// `self` and `CI`. + /// + /// Iterator element type is `(Self::Item, CI::Item)` + /// + /// # Examples + /// + /// ``` + /// use rayon::prelude::*; + /// + /// let set: Vec<_> = (0..2).into_par_iter().cartesian_product(0..2).collect(); + /// assert!(set.contains(&(0, 0))); + /// assert!(set.contains(&(0, 1))); + /// assert!(set.contains(&(1, 0))); + /// assert!(set.contains(&(1, 1))); + /// ``` + fn cartesian_product(self, other: CI) -> CartesianProduct + where + Self::Item: Clone + Send, + CI: IntoParallelIterator, + CI::Iter: IndexedParallelIterator + Clone + Sync, + { + CartesianProduct::new(self, other.into_par_iter()) + } + /// Applies `map_op` to the given `init` value with each item of this /// iterator, producing a new iterator with the results. /// diff --git a/tests/cartesian_product.rs b/tests/cartesian_product.rs new file mode 100644 index 000000000..71dd787fa --- /dev/null +++ b/tests/cartesian_product.rs @@ -0,0 +1,53 @@ +use rayon::prelude::*; +use std::collections::HashSet; + +#[test] +fn cartesian_ranges() { + let a: Vec<(usize, usize)> = (0..3).into_par_iter().cartesian_product(0..3).collect(); + let b: HashSet<(usize, usize)> = HashSet::from_iter(a.into_iter()); + assert!(b.contains(&(0, 0))); + assert!(b.contains(&(0, 1))); + assert!(b.contains(&(0, 2))); + assert!(b.contains(&(1, 0))); + assert!(b.contains(&(1, 1))); + assert!(b.contains(&(1, 2))); + assert!(b.contains(&(2, 0))); + assert!(b.contains(&(2, 1))); + assert!(b.contains(&(2, 2))); +} + +#[test] +fn cartesian_vecs() { + let a: Vec<(f64, f64)> = vec![0.1, 1.2, 2.4] + .into_par_iter() + .cartesian_product(vec![4.8, 16.32, 32.64]) + .collect(); + assert!(a.contains(&(0.1, 4.8))); + assert!(a.contains(&(0.1, 16.32))); + assert!(a.contains(&(0.1, 32.64))); + assert!(a.contains(&(1.2, 4.8))); + assert!(a.contains(&(1.2, 16.32))); + assert!(a.contains(&(1.2, 32.64))); + assert!(a.contains(&(2.4, 4.8))); + assert!(a.contains(&(2.4, 16.32))); + assert!(a.contains(&(2.4, 32.64))); +} + +#[test] +fn cartesian_chain() { + let a: Vec<(usize, usize, usize)> = (0..2) + .into_par_iter() + .cartesian_product(0..2) + .cartesian_product(0..2) + .map(|((a, b), c)| (a, b, c)) + .collect(); + let b: HashSet<(usize, usize, usize)> = HashSet::from_iter(a.into_iter()); + assert!(b.contains(&(0, 0, 0))); + assert!(b.contains(&(0, 0, 1))); + assert!(b.contains(&(0, 1, 0))); + assert!(b.contains(&(0, 1, 1))); + assert!(b.contains(&(1, 0, 0))); + assert!(b.contains(&(1, 0, 1))); + assert!(b.contains(&(1, 1, 0))); + assert!(b.contains(&(1, 1, 1))); +} From 1520e406aa289302743a5298f8d638d18f4db232 Mon Sep 17 00:00:00 2001 From: Nathan Scheinkman Date: Sat, 20 Jul 2024 00:16:55 -0700 Subject: [PATCH 2/3] Cleanup cartesian_product implementation --- src/iter/cartesian_product.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/iter/cartesian_product.rs b/src/iter/cartesian_product.rs index 8b2b865da..bd1c792e6 100644 --- a/src/iter/cartesian_product.rs +++ b/src/iter/cartesian_product.rs @@ -1,6 +1,6 @@ use super::plumbing::*; use super::repeat::*; -use super::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator}; +use super::{IndexedParallelIterator, ParallelIterator}; /// `CartesianProduct` is an iterator that combines `i` and `j` into a single iterator that /// iterates over the cartesian product of the elements of `i` and `j`. This struct is created by @@ -40,9 +40,7 @@ where C: UnindexedConsumer, { self.i - .into_par_iter() - .map(|i_item| repeat(i_item.clone()).zip(self.j.clone())) - .flatten() + .flat_map(|i_item| repeat(i_item).zip(self.j.clone())) .drive_unindexed(consumer) } } From 323b4fb5d1f67c0087d6c90f327a1f2b5154cdf3 Mon Sep 17 00:00:00 2001 From: Nathan Scheinkman Date: Sat, 20 Jul 2024 00:44:11 -0700 Subject: [PATCH 3/3] Remove 'Send' restriction on cartesian_product item --- src/iter/cartesian_product.rs | 2 +- src/iter/mod.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/iter/cartesian_product.rs b/src/iter/cartesian_product.rs index bd1c792e6..7cf4be1ab 100644 --- a/src/iter/cartesian_product.rs +++ b/src/iter/cartesian_product.rs @@ -31,7 +31,7 @@ impl ParallelIterator for CartesianProduct where I: ParallelIterator, J: IndexedParallelIterator + Clone + Sync, - I::Item: Clone + Send, + I::Item: Clone, { type Item = (I::Item, J::Item); diff --git a/src/iter/mod.rs b/src/iter/mod.rs index 02b365da1..9d8c55714 100644 --- a/src/iter/mod.rs +++ b/src/iter/mod.rs @@ -626,7 +626,7 @@ pub trait ParallelIterator: Sized + Send { /// ``` fn cartesian_product(self, other: CI) -> CartesianProduct where - Self::Item: Clone + Send, + Self::Item: Clone, CI: IntoParallelIterator, CI::Iter: IndexedParallelIterator + Clone + Sync, {