dbsp/operator/dynamic/time_series/radix_tree/
treenode.rs

1use std::{
2    fmt::{self, Debug, Display, Formatter, Write},
3    mem::take,
4};
5
6use feldera_macros::IsNone;
7use num::PrimInt;
8use rkyv::{Archive, Deserialize, Serialize};
9use size_of::SizeOf;
10
11use super::{Prefix, RADIX};
12use crate::{
13    DBData, declare_trait_object,
14    dynamic::{Data, DataTrait, DowncastTrait, DynOpt, Erase},
15    operator::dynamic::aggregate::AggCombineFunc,
16};
17
18/// Pointer to a child node.
19#[derive(
20    Clone,
21    Default,
22    Debug,
23    SizeOf,
24    PartialEq,
25    Eq,
26    Hash,
27    PartialOrd,
28    Ord,
29    Archive,
30    Serialize,
31    Deserialize,
32    IsNone,
33)]
34#[archive_attr(derive(Ord, Eq, PartialEq, PartialOrd))]
35#[archive(compare(PartialEq, PartialOrd))]
36#[archive(bound(archive = "Prefix<TS>: DBData"))]
37pub struct ChildPtr<TS: DBData, A: DBData> {
38    /// Unique prefix of a child subtree, which serves as a pointer
39    /// to the child node.  Given this prefix the child node can
40    /// be located using `Cursor::seek_key`, unless
41    /// `child_prefix.is_leaf()`, in which case we're at the bottom
42    /// of the tree.
43    child_prefix: Prefix<TS>,
44    /// Aggregate over all timestamps covered by the child subtree.
45    child_agg: A,
46}
47
48impl<TS: DBData, A: DBData> ChildPtr<TS, A> {
49    #[cfg(test)]
50    fn new(child_prefix: Prefix<TS>, child_agg: A) -> Self {
51        Self {
52            child_prefix,
53            child_agg,
54        }
55    }
56}
57
58impl<TS, A> Display for ChildPtr<TS, A>
59where
60    TS: DBData + PrimInt,
61    A: DBData,
62{
63    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), fmt::Error> {
64        write!(f, "[{}->{:?}]", self.child_prefix, self.child_agg)
65    }
66}
67
68pub trait ChildPtrTrait<TS, A>: Data
69where
70    TS: DBData + PrimInt,
71    A: DataTrait + ?Sized,
72{
73    #[allow(clippy::wrong_self_convention)]
74    fn from_timestamp(&mut self, key: TS, child_agg: &mut A);
75
76    #[allow(clippy::wrong_self_convention)]
77    fn from_prefix(&mut self, child_prefix: Prefix<TS>);
78
79    fn child_prefix(&self) -> Prefix<TS>;
80
81    fn child_agg(&self) -> &A;
82
83    fn child_agg_mut(&mut self) -> &mut A;
84}
85
86impl<TS, A, AType> ChildPtrTrait<TS, A> for ChildPtr<TS, AType>
87where
88    TS: DBData + PrimInt,
89    A: DataTrait + ?Sized,
90    AType: DBData + Erase<A>,
91{
92    fn from_timestamp(&mut self, key: TS, child_agg: &mut A) {
93        let child_agg = unsafe { child_agg.downcast_mut::<AType>() };
94
95        self.child_prefix = Prefix::from_timestamp(key);
96        self.child_agg = take(child_agg);
97    }
98
99    fn from_prefix(&mut self, child_prefix: Prefix<TS>) {
100        self.child_prefix = child_prefix;
101        self.child_agg = Default::default();
102    }
103
104    fn child_prefix(&self) -> Prefix<TS> {
105        self.child_prefix.clone()
106    }
107
108    fn child_agg(&self) -> &A {
109        self.child_agg.erase()
110    }
111
112    fn child_agg_mut(&mut self) -> &mut A {
113        self.child_agg.erase_mut()
114    }
115}
116
117declare_trait_object!(DynChildPtr<TS, A> = dyn ChildPtrTrait<TS, A>
118where
119    A: DataTrait + ?Sized,
120    TS: PrimInt + DBData);
121
122/// Radix tree node.
123#[derive(
124    Clone,
125    Debug,
126    Default,
127    SizeOf,
128    PartialEq,
129    Eq,
130    Hash,
131    PartialOrd,
132    Ord,
133    Archive,
134    Serialize,
135    Deserialize,
136    IsNone,
137)]
138//#[archive_attr(derive(Ord, Eq, PartialEq, PartialOrd))]
139//#[archive(compare(PartialEq, PartialOrd))]
140//#[archive(bound(archive = "[Option<ChildPtr<TS, A>>; RADIX]: DBData"))]
141pub struct TreeNode<TS: DBData, A: DBData> {
142    /// Array of children.
143    // `Option` doesn't introduce space overhead.
144    children: [Option<ChildPtr<TS, A>>; RADIX],
145}
146
147// For some unknown reason we cant jsut use `#[archive(compare(PartialEq,
148// PartialOrd))]` and `#[archive_attr(derive(Clone, Ord, PartialOrd, Eq,
149// PartialEq))]` so we define these traits for now. They don't get called so
150// leaving it unimplemented!().
151
152impl<TS, A> PartialEq<TreeNode<TS, A>> for ArchivedTreeNode<TS, A>
153where
154    TS: DBData,
155    A: DBData,
156    //[Option<ChildPtr<TS, A>>; RADIX]: DBData,
157{
158    fn eq(&self, _other: &TreeNode<TS, A>) -> bool {
159        unimplemented!()
160    }
161}
162impl<TS, A> PartialOrd<TreeNode<TS, A>> for ArchivedTreeNode<TS, A>
163where
164    TS: DBData,
165    A: DBData,
166    //[Option<ChildPtr<TS, A>>; RADIX]: DBData,
167{
168    fn partial_cmp(&self, _other: &TreeNode<TS, A>) -> Option<std::cmp::Ordering> {
169        unimplemented!()
170    }
171}
172
173impl<TS, A> PartialEq for ArchivedTreeNode<TS, A>
174where
175    TS: DBData,
176    A: DBData,
177{
178    fn eq(&self, _other: &Self) -> bool {
179        unimplemented!()
180    }
181}
182
183impl<TS, A> Eq for ArchivedTreeNode<TS, A>
184where
185    TS: DBData,
186    A: DBData,
187{
188}
189
190impl<TS, A> PartialOrd for ArchivedTreeNode<TS, A>
191where
192    TS: DBData,
193    A: DBData,
194{
195    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
196        Some(self.cmp(other))
197    }
198}
199
200impl<TS, A> Ord for ArchivedTreeNode<TS, A>
201where
202    TS: DBData,
203    A: DBData,
204{
205    fn cmp(&self, _other: &Self) -> std::cmp::Ordering {
206        unimplemented!()
207    }
208}
209
210// end of not implemented, hopefully derivable, traits
211
212impl<TS, A> Display for TreeNode<TS, A>
213where
214    TS: DBData + PrimInt,
215    A: DBData,
216{
217    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), fmt::Error> {
218        for child in self.children.iter() {
219            match child {
220                None => f.write_char('.')?,
221                Some(child) => write!(f, "{child}")?,
222            }
223        }
224
225        Ok(())
226    }
227}
228
229pub trait TreeNodeTrait<TS, A>: Data
230where
231    TS: PrimInt + DBData,
232    A: DataTrait + ?Sized,
233{
234    fn clear(&mut self);
235
236    /// Returns a reference to a child pointer number `slot`.
237    fn slot(&self, slot: usize) -> &DynOpt<DynChildPtr<TS, A>>;
238
239    /// Returns a mutable reference to a child pointer number `slot`.
240    fn slot_mut(&mut self, slot: usize) -> &mut DynOpt<DynChildPtr<TS, A>>;
241
242    /// Counts the number of non-empty slots.
243    fn occupied_slots(&self) -> usize;
244
245    /// The first non-empty slot or `None` if all slots are empty.
246    fn first_occupied_slot(&self) -> Option<&DynChildPtr<TS, A>>;
247
248    /// Computes aggregate of the entire subtree under `self` as a
249    /// sum of aggregates of its children.
250    fn aggregate(&self, combine: &dyn AggCombineFunc<A>, result: &mut DynOpt<A>);
251
252    fn display(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result;
253}
254
255impl<TS, A, AType> TreeNodeTrait<TS, A> for TreeNode<TS, AType>
256where
257    TS: PrimInt + DBData,
258    A: DataTrait + ?Sized,
259    AType: DBData + Erase<A>,
260    //[Option<ChildPtr<TS, AType>>; RADIX]: DBData,
261{
262    fn clear(&mut self) {
263        *self = Default::default();
264    }
265
266    fn slot(&self, slot: usize) -> &DynOpt<DynChildPtr<TS, A>> {
267        self.children[slot].erase()
268    }
269
270    fn slot_mut(&mut self, slot: usize) -> &mut DynOpt<DynChildPtr<TS, A>> {
271        self.children[slot].erase_mut()
272    }
273
274    fn occupied_slots(&self) -> usize {
275        let mut res = 0;
276        for child in self.children.iter() {
277            if child.is_some() {
278                res += 1;
279            }
280        }
281        res
282    }
283
284    fn first_occupied_slot(&self) -> Option<&DynChildPtr<TS, A>> {
285        if let Some(child) = self.children.iter().flatten().next() {
286            return Some(child.erase());
287        }
288        None
289    }
290
291    fn aggregate(&self, combine: &dyn AggCombineFunc<A>, acc: &mut DynOpt<A>) {
292        for child in self.children.iter().flatten() {
293            if let Some(acc) = acc.get_mut() {
294                combine(acc, child.child_agg.erase());
295            } else {
296                acc.from_ref(child.child_agg.erase());
297            }
298        }
299    }
300
301    fn display(&self, f: &mut Formatter<'_>) -> fmt::Result {
302        Display::fmt(self, f)
303    }
304}
305
306declare_trait_object!(DynTreeNode<TS, A> = dyn TreeNodeTrait<TS, A>
307where
308    TS: PrimInt + DBData,
309    A: DataTrait + ?Sized);
310
311impl<TS, A> Display for DynTreeNode<TS, A>
312where
313    TS: PrimInt + DBData,
314    A: DataTrait + ?Sized,
315{
316    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), fmt::Error> {
317        self.display(f)
318    }
319}
320
321/// Describes incremental update to a radix tree node.
322#[derive(
323    Clone,
324    Default,
325    Debug,
326    SizeOf,
327    PartialEq,
328    Eq,
329    Hash,
330    PartialOrd,
331    Ord,
332    Archive,
333    Serialize,
334    Deserialize,
335    IsNone,
336)]
337#[archive_attr(derive(Ord, Eq, PartialEq, PartialOrd))]
338#[archive(bound(archive = "Option<TreeNode<TS,A>>: DBData, Prefix<TS>: DBData"))]
339#[archive(compare(PartialEq, PartialOrd))]
340pub struct TreeNodeUpdate<TS: DBData, A: DBData> {
341    /// Prefix that uniquely identifies the node.
342    pub prefix: Prefix<TS>,
343    /// Old value of the node or `None` if we are creating a new node.
344    pub old: Option<TreeNode<TS, A>>,
345    /// New value of the node or `None` if we are deleting an existing node.
346    pub new: Option<TreeNode<TS, A>>,
347}
348
349pub trait TreeNodeUpdateTrait<TS, A>: Data
350where
351    TS: PrimInt + DBData,
352    A: DataTrait + ?Sized,
353{
354    fn prefix(&self) -> Prefix<TS>;
355
356    fn old(&self) -> &DynOpt<DynTreeNode<TS, A>>;
357
358    fn old_mut(&mut self) -> &mut DynOpt<DynTreeNode<TS, A>>;
359
360    #[allow(clippy::wrong_self_convention, clippy::new_ret_no_self)]
361    fn new(&self) -> &DynOpt<DynTreeNode<TS, A>>;
362
363    fn new_mut(&mut self) -> &mut DynOpt<DynTreeNode<TS, A>>;
364
365    #[allow(clippy::wrong_self_convention)]
366    fn from_existing_node(&mut self, prefix: Prefix<TS>, node: &mut DynTreeNode<TS, A>);
367
368    #[allow(clippy::wrong_self_convention)]
369    fn from_new_node(&mut self, prefix: Prefix<TS>, node: &mut DynTreeNode<TS, A>);
370}
371
372impl<TS, A, AType> TreeNodeUpdateTrait<TS, A> for TreeNodeUpdate<TS, AType>
373where
374    TS: PrimInt + DBData,
375    A: DataTrait + ?Sized,
376    AType: DBData + Erase<A>,
377    //[Option<ChildPtr<TS, AType>>; RADIX]: DBData,
378{
379    fn prefix(&self) -> Prefix<TS> {
380        self.prefix.clone()
381    }
382
383    fn old(&self) -> &DynOpt<DynTreeNode<TS, A>> {
384        self.old.erase()
385    }
386
387    fn old_mut(&mut self) -> &mut DynOpt<DynTreeNode<TS, A>> {
388        self.old.erase_mut()
389    }
390
391    fn new(&self) -> &DynOpt<DynTreeNode<TS, A>> {
392        self.new.erase()
393    }
394
395    fn new_mut(&mut self) -> &mut DynOpt<DynTreeNode<TS, A>> {
396        self.new.erase_mut()
397    }
398
399    fn from_existing_node(&mut self, prefix: Prefix<TS>, node: &mut DynTreeNode<TS, A>) {
400        let node = unsafe { node.downcast_mut::<TreeNode<TS, AType>>() };
401        self.prefix = prefix;
402        self.old = Some(node.clone());
403        self.new = Some(take(node));
404    }
405
406    fn from_new_node(&mut self, prefix: Prefix<TS>, node: &mut DynTreeNode<TS, A>) {
407        let node = unsafe { node.downcast_mut::<TreeNode<TS, AType>>() };
408
409        self.prefix = prefix;
410        self.old = None;
411        self.new = Some(take(node));
412    }
413}
414declare_trait_object!(DynTreeNodeUpdate<TS, A> = dyn TreeNodeUpdateTrait<TS, A>
415where
416    A: DataTrait + ?Sized,
417    TS: PrimInt + DBData);
418
419#[cfg(test)]
420mod test {
421    use crate::{
422        DBData,
423        dynamic::{DowncastTrait, DynData, Erase},
424        operator::dynamic::time_series::radix_tree::{
425            DynChildPtr, DynTreeNode, Prefix,
426            treenode::{ChildPtr, TreeNode, TreeNodeTrait},
427        },
428    };
429    use rkyv::{Deserialize, Infallible, archived_root, to_bytes};
430
431    fn aggregate_default(node: &mut DynTreeNode<u64, DynData /* <u64> */>) -> Option<u64> {
432        let mut result: Option<u64> = None;
433        node.aggregate(
434            &|acc, x| *acc.downcast_mut_checked::<u64>() += *x.downcast_checked::<u64>(),
435            result.erase_mut(),
436        );
437        result
438    }
439
440    fn child_from_timestamp<A: DBData>(key: u64, mut val: A) -> ChildPtr<u64, A> {
441        let mut child: ChildPtr<u64, A> = Default::default();
442        let dyn_child: &mut DynChildPtr<u64, DynData /* <A> */> = child.erase_mut();
443
444        dyn_child.from_timestamp(key, val.erase_mut());
445
446        child
447    }
448
449    #[test]
450    fn test_tree_node() {
451        let mut node = TreeNode::<u64, u64>::default();
452        let node: &mut DynTreeNode<u64, DynData /* <u64> */> = node.erase_mut();
453
454        assert_eq!(node.occupied_slots(), 0);
455        assert_eq!(node.first_occupied_slot(), None);
456        assert_eq!(aggregate_default(node), None);
457
458        node.slot_mut(1).set_some_with(&mut |ptr| {
459            ptr.from_timestamp(0x8000_0000_0000_0000u64, 10u64.erase_mut())
460        });
461        assert_eq!(node.occupied_slots(), 1);
462
463        assert_eq!(
464            node.first_occupied_slot()
465                .unwrap()
466                .downcast_checked::<ChildPtr<_, _>>(),
467            &child_from_timestamp(0x8000_0000_0000_0000u64, 10u64)
468        );
469
470        assert_eq!(aggregate_default(node), Some(10));
471
472        // Used for testing with RADIX=16.
473        // node.slot_mut(4).set_some_with(&mut |ptr| {
474        //     ptr.from_timestamp(0x4000_0000_0000_0000u64, 40u64.erase_mut())
475        // });
476        // assert_eq!(node.occupied_slots(), 2);
477        // assert_eq!(
478        //     node.first_occupied_slot()
479        //         .unwrap()
480        //         .downcast_checked::<ChildPtr<_, _>>(),
481        //     &child_from_timestamp(0x1000_0000_0000_0000u64, 10u64)
482        // );
483        // assert_eq!(aggregate_default(node), Some(50));
484
485        node.slot_mut(0).set_some_with(&mut |ptr| {
486            *ptr.downcast_mut_checked() = ChildPtr::new(
487                Prefix {
488                    key: 0x0fff_ffff_ffff_ffffu64,
489                    prefix_len: 4,
490                },
491                80u64,
492            )
493        });
494        assert_eq!(node.occupied_slots(), 2);
495        assert_eq!(
496            node.first_occupied_slot()
497                .unwrap()
498                .downcast_checked::<ChildPtr<_, _>>(),
499            &ChildPtr::new(
500                Prefix {
501                    key: 0x0fff_ffff_ffff_ffffu64,
502                    prefix_len: 4,
503                },
504                80u64,
505            )
506        );
507        assert_eq!(aggregate_default(node), Some(90));
508    }
509
510    #[test]
511    fn childptr_decode_encode() {
512        type Type = ChildPtr<u64, i32>;
513        for input in [
514            child_from_timestamp(u64::MIN, -1),
515            child_from_timestamp(0x1000_0000_0000_0000u64, 10),
516            child_from_timestamp(u64::MAX, 3),
517        ] {
518            let input: Type = input;
519            let encoded = to_bytes::<_, 4096>(&input).unwrap();
520            let archived = unsafe { archived_root::<Type>(&encoded[..]) };
521            let decoded: Type = archived.deserialize(&mut Infallible).unwrap();
522            assert_eq!(decoded, input);
523        }
524    }
525
526    #[test]
527    fn treenode_decode_encode() {
528        type Type = TreeNode<u64, i32>;
529
530        let mut input: Type = TreeNode::default();
531        input.slot_mut(1).set_some_with(&mut |ptr: &mut DynChildPtr<
532            u64,
533            DynData, /* <i32> */
534        >| {
535            *ptr.downcast_mut_checked() = child_from_timestamp(0x1000_0000_0000_0000u64, 10);
536        });
537
538        let encoded = to_bytes::<_, 4096>(&input).unwrap();
539        let archived = unsafe { archived_root::<Type>(&encoded[..]) };
540        let decoded: Type = archived.deserialize(&mut Infallible).unwrap();
541        assert_eq!(decoded, input);
542    }
543}
dbsp/operator/dynamic/time_series/radix_tree/treenode.rs

dbsp/operator/dynamic/time_series/radix_tree/
treenode.rs
