zrx_stream/stream/operator/
group.rs

1// Copyright (c) Zensical LLC <https://zensical.org>
2
3// SPDX-License-Identifier: MIT
4// Third-party contributions licensed under CLA
5
6// Permission is hereby granted, free of charge, to any person obtaining a copy
7// of this software and associated documentation files (the "Software"), to
8// deal in the Software without restriction, including without limitation the
9// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
10// sell copies of the Software, and to permit persons to whom the Software is
11// furnished to do so, subject to the following conditions:
12
13// The above copyright notice and this permission notice shall be included in
14// all copies or substantial portions of the Software.
15
16// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18// FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT. IN NO EVENT SHALL THE
19// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
21// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
22// IN THE SOFTWARE.
23
24// ----------------------------------------------------------------------------
25
26//! Group operator.
27
28use ahash::HashMap;
29use zrx_scheduler::action::output::IntoOutputs;
30use zrx_scheduler::action::{Descriptor, Report};
31use zrx_scheduler::effect::Item;
32use zrx_scheduler::{Id, Value};
33
34use crate::stream::function::SelectFn;
35use crate::stream::value::Delta;
36use crate::stream::Stream;
37
38use super::Operator;
39
40// ----------------------------------------------------------------------------
41// Structs
42// ----------------------------------------------------------------------------
43
44/// Group operator.
45struct Group<F, I> {
46    /// Operator function.
47    function: F,
48    /// Store of group identifiers.
49    store: HashMap<I, I>,
50}
51
52// ----------------------------------------------------------------------------
53// Implementations
54// ----------------------------------------------------------------------------
55
56impl<I, T> Stream<I, T>
57where
58    I: Id,
59    T: Value + Clone,
60{
61    pub fn group<F, U>(&self, f: F) -> Stream<I, Delta<I, U>>
62    where
63        F: SelectFn<I, T, I>,
64        U: Value,
65    {
66        self.workflow.add_operator(
67            [self.id],
68            Group {
69                function: f,
70                store: HashMap::default(),
71            },
72        )
73    }
74}
75
76// ----------------------------------------------------------------------------
77// Trait implementations
78// ----------------------------------------------------------------------------
79
80impl<I, T, F> Operator<I, T> for Group<F, I>
81where
82    I: Id,
83    T: Value + Clone,
84    F: SelectFn<I, T, I>,
85{
86    type Item<'a> = Item<&'a I, Option<&'a T>>;
87
88    /// Handles the given item.
89    ///
90    /// Grouping is achieved by applying the operator function to each incoming
91    /// item to determine its group identifier. An internal mapping of items to
92    /// the computed group identifiers is maintained, as items might need to be
93    /// migrated between groups when their associated data changes.
94    #[cfg_attr(
95        feature = "tracing",
96        tracing::instrument(level = "debug", skip_all, fields(id = %item.id))
97    )]
98    fn handle(&mut self, item: Self::Item<'_>) -> impl IntoOutputs<I> {
99        if let Some(data) = item.data {
100            // When new data arrives, we apply the operator function to compute
101            // the group identifier. If the item was previously associated with
102            // another group, we emit a deletion for the prior and an insertion
103            // for the new group. Otherwise, we only emit an insertion.
104            self.function.execute(item.id, data).map(|report| {
105                report.map(|to| {
106                    let group = self.store.insert(item.id.clone(), to.clone());
107                    let prior = group.filter(|id| *id != to).map(|id| {
108                        let inner = Item::new(item.id.clone(), None);
109                        Item::new(id, Some(Delta::from([inner])))
110                    });
111
112                    // The item returned receives the group identifier, wrapping
113                    // the inner item, which is the item passed to the operator
114                    let inner = item.into_owned().map(Some);
115                    let delta = Delta::from([inner]);
116                    prior
117                        .into_iter()
118                        .chain(Some(Item::new(to, Some(delta))))
119                        .collect()
120                })
121            })
122        } else {
123            // If the incoming item has no data, interpret this as a deletion,
124            // removing the item from its previously allocated group, if any
125            let prior = self.store.remove(item.id).map(|id| {
126                let inner = Item::new(item.id.clone(), None);
127                Item::new(id, Some(Delta::from([inner])))
128            });
129
130            // Return delta of items
131            Ok(Report::new(prior.into_iter().collect::<Vec<_>>()))
132        }
133    }
134
135    /// Returns the descriptor.
136    #[inline]
137    fn descriptor(&self) -> Descriptor {
138        Descriptor::default()
139    }
140}
zrx_stream/stream/operator/group.rs

zrx_stream/stream/operator/
group.rs
