krabka_client_streams/dsl/
session_windowed_cogrouped.rs

1//! `SessionWindowedCogroupedStream<K, VOut>`: session-windowed KIP-150 cogroup.
2//!
3//! `CogroupedKStream::windowed_by_session(SessionWindows)` builds it. The
4//! terminal `aggregate_explicit` produces a `KTable<Windowed<K>, VOut>` over a
5//! shared session store. Unlike the time-windowed and sliding-windowed variants,
6//! the session aggregate needs a **merger** that combines two sessions when the
7//! runtime merges them.
8
9use std::{any::Any, cell::RefCell, marker::PhantomData, rc::Rc, sync::Arc};
10
11use crate::{
12    dsl::{
13        builder::InternalStreamsBuilder,
14        cogrouped::{
15            CogroupInput, CogroupKind, CogroupSpec, CogroupedKStream, StoreRegistrarFn,
16            lower_cogroup,
17        },
18        config::Materialized,
19        kgrouped::mint_store_name,
20        ktable::KTable,
21        names,
22        session_windowed_kgrouped::session_suppress_factory,
23        windows::{SessionWindowedSerde, SessionWindows, Windowed},
24    },
25    processor::serde::Serde,
26};
27
28impl<K, VOut> CogroupedKStream<K, VOut>
29where
30    K: Any + Send + Sync + Clone,
31    VOut: Any + Send + Sync + Clone,
32{
33    /// `windowedBy(SessionWindows)`: switch to a session-windowed cogroup
34    /// (KIP-150).
35    ///
36    /// Unlike the time-windowed and sliding-windowed cogroup, the terminal
37    /// `aggregate_explicit` needs a `merger`. The merger combines two session
38    /// aggregates when the runtime merges sessions after an inactivity gap
39    /// expires.
40    #[must_use]
41    pub fn windowed_by_session(
42        self,
43        windows: SessionWindows,
44    ) -> SessionWindowedCogroupedStream<K, VOut> {
45        SessionWindowedCogroupedStream {
46            builder: self.builder,
47            inputs: self.inputs,
48            windows,
49            _pd: PhantomData,
50        }
51    }
52}
53
54/// Handle produced by [`CogroupedKStream::windowed_by_session`]. The terminal
55/// session-windowed aggregation consumes it.
56pub struct SessionWindowedCogroupedStream<K, VOut> {
57    builder: Rc<RefCell<InternalStreamsBuilder>>,
58    inputs: Vec<CogroupInput<K, VOut>>,
59    windows: SessionWindows,
60    _pd: PhantomData<fn() -> (K, VOut)>,
61}
62
63impl<K, VOut> SessionWindowedCogroupedStream<K, VOut>
64where
65    K: Any + Send + Sync + Clone,
66    VOut: Any + Send + Sync + Clone,
67{
68    /// Session-windowed terminal aggregation. It produces a
69    /// `KTable<Windowed<K>, VOut>`.
70    ///
71    /// The `merger` combines two session aggregates when sessions are merged
72    /// (required for session windows — no default merger exists).
73    pub fn aggregate_explicit<KS, VS, I, M>(
74        self,
75        init: I,
76        merger: M,
77        materialized: impl Into<Materialized<KS, VS>>,
78    ) -> KTable<Windowed<K>, VOut, SessionWindowedSerde<KS>, VS>
79    where
80        KS: Serde<K> + Clone + 'static,
81        VS: Serde<VOut> + Clone + 'static,
82        I: Fn() -> VOut + Send + Sync + 'static,
83        M: Fn(&K, VOut, VOut) -> VOut + Send + Sync + 'static,
84    {
85        let materialized = materialized.into();
86        let store_name = mint_store_name(&self.builder, &materialized, names::AGGREGATE_STORE);
87        let suppress_factory = session_suppress_factory::<K, VOut, KS, VS>(
88            materialized.key_serde.clone(),
89            materialized.value_serde.clone(),
90        );
91        let Materialized {
92            key_serde,
93            value_serde,
94            logging,
95            caching,
96            ..
97        } = materialized;
98        let spec = CogroupSpec::<K, VOut> {
99            kind: CogroupKind::Session(self.windows),
100            init: Arc::new(init),
101            merger: Some(Arc::new(merger)),
102        };
103        let ks = key_serde.clone();
104        let vs = value_serde.clone();
105        let store_for_reg = store_name.clone();
106        let gap = self.windows.gap;
107        let grace = self.windows.grace;
108        let registrar: StoreRegistrarFn = Box::new(move |state, procs| {
109            state.topology.add_session_store::<K, VOut, KS, VS>(
110                store_for_reg.clone(),
111                ks.clone(),
112                vs.clone(),
113                gap,
114                grace,
115                procs,
116            );
117            state.topology.mark_store_caching(&store_for_reg, caching);
118        });
119        let merge_id = lower_cogroup::<K, VOut, Windowed<K>>(
120            &self.builder,
121            self.inputs,
122            &store_name,
123            &spec,
124            logging,
125            registrar,
126        );
127        KTable::new(
128            Rc::clone(&self.builder),
129            merge_id,
130            Some(store_name),
131            None,
132            SessionWindowedSerde::new(key_serde),
133            value_serde,
134        )
135        .with_window_grace(Some(self.windows.grace))
136        .with_suppress_factory(Some(suppress_factory))
137    }
138}
139
140#[cfg(test)]
141mod caching_tests {
142    use assert2::check;
143    use krabka_units::prelude::*;
144
145    use crate::{
146        BufferConfig, I64Serde, Materialized, Produced, SessionWindows, StringSerde, Suppressed,
147        dsl::{StreamsBuilder, windows::SessionWindowedSerde},
148        store::backend::StoreBackend,
149    };
150
151    #[test]
152    fn session_windowed_cogroup_marks_store_cached() {
153        let b = StreamsBuilder::new();
154        let g1 = b.stream::<String, String>(["in1"]).group_by_key();
155        let g2 = b.stream::<String, String>(["in2"]).group_by_key();
156        g1.cogroup::<i64, _>(|_k, v: &String, acc| {
157            acc + i64::try_from(v.len()).unwrap_or(i64::MAX)
158        })
159        .cogroup(g2, |_k, _v: &String, acc| acc + 1)
160        .windowed_by_session(SessionWindows::of_inactivity_gap(millis(100)))
161        .aggregate_explicit(
162            || 0i64,
163            |_k: &String, a: i64, b: i64| a + b,
164            Materialized::with(StringSerde, I64Serde).as_store("cg"),
165        )
166        .suppress(Suppressed::until_window_closes(BufferConfig::unbounded()))
167        .to_stream()
168        .to_explicit(
169            "out",
170            Produced::with(SessionWindowedSerde::new(StringSerde), I64Serde),
171        );
172        let built = b.build("app").unwrap();
173        let g = pollster::block_on(built.instantiate(&StoreBackend::InMemory, "app", kibibytes(1)))
174            .unwrap();
175        check!(g.cache_owner.contains_key("cg"));
176    }
177
178    #[test]
179    fn session_windowed_cogroup_uncached_when_off() {
180        let b = StreamsBuilder::new();
181        let g1 = b.stream::<String, String>(["in1"]).group_by_key();
182        let g2 = b.stream::<String, String>(["in2"]).group_by_key();
183        g1.cogroup::<i64, _>(|_k, v: &String, acc| {
184            acc + i64::try_from(v.len()).unwrap_or(i64::MAX)
185        })
186        .cogroup(g2, |_k, _v: &String, acc| acc + 1)
187        .windowed_by_session(SessionWindows::of_inactivity_gap(millis(100)))
188        .aggregate_explicit(
189            || 0i64,
190            |_k: &String, a: i64, b: i64| a + b,
191            Materialized::with(StringSerde, I64Serde)
192                .as_store("cg")
193                .with_caching(false),
194        )
195        .to_stream()
196        .to_explicit(
197            "out",
198            Produced::with(SessionWindowedSerde::new(StringSerde), I64Serde),
199        );
200        let built = b.build("app").unwrap();
201        let g = pollster::block_on(built.instantiate(&StoreBackend::InMemory, "app", kibibytes(1)))
202            .unwrap();
203        check!(!g.cache_owner.contains_key("cg"));
204    }
205}
krabka_client_streams/dsl/session_windowed_cogrouped.rs

krabka_client_streams/dsl/
session_windowed_cogrouped.rs
