forked from nivaldoh/beam
-
Notifications
You must be signed in to change notification settings - Fork 6
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge remote-tracking branch 'robertwb/rust-dofn' into rust_sdk
- Loading branch information
Showing
17 changed files
with
938 additions
and
116 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -18,6 +18,7 @@ | |
|
|
||
| pub mod pipeline; | ||
| pub mod pvalue; | ||
| pub mod serialize; | ||
| pub mod urns; | ||
|
|
||
| pub mod utils { | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,124 @@ | ||
| use std::collections::HashMap; | ||
| use std::iter::Iterator; | ||
| use std::marker::PhantomData; | ||
|
|
||
| use std::any::Any; | ||
| use std::boxed::Box; | ||
| use std::sync::Mutex; | ||
|
|
||
| use once_cell::sync::Lazy; | ||
|
|
||
| static SERIALIZED_FNS: Lazy<Mutex<HashMap<String, Box<dyn Any + Sync + Send>>>> = | ||
| Lazy::new(|| Mutex::new(HashMap::new())); | ||
|
|
||
| pub fn serialize_fn<T: Any + Sync + Send>(obj: Box<T>) -> String { | ||
| let name = format!("object{}", SERIALIZED_FNS.lock().unwrap().len()); | ||
| SERIALIZED_FNS.lock().unwrap().insert(name.to_string(), obj); | ||
| name | ||
| } | ||
|
|
||
| pub fn deserialize_fn<T: Any + Sync + Send>(name: &String) -> Option<&'static T> { | ||
| let binding = SERIALIZED_FNS.lock().unwrap(); | ||
| let untyped = binding.get(name); | ||
| let typed = match untyped { | ||
| Some(x) => x.downcast_ref::<T>(), | ||
| None => None, | ||
| }; | ||
|
|
||
| unsafe { | ||
| return std::mem::transmute::<Option<&T>, Option<&'static T>>(typed); | ||
| } | ||
| } | ||
|
|
||
| // ******* DoFn Wrappers, perhaps move elsewhere? ******* | ||
|
|
||
| // TODO: Give these start/finish_bundles, etc. | ||
| pub type GenericDoFn = | ||
| Box<dyn Fn(&dyn Any) -> Box<dyn Iterator<Item = Box<dyn Any>>> + Sync + Send>; | ||
|
|
||
| struct GenericDoFnWrapper { | ||
| func: GenericDoFn, | ||
| } | ||
|
|
||
| unsafe impl std::marker::Send for GenericDoFnWrapper {} | ||
|
|
||
| struct BoxedIter<O: Any, I: IntoIterator<Item = O>> { | ||
| typed_iter: I::IntoIter, | ||
| } | ||
|
|
||
| impl<O: Any, I: IntoIterator<Item = O>> Iterator for BoxedIter<O, I> { | ||
| type Item = Box<dyn Any>; | ||
|
|
||
| fn next(&mut self) -> Option<Box<dyn Any>> { | ||
| if let Some(x) = self.typed_iter.next() { | ||
| return Some(Box::new(x)); | ||
| } else { | ||
| return None; | ||
| } | ||
| } | ||
| } | ||
|
|
||
| pub fn to_generic_dofn<T: Any, O: Any, I: IntoIterator<Item = O> + 'static>( | ||
| func: fn(&T) -> I, | ||
| ) -> GenericDoFn { | ||
| Box::new( | ||
| move |untyped_input: &dyn Any| -> Box<dyn Iterator<Item = Box<dyn Any>>> { | ||
| let typed_input: &T = untyped_input.downcast_ref::<T>().unwrap(); | ||
| Box::new(BoxedIter::<O, I> { | ||
| typed_iter: func(typed_input).into_iter(), | ||
| }) | ||
| }, | ||
| ) | ||
| } | ||
|
|
||
| pub fn to_generic_dofn_dyn<T: Any, O: Any, I: IntoIterator<Item = O> + 'static>( | ||
| func: Box<dyn Fn(&T) -> I + Sync + Send>, | ||
| ) -> GenericDoFn { | ||
| Box::new( | ||
| move |untyped_input: &dyn Any| -> Box<dyn Iterator<Item = Box<dyn Any>>> { | ||
| let typed_input: &T = untyped_input.downcast_ref::<T>().unwrap(); | ||
| Box::new(BoxedIter::<O, I> { | ||
| typed_iter: func(typed_input).into_iter(), | ||
| }) | ||
| }, | ||
| ) | ||
| } | ||
|
|
||
| pub trait KeyExtractor: Sync + Send { | ||
| fn extract(&self, kv: &dyn Any) -> (String, Box<dyn Any + Sync + Send>); | ||
| fn recombine( | ||
| &self, | ||
| key: &String, | ||
| values: &Box<Vec<Box<dyn Any + Sync + Send>>>, | ||
| ) -> Box<dyn Any + Sync + Send>; | ||
| } | ||
|
|
||
| pub struct TypedKeyExtractor<V: Clone + Sync + Send + 'static> { | ||
| phantom_data: PhantomData<V>, | ||
| } | ||
|
|
||
| impl<V: Clone + Sync + Send + 'static> TypedKeyExtractor<V> { | ||
| pub fn default() -> Self { | ||
| Self { | ||
| phantom_data: PhantomData, | ||
| } | ||
| } | ||
| } | ||
|
|
||
| impl<V: Clone + Sync + Send + 'static> KeyExtractor for TypedKeyExtractor<V> { | ||
| fn extract(&self, kv: &dyn Any) -> (String, Box<dyn Any + Sync + Send>) { | ||
| let typed_kv = kv.downcast_ref::<(String, V)>().unwrap(); | ||
| return (typed_kv.0.clone(), Box::new(typed_kv.1.clone())); | ||
| } | ||
| fn recombine( | ||
| &self, | ||
| key: &String, | ||
| values: &Box<Vec<Box<dyn Any + Sync + Send>>>, | ||
| ) -> Box<dyn Any + Sync + Send> { | ||
| let mut typed_values: Vec<V> = Vec::new(); | ||
| for untyped_value in values.iter() { | ||
| typed_values.push(untyped_value.downcast_ref::<V>().unwrap().clone()); | ||
| } | ||
| return Box::new((key.clone(), typed_values)); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.