diff --git a/sdks/rust/src/transforms/create.rs b/sdks/rust/src/transforms/create.rs index a9ec0d1ceeb80..c7cf25cf88d19 100644 --- a/sdks/rust/src/transforms/create.rs +++ b/sdks/rust/src/transforms/create.rs @@ -46,8 +46,8 @@ impl PTransform for Create { input .clone() .apply(Impulse::new()) - .apply(ParDo::from_dyn_flat_map(Box::new(move |_x| -> Vec { - elements.to_vec() - }))) + .apply(ParDo::from_flatmap_with_context(Box::new( + move |_x| -> Vec { elements.to_vec() }, + ))) } } diff --git a/sdks/rust/src/transforms/pardo.rs b/sdks/rust/src/transforms/pardo.rs index 1ddcc31b3ae91..1d7099b62faf9 100644 --- a/sdks/rust/src/transforms/pardo.rs +++ b/sdks/rust/src/transforms/pardo.rs @@ -52,16 +52,27 @@ pub struct ParDo { // TODO: Is the Sync + Send stuff needed? impl ParDo { // TODO: These should correspond to methods on PCollection (but not on PValue). + /// Creates a ParDo from a function that maps a single input element to a single output element. + /// The function must not have any context. See [Understanding Closures in Rust](https://medium.com/swlh/understanding-closures-in-rust-21f286ed1759) for detail. pub fn from_map(func: fn(&T) -> O) -> Self { - Self::from_dyn_map(Box::new(func)) + Self::from_map_with_context(Box::new(func)) } - pub fn from_dyn_map(func: Box O + Send + Sync>) -> Self { - Self::from_dyn_flat_map(Box::new(move |x: &T| -> Vec { vec![func(x)] })) + + /// Creates a ParDo from a function that maps a single input element to a single output element. + /// The function may have an immutable context. See [Understanding Closures in Rust](https://medium.com/swlh/understanding-closures-in-rust-21f286ed1759) for detail. + pub fn from_map_with_context(func: Box O + Send + Sync>) -> Self { + Self::from_flatmap_with_context(Box::new(move |x: &T| -> Vec { vec![func(x)] })) } - pub fn from_flat_map + 'static>(func: fn(&T) -> I) -> Self { - Self::from_dyn_flat_map(Box::new(func)) + + /// Creates a ParDo from a function that maps a single input element to a collection of output elements. + /// The function must not have any context. See [Understanding Closures in Rust](https://medium.com/swlh/understanding-closures-in-rust-21f286ed1759) for detail. + pub fn from_flatmap + 'static>(func: fn(&T) -> I) -> Self { + Self::from_flatmap_with_context(Box::new(func)) } - pub fn from_dyn_flat_map + 'static>( + + /// Creates a ParDo from a function that maps a single input element to a collection of output elements. + /// The function may have an immutable context. See [Understanding Closures in Rust](https://medium.com/swlh/understanding-closures-in-rust-21f286ed1759) for detail. + pub fn from_flatmap_with_context + 'static>( func: Box I + Send + Sync>, ) -> Self { Self { diff --git a/sdks/rust/src/transforms/testing.rs b/sdks/rust/src/transforms/testing.rs index eddd4390a59a5..1f2d7ad880b4b 100644 --- a/sdks/rust/src/transforms/testing.rs +++ b/sdks/rust/src/transforms/testing.rs @@ -55,7 +55,7 @@ impl PTransform for AssertEqu KV::new("".to_string(), x.clone()) })) .apply(GroupByKey::default()) - .apply(ParDo::from_dyn_map(Box::new( + .apply(ParDo::from_map_with_context(Box::new( move |kvs: &KV>>| { let mut actual: Vec = kvs .as_values()