hyphae/traits/collections/
project_many.rs1use std::{hash::Hash, marker::PhantomData};
4
5use crate::{
6 map_query::{
7 BuildQueryRuntime, MapQuery,
8 properties::{Many, PlanProperties, Repartition},
9 },
10 subscription::SubscriptionGuard,
11 traits::{CellValue, collections::internal::map_runtime::install_map_runtime_via_query},
12};
13
14impl<S, SK, SV, LK, OV, F> PlanProperties for FlatMapEntriesPlan<S, SK, SV, LK, OV, F>
15where
16 S: MapQuery<Key = SK, Value = SV> + PlanProperties,
17 SK: Hash + Eq + CellValue,
18 SV: CellValue,
19 LK: Hash + Eq + CellValue,
20 OV: CellValue,
21 F: Fn(&SK, &SV) -> Vec<(LK, OV)> + Send + Sync + 'static,
22{
23 type Cardinality = Many;
24 type InputPartition = S::OutputPartition;
25 type OutputPartition = Repartition<(SK, LK)>;
26}
27
28pub struct FlatMapEntriesPlan<S, SK, SV, LK, OV, F>
31where
32 S: MapQuery<Key = SK, Value = SV>,
33 SK: Hash + Eq + CellValue,
34 SV: CellValue,
35 LK: Hash + Eq + CellValue,
36 OV: CellValue,
37 F: Fn(&SK, &SV) -> Vec<(LK, OV)> + Send + Sync + 'static,
38{
39 pub(crate) source: S,
40 pub(crate) f: F,
41 pub(crate) _types: PhantomData<fn() -> (SK, SV, LK, OV)>,
42}
43
44impl<S, SK, SV, LK, OV, F> BuildQueryRuntime<(SK, LK), OV>
45 for FlatMapEntriesPlan<S, SK, SV, LK, OV, F>
46where
47 S: MapQuery<Key = SK, Value = SV>,
48 SK: Hash + Eq + CellValue,
49 SV: CellValue,
50 LK: Hash + Eq + CellValue,
51 OV: CellValue,
52 F: Fn(&SK, &SV) -> Vec<(LK, OV)> + Send + Sync + 'static,
53{
54 fn build_into(
55 self,
56 cx: &mut crate::map_query::compiler::CompileContext,
57 sink: crate::map_query::BoxedMapDiffSink<(SK, LK), OV>,
58 ) -> Vec<SubscriptionGuard> {
59 let f = self.f;
60 install_map_runtime_via_query(
61 cx,
62 self.source,
63 move |source_key, value| {
64 f(source_key, value)
65 .into_iter()
66 .map(|(local_key, output)| ((source_key.clone(), local_key), output))
67 .collect()
68 },
69 sink,
70 )
71 }
72}
73
74#[allow(private_bounds)]
75impl<S, SK, SV, LK, OV, F> MapQuery for FlatMapEntriesPlan<S, SK, SV, LK, OV, F>
76where
77 S: MapQuery<Key = SK, Value = SV>,
78 SK: Hash + Eq + CellValue,
79 SV: CellValue,
80 LK: Hash + Eq + CellValue,
81 OV: CellValue,
82 F: Fn(&SK, &SV) -> Vec<(LK, OV)> + Send + Sync + 'static,
83{
84 type Key = (SK, LK);
85 type Value = OV;
86}
87
88pub trait FlatMapEntriesExt<K, V>: MapQuery<Key = K, Value = V>
90where
91 K: Hash + Eq + CellValue,
92 V: CellValue,
93{
94 fn flat_map_entries<LK, V2, F>(self, f: F) -> impl MapQuery<Key = (K, LK), Value = V2>
102 where
103 LK: Hash + Eq + CellValue,
104 V2: CellValue,
105 F: Fn(&K, &V) -> Vec<(LK, V2)> + Send + Sync + 'static,
106 {
107 FlatMapEntriesPlan {
108 source: self,
109 f,
110 _types: PhantomData,
111 }
112 }
113}
114
115impl<K, V, M> FlatMapEntriesExt<K, V> for M
116where
117 K: Hash + Eq + CellValue,
118 V: CellValue,
119 M: MapQuery<Key = K, Value = V>,
120{
121}
122
123#[cfg(test)]
124mod tests {
125 use super::*;
126 use crate::CellMap;
127
128 #[test]
129 fn flat_map_entries_emits_multiple_rows_per_source() {
130 let source = CellMap::<String, i32>::new();
131 let out = source
132 .clone()
133 .flat_map_entries(|_, value| {
134 if *value <= 0 {
135 return Vec::new();
136 }
137 vec![
138 ("a".to_string(), value * 10),
139 ("b".to_string(), value * 100),
140 ]
141 })
142 .materialize();
143
144 source.insert("x".to_string(), 2);
145 assert_eq!(out.get_value(&("x".to_string(), "a".to_string())), Some(20));
146 assert_eq!(
147 out.get_value(&("x".to_string(), "b".to_string())),
148 Some(200)
149 );
150
151 source.insert("x".to_string(), 0);
152 assert_eq!(out.get_value(&("x".to_string(), "a".to_string())), None);
153 assert_eq!(out.get_value(&("x".to_string(), "b".to_string())), None);
154 }
155
156 #[test]
157 fn flat_map_entries_scopes_local_keys_by_source() {
158 let source = CellMap::<String, i32>::new();
159 let out = source
160 .clone()
161 .flat_map_entries(|_key, value| vec![("same", value * 10)])
162 .materialize();
163
164 source.insert("a".to_string(), 1);
165 source.insert("b".to_string(), 2);
166
167 assert_eq!(out.get_value(&("a".to_string(), "same")), Some(10));
168 assert_eq!(out.get_value(&("b".to_string(), "same")), Some(20));
169 }
170}