krabka_client_streams/dsl/
session_windowed_kgrouped.rs

1//! `SessionWindowedKGroupedStream<K,V>`: the handle between
2//! `KGroupedStream::windowed_by_session(SessionWindows)` and a terminal session
3//! aggregation (`count`/`reduce`/`aggregate`).
4//!
5//! This type is the session analogue of
6//! [`crate::dsl::windowed_kgrouped::TimeWindowedKGroupedStream`]. It holds the
7//! same grouped lineage plus the [`SessionWindows`] spec. Its terminal ops emit
8//! `Windowed<K>` keys and materialize a **session store** with
9//! `add_session_store`. The result is a `KTable<Windowed<K>, _>` with a
10//! changelog-backed session store.
11use std::{any::Any, cell::RefCell, marker::PhantomData, rc::Rc};
12
13use crate::{
14    dsl::{
15        builder::InternalStreamsBuilder,
16        config::Materialized,
17        emit::EmitStrategy,
18        graph::{GraphNodeKind, LowerState, NodeId},
19        kgrouped::{KGroupedStream, RepartitionLowerFn, mint_store_name},
20        ktable::{KTable, SuppressStoreFactory},
21        names,
22        processors::session_aggregate::{
23            KStreamSessionAggregateProcessor, KStreamSessionReduceProcessor,
24        },
25        windows::{SessionWindowedSerde, SessionWindows, Windowed},
26    },
27    processor::serde::{DefaultSerde, Serde},
28    topology::NodeHandle,
29};
30
31/// Handle produced by [`KGroupedStream::windowed_by_session`].
32///
33/// [`KGroupedStream::windowed_by_session`]: crate::dsl::kgrouped::KGroupedStream::windowed_by_session
34pub struct SessionWindowedKGroupedStream<K, V> {
35    builder: Rc<RefCell<InternalStreamsBuilder>>,
36    parent: NodeId,
37    key_changing_upstream: bool,
38    #[allow(dead_code)]
39    grouped_name: Option<String>,
40    repartition_lower: Option<RepartitionLowerFn>,
41    windows: SessionWindows,
42    emit: EmitStrategy,
43    _pd: PhantomData<fn() -> (K, V)>,
44}
45
46impl<K, V> SessionWindowedKGroupedStream<K, V>
47where
48    K: Any + Send + Sync + Clone,
49    V: Any + Send + Sync + Clone,
50{
51    pub(crate) fn new(
52        builder: Rc<RefCell<InternalStreamsBuilder>>,
53        parent: NodeId,
54        key_changing_upstream: bool,
55        grouped_name: Option<String>,
56        repartition_lower: Option<RepartitionLowerFn>,
57        windows: SessionWindows,
58    ) -> Self {
59        Self {
60            builder,
61            parent,
62            key_changing_upstream,
63            grouped_name,
64            repartition_lower,
65            windows,
66            emit: EmitStrategy::default(),
67            _pd: PhantomData,
68        }
69    }
70
71    /// Emit on every update (default) or only on window close (KIP-825).
72    /// Records whose merged session has already closed at the current stream
73    /// time are dropped for both strategies.
74    #[must_use]
75    pub fn emit_strategy(mut self, emit: EmitStrategy) -> Self {
76        self.emit = emit;
77        self
78    }
79
80    /// `count`: count records per session into a `KTable<Windowed<K>, i64>`.
81    pub fn count_explicit<KS, VS>(
82        self,
83        materialized: impl Into<Materialized<KS, VS>>,
84    ) -> KTable<Windowed<K>, i64, SessionWindowedSerde<KS>, VS>
85    where
86        KS: Serde<K> + Clone + 'static,
87        VS: Serde<i64> + Clone + 'static,
88    {
89        let materialized = materialized.into();
90        let store_name = mint_store_name(&self.builder, &materialized, names::AGGREGATE_STORE);
91        // JVM `count` burns an extra store-name counter index when unnamed.
92        if materialized.store_name.is_none() {
93            self.builder
94                .borrow_mut()
95                .new_processor_name(names::AGGREGATE_STORE);
96        }
97        self.lower_aggregate::<KS, VS, i64, _, _, _>(
98            materialized,
99            store_name,
100            || 0i64,
101            |_k: &K, _v: &V, acc: i64| acc + 1,
102            |_k: &K, a: i64, b: i64| a + b,
103        )
104    }
105
106    /// `aggregate`: general session aggregation with `init`, `agg`, and the
107    /// session `merger`. The merger combines two session aggregates on a merge.
108    pub fn aggregate_explicit<KS, VS, VA, I, A, M>(
109        self,
110        init: I,
111        agg: A,
112        merger: M,
113        materialized: impl Into<Materialized<KS, VS>>,
114    ) -> KTable<Windowed<K>, VA, SessionWindowedSerde<KS>, VS>
115    where
116        VA: Any + Send + Sync + Clone,
117        KS: Serde<K> + Clone + 'static,
118        VS: Serde<VA> + Clone + 'static,
119        I: Fn() -> VA + Clone + Send + Sync + 'static,
120        A: Fn(&K, &V, VA) -> VA + Clone + Send + Sync + 'static,
121        M: Fn(&K, VA, VA) -> VA + Clone + Send + Sync + 'static,
122    {
123        let materialized = materialized.into();
124        let store_name = mint_store_name(&self.builder, &materialized, names::AGGREGATE_STORE);
125        self.lower_aggregate::<KS, VS, VA, I, A, M>(materialized, store_name, init, agg, merger)
126    }
127
128    /// `reduce`: combine values per session with `reducer` into a
129    /// `KTable<Windowed<K>, V>`.
130    pub fn reduce_explicit<KS, VS, R>(
131        self,
132        reducer: R,
133        materialized: impl Into<Materialized<KS, VS>>,
134    ) -> KTable<Windowed<K>, V, SessionWindowedSerde<KS>, VS>
135    where
136        KS: Serde<K> + Clone + 'static,
137        VS: Serde<V> + Clone + 'static,
138        R: Fn(&V, &V) -> V + Clone + Send + Sync + 'static,
139    {
140        let materialized = materialized.into();
141        let store_name = mint_store_name(&self.builder, &materialized, names::REDUCE_STORE);
142        self.lower_reduce::<KS, VS, R>(materialized, store_name, reducer)
143    }
144
145    pub fn count(
146        self,
147        store_name: impl Into<String>,
148    ) -> KTable<
149        Windowed<K>,
150        i64,
151        SessionWindowedSerde<<K as DefaultSerde>::Serde>,
152        crate::processor::serde::I64Serde,
153    >
154    where
155        K: DefaultSerde,
156        <K as DefaultSerde>::Serde: Serde<K> + Clone,
157    {
158        self.count_explicit(
159            Materialized::with(
160                <K as DefaultSerde>::Serde::default(),
161                crate::processor::serde::I64Serde,
162            )
163            .as_store(store_name),
164        )
165    }
166
167    pub fn reduce<R>(
168        self,
169        reducer: R,
170        store_name: impl Into<String>,
171    ) -> KTable<
172        Windowed<K>,
173        V,
174        SessionWindowedSerde<<K as DefaultSerde>::Serde>,
175        <V as DefaultSerde>::Serde,
176    >
177    where
178        K: DefaultSerde,
179        V: DefaultSerde,
180        <K as DefaultSerde>::Serde: Serde<K> + Clone,
181        <V as DefaultSerde>::Serde: Serde<V> + Clone,
182        R: Fn(&V, &V) -> V + Clone + Send + Sync + 'static,
183    {
184        self.reduce_explicit(
185            reducer,
186            Materialized::with(
187                <K as DefaultSerde>::Serde::default(),
188                <V as DefaultSerde>::Serde::default(),
189            )
190            .as_store(store_name),
191        )
192    }
193
194    pub fn aggregate<VA, I, A, M>(
195        self,
196        init: I,
197        agg: A,
198        merger: M,
199        store_name: impl Into<String>,
200    ) -> KTable<
201        Windowed<K>,
202        VA,
203        SessionWindowedSerde<<K as DefaultSerde>::Serde>,
204        <VA as DefaultSerde>::Serde,
205    >
206    where
207        VA: DefaultSerde + Any + Send + Sync + Clone,
208        K: DefaultSerde,
209        <K as DefaultSerde>::Serde: Serde<K> + Clone,
210        <VA as DefaultSerde>::Serde: Serde<VA> + Clone,
211        I: Fn() -> VA + Clone + Send + Sync + 'static,
212        A: Fn(&K, &V, VA) -> VA + Clone + Send + Sync + 'static,
213        M: Fn(&K, VA, VA) -> VA + Clone + Send + Sync + 'static,
214    {
215        self.aggregate_explicit(
216            init,
217            agg,
218            merger,
219            Materialized::with(
220                <K as DefaultSerde>::Serde::default(),
221                <VA as DefaultSerde>::Serde::default(),
222            )
223            .as_store(store_name),
224        )
225    }
226
227    fn lower_aggregate<KS, VS, VA, I, A, M>(
228        mut self,
229        materialized: Materialized<KS, VS>,
230        store_name: String,
231        init: I,
232        agg: A,
233        merger: M,
234    ) -> KTable<Windowed<K>, VA, SessionWindowedSerde<KS>, VS>
235    where
236        VA: Any + Send + Sync + Clone,
237        KS: Serde<K> + Clone + 'static,
238        VS: Serde<VA> + Clone + 'static,
239        I: Fn() -> VA + Clone + Send + Sync + 'static,
240        A: Fn(&K, &V, VA) -> VA + Clone + Send + Sync + 'static,
241        M: Fn(&K, VA, VA) -> VA + Clone + Send + Sync + 'static,
242    {
243        let Materialized {
244            key_serde,
245            value_serde,
246            caching,
247            ..
248        } = materialized;
249        let suppress_factory =
250            session_suppress_factory::<K, VA, KS, VS>(key_serde.clone(), value_serde.clone());
251        let parent = self.parent;
252        let key_changing = self.key_changing_upstream;
253        let rp_lower = self.repartition_lower.take();
254        let windows = self.windows;
255        let emit = self.emit;
256        let mut g = self.builder.borrow_mut();
257        let agg_parent = KGroupedStream::<K, V>::record_repartition(
258            &mut g,
259            &store_name,
260            parent,
261            key_changing,
262            rp_lower,
263        );
264
265        let agg_name = g.new_processor_name(names::AGGREGATE);
266        let agg_id = g.graph.add(
267            agg_name.clone(),
268            GraphNodeKind::Aggregate {
269                store_name: store_name.clone(),
270                changelog: true,
271            },
272            vec![agg_parent],
273        );
274        let store_for_thunk = store_name.clone();
275        let key_serde_for_lower = key_serde.clone();
276        let value_serde_for_lower = value_serde.clone();
277        g.graph.nodes[agg_id].lower = Some(Box::new(move |state: &mut LowerState| {
278            let parent = NodeHandle::<K, V>::from_name(state.handle_name[&agg_parent].clone());
279            let store_for_proc = store_for_thunk.clone();
280            let init = init.clone();
281            let agg = agg.clone();
282            let merger = merger.clone();
283            let h = state
284                .topology
285                .add_processor::<K, V, Windowed<K>, crate::dsl::processors::change::Change<VA>, _, _, _>(
286                    agg_name.clone(),
287                    move || KStreamSessionAggregateProcessor {
288                        store_name: store_for_proc.clone(),
289                        gap: windows.gap,
290                        init: init.clone(),
291                        agg: agg.clone(),
292                        merger: merger.clone(),
293                        emit,
294                        grace: windows.grace,
295                        stream_time: i64::MIN,
296                        last_emitted_close: i64::MIN,
297                        forwarder: crate::dsl::processors::tuple_forwarder::TupleForwarder::default(),
298                        _pd: PhantomData,
299                    },
300                    [parent],
301                );
302            state.topology.add_session_store::<K, VA, KS, VS>(
303                store_for_thunk.clone(),
304                key_serde_for_lower.clone(),
305                value_serde_for_lower.clone(),
306                windows.gap,
307                windows.grace,
308                [h.name().to_string()],
309            );
310            // Cache only emit-on-update session aggregates: emit-final stays
311            // uncached or the flush would emit the per-update tombstones+updates
312            // emit-final suppresses.
313            state
314                .topology
315                .mark_store_caching(&store_for_thunk, caching && emit.is_on_update());
316            state.handle_name.insert(agg_id, h.name().to_string());
317        }));
318
319        drop(g);
320        KTable::new(
321            Rc::clone(&self.builder),
322            agg_id,
323            Some(store_name),
324            None,
325            SessionWindowedSerde::new(key_serde.clone()),
326            value_serde.clone(),
327        )
328        .with_window_grace(Some(windows.grace))
329        .with_suppress_factory(Some(suppress_factory))
330    }
331
332    fn lower_reduce<KS, VS, R>(
333        mut self,
334        materialized: Materialized<KS, VS>,
335        store_name: String,
336        reducer: R,
337    ) -> KTable<Windowed<K>, V, SessionWindowedSerde<KS>, VS>
338    where
339        KS: Serde<K> + Clone + 'static,
340        VS: Serde<V> + Clone + 'static,
341        R: Fn(&V, &V) -> V + Clone + Send + Sync + 'static,
342    {
343        let Materialized {
344            key_serde,
345            value_serde,
346            caching,
347            ..
348        } = materialized;
349        let suppress_factory =
350            session_suppress_factory::<K, V, KS, VS>(key_serde.clone(), value_serde.clone());
351        let parent = self.parent;
352        let key_changing = self.key_changing_upstream;
353        let rp_lower = self.repartition_lower.take();
354        let windows = self.windows;
355        let emit = self.emit;
356        let mut g = self.builder.borrow_mut();
357        let agg_parent = KGroupedStream::<K, V>::record_repartition(
358            &mut g,
359            &store_name,
360            parent,
361            key_changing,
362            rp_lower,
363        );
364
365        let red_name = g.new_processor_name(names::REDUCE);
366        let red_id = g.graph.add(
367            red_name.clone(),
368            GraphNodeKind::Aggregate {
369                store_name: store_name.clone(),
370                changelog: true,
371            },
372            vec![agg_parent],
373        );
374        let store_for_thunk = store_name.clone();
375        let key_serde_for_lower = key_serde.clone();
376        let value_serde_for_lower = value_serde.clone();
377        g.graph.nodes[red_id].lower = Some(Box::new(move |state: &mut LowerState| {
378            let parent = NodeHandle::<K, V>::from_name(state.handle_name[&agg_parent].clone());
379            let store_for_proc = store_for_thunk.clone();
380            let reducer = reducer.clone();
381            let h = state
382                .topology
383                .add_processor::<K, V, Windowed<K>, crate::dsl::processors::change::Change<V>, _, _, _>(
384                    red_name.clone(),
385                    move || KStreamSessionReduceProcessor {
386                        store_name: store_for_proc.clone(),
387                        gap: windows.gap,
388                        reducer: reducer.clone(),
389                        emit,
390                        grace: windows.grace,
391                        stream_time: i64::MIN,
392                        last_emitted_close: i64::MIN,
393                        forwarder: crate::dsl::processors::tuple_forwarder::TupleForwarder::default(),
394                        _pd: PhantomData,
395                    },
396                    [parent],
397                );
398            state.topology.add_session_store::<K, V, KS, VS>(
399                store_for_thunk.clone(),
400                key_serde_for_lower.clone(),
401                value_serde_for_lower.clone(),
402                windows.gap,
403                windows.grace,
404                [h.name().to_string()],
405            );
406            // Cache only emit-on-update session reduces (see aggregate lower).
407            state
408                .topology
409                .mark_store_caching(&store_for_thunk, caching && emit.is_on_update());
410            state.handle_name.insert(red_id, h.name().to_string());
411        }));
412
413        drop(g);
414        KTable::new(
415            Rc::clone(&self.builder),
416            red_id,
417            Some(store_name),
418            None,
419            SessionWindowedSerde::new(key_serde.clone()),
420            value_serde.clone(),
421        )
422        .with_window_grace(Some(windows.grace))
423        .with_suppress_factory(Some(suppress_factory))
424    }
425}
426
427/// Build the suppress-store factory for a session-aggregation result table.
428/// Captures the session key serde ([`SessionWindowedSerde`]) + the aggregate value
429/// serde so a downstream `suppress` registers a `SuppressBytesStore<Windowed<K>, VA>`
430/// with the session-windowed key serde + the matching changelog config.
431pub(crate) fn session_suppress_factory<K, VA, KS, VS>(
432    key_serde: KS,
433    value_serde: VS,
434) -> SuppressStoreFactory
435where
436    K: Any + Send + Sync + Clone,
437    VA: Any + Send + Clone,
438    KS: Serde<K> + Clone + 'static,
439    VS: Serde<VA> + Clone + 'static,
440{
441    std::sync::Arc::new(
442        move |state: &mut LowerState, store_name: &str, proc_name: &str, logging: bool| {
443            // The suppress buffer's changelog is a plain compacted KV changelog
444            // (the JVM suppress buffer is a compacted KV store) — no retention arg.
445            state
446                .topology
447                .add_suppress_store::<Windowed<K>, VA, SessionWindowedSerde<KS>, VS>(
448                    store_name.to_string(),
449                    SessionWindowedSerde::new(key_serde.clone()),
450                    value_serde.clone(),
451                    logging,
452                    [proc_name.to_string()],
453                );
454        },
455    )
456}
457
458#[cfg(test)]
459mod tests {
460    use assert2::check;
461    use krabka_units::prelude::*;
462
463    use crate::{
464        dsl::{
465            StreamsBuilder,
466            emit::EmitStrategy,
467            windows::{SessionWindowedSerde, SessionWindows, Window, Windowed},
468        },
469        processor::serde::{Consumed, I64Serde, Produced, StringSerde},
470        test_driver::TopologyTestDriver,
471    };
472
473    /// Sub-task 3d-ii: an emit-on-update session aggregate store, the default, is
474    /// marked caching. With a positive cache budget it lands in `cache_owner`.
475    #[test]
476    fn emit_on_update_session_store_is_cached() {
477        let b = StreamsBuilder::new();
478        b.stream::<String, String>(["in"])
479            .group_by_key()
480            .windowed_by_session(SessionWindows::of_inactivity_gap(millis(60)))
481            .count("s");
482        let built = b.build("app").unwrap();
483        let g = pollster::block_on(built.instantiate(
484            &crate::store::backend::StoreBackend::InMemory,
485            "app",
486            mebibytes(10),
487        ))
488        .unwrap();
489        check!(
490            g.cache_owner.contains_key("s"),
491            "emit-on-update session store must be cached, cache_owner = {:?}",
492            g.cache_owner
493        );
494    }
495
496    /// Sub-task 3d-ii: an emit-FINAL session aggregate store must stay UNCACHED.
497    /// A cache flush would emit the per-update tombstones and updates that
498    /// emit-final suppresses. The `emit.is_on_update()` mark guard enforces this.
499    #[test]
500    fn emit_final_session_store_is_not_cached() {
501        let b = StreamsBuilder::new();
502        b.stream::<String, String>(["in"])
503            .group_by_key()
504            .windowed_by_session(SessionWindows::of_inactivity_gap(millis(60)).grace(millis(10)))
505            .emit_strategy(EmitStrategy::on_window_close())
506            .count("s");
507        let built = b.build("app").unwrap();
508        let g = pollster::block_on(built.instantiate(
509            &crate::store::backend::StoreBackend::InMemory,
510            "app",
511            mebibytes(10),
512        ))
513        .unwrap();
514        check!(
515            !g.cache_owner.contains_key("s"),
516            "emit-final session store must NOT be cached, cache_owner = {:?}",
517            g.cache_owner
518        );
519    }
520
521    /// `windowedBy(SessionWindows).emit_strategy(on_window_close).count`:
522    /// emit-final (KIP-825) suppresses the per-update and merge-tombstone emits.
523    /// It forwards a single final count for each session only after stream-time
524    /// advances past that session's close. Records a@1 and a@4 fall within gap 60
525    /// and form session `[1,4]` with count 2, silently. The a@1000 record opens
526    /// session `[1000,1000]` and closes `[1,4]`, because end 4 <=
527    /// `window_close_time` = 1000 - grace 10 = 990. It emits exactly one record:
528    /// session `[1,4]` with value 2. A grace of 10 keeps the data-defined session
529    /// open at its own inclusive end until stream-time jumps ahead.
530    #[test]
531    fn dsl_session_count_emit_final_emits_once_on_close() {
532        let b = StreamsBuilder::new();
533        b.stream::<String, String>(["in"])
534            .group_by_key()
535            .windowed_by_session(SessionWindows::of_inactivity_gap(millis(60)).grace(millis(10)))
536            .emit_strategy(EmitStrategy::on_window_close())
537            .count("s")
538            .to_stream()
539            .to_explicit(
540                "out",
541                Produced::with(SessionWindowedSerde::new(StringSerde), I64Serde),
542            );
543        let built = b.build("app").unwrap();
544        let mut d = TopologyTestDriver::new(&built).unwrap();
545        for ts in [1, 4, 1000] {
546            d.pipe_input(
547                "in",
548                Consumed::with(StringSerde, StringSerde),
549                Some("a".to_string()),
550                "x".to_string(),
551                ts,
552            );
553        }
554        let p = || Produced::with(SessionWindowedSerde::new(StringSerde), I64Serde);
555        assert_eq!(
556            d.read_output("out", p()),
557            Some((
558                Some(Windowed {
559                    key: "a".into(),
560                    window: Window { start: 1, end: 4 }
561                }),
562                2
563            )),
564            "emit-final forwards session [1,4] with final count 2 on close"
565        );
566        // No per-update / merge-tombstone emits and the still-open session
567        // [1000,1000] is not emitted.
568        assert_eq!(
569            d.read_output("out", p()),
570            None,
571            "exactly one emit-final record"
572        );
573    }
574}
krabka_client_streams/dsl/session_windowed_kgrouped.rs

krabka_client_streams/dsl/
session_windowed_kgrouped.rs
