krabka_client_streams/dsl/
windowed_kgrouped.rs

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

krabka_client_streams/dsl/
windowed_kgrouped.rs
