Skip to main content

hyphae/traits/collections/
project_many.rs

1//! Collision-safe one-to-many projection plans implementing [`MapQuery`].
2
3use 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
28/// One-to-many projection whose `(source key, local key)` output identity is
29/// collision-free across distinct source rows.
30pub 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
88/// Semantic one-to-many projection operator.
89pub trait FlatMapEntriesExt<K, V>: MapQuery<Key = K, Value = V>
90where
91    K: Hash + Eq + CellValue,
92    V: CellValue,
93{
94    /// Expand each source row into locally keyed rows. The output identity is
95    /// `(source_key, local_key)`, preventing collisions between source rows.
96    ///
97    /// Each source row must emit a `local_key` at most once; pairing with the
98    /// source key does not resolve duplicates within that row. The closure
99    /// follows [`MapQuery`]'s purity/invocation contract and may run repeatedly
100    /// or concurrently without an invocation-order guarantee.
101    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}