krabka_client_streams/dsl/
time_windowed_cogrouped.rs

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

krabka_client_streams/dsl/
time_windowed_cogrouped.rs
