arcon 0.2.1

A runtime for writing streaming applications
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

use arrow::{datatypes::Schema, record_batch::RecordBatch};

use crate::{
    index::{IndexOps, WindowIndex},
    prelude::*,
    stream::operator::window::WindowContext,
    table::{to_record_batches, ImmutableTable, RawRecordBatch},
    util::ArconFnBounds,
};
use arcon_state::{backend::handles::ActiveHandle, Backend, VecState};
use std::marker::PhantomData;

/// A window index for Arrow Data
///
/// Elements are appended into RecordBatches and once a window is triggered,
/// the underlying Arrow Schema and Vec<RecordBatch> is exposed.
pub struct ArrowWindow<IN, OUT, F, B>
where
    IN: ArconType + ToArrow,
    OUT: ArconType,
    F: Fn(Arc<Schema>, Vec<RecordBatch>) -> ArconResult<OUT> + ArconFnBounds,
    B: Backend,
{
    handle: ActiveHandle<B, VecState<RawRecordBatch>, u64, u64>,
    map: std::collections::HashMap<WindowContext, MutableTable>,
    udf: F,
    _marker: std::marker::PhantomData<IN>,
}

impl<IN, OUT, F, B> ArrowWindow<IN, OUT, F, B>
where
    IN: ArconType + ToArrow,
    OUT: ArconType,
    F: Fn(Arc<Schema>, Vec<RecordBatch>) -> ArconResult<OUT> + ArconFnBounds,
    B: Backend,
{
    pub fn new(backend: Arc<B>, udf: F) -> Self {
        let mut handle = Handle::vec("window_handle")
            .with_item_key(0)
            .with_namespace(0);

        backend.register_vec_handle(&mut handle);

        let handle = handle.activate(backend);

        Self {
            handle,
            map: std::collections::HashMap::new(),
            udf,
            _marker: PhantomData,
        }
    }
}

impl<IN, OUT, F, B> WindowIndex for ArrowWindow<IN, OUT, F, B>
where
    IN: ArconType + ToArrow,
    OUT: ArconType,
    F: Fn(Arc<Schema>, Vec<RecordBatch>) -> ArconResult<OUT> + ArconFnBounds,
    B: Backend,
{
    type IN = IN;
    type OUT = OUT;

    fn on_element(&mut self, element: Self::IN, ctx: WindowContext) -> ArconResult<()> {
        let table = self.map.entry(ctx).or_insert_with(IN::table);
        table.append(element, None)?;

        Ok(())
    }

    fn result(&mut self, ctx: WindowContext) -> ArconResult<Self::OUT> {
        let table = self.map.entry(ctx).or_insert_with(IN::table);
        self.handle.set_item_key(ctx.key);
        self.handle.set_namespace(ctx.index);

        // fetch in-memory batches
        let mut batches = table.batches()?;
        // fetch if any raw batches and append to the vector...
        let raw_batches = self.handle.get()?;
        batches.append(&mut to_record_batches(Arc::new(IN::schema()), raw_batches)?);

        (self.udf)(Arc::new(IN::schema()), batches)
    }

    fn clear(&mut self, ctx: WindowContext) -> ArconResult<()> {
        // clear from memory layer
        let _ = self.map.remove(&ctx);

        // clear everything in the backend
        self.handle.set_item_key(ctx.key);
        self.handle.set_namespace(ctx.index);

        self.handle.clear()?;

        Ok(())
    }
}
impl<IN, OUT, F, B> IndexOps for ArrowWindow<IN, OUT, F, B>
where
    IN: ArconType + ToArrow,
    OUT: ArconType,
    F: Fn(Arc<Schema>, Vec<RecordBatch>) -> ArconResult<OUT> + ArconFnBounds,
    B: Backend,
{
    fn persist(&mut self) -> ArconResult<()> {
        for (ctx, table) in self.map.iter_mut() {
            self.handle.set_item_key(ctx.key);
            self.handle.set_namespace(ctx.index);

            let batches = table.raw_batches()?;
            for batch in batches {
                self.handle.append(batch)?;
            }
        }
        Ok(())
    }
    fn set_key(&mut self, _: u64) {}
    fn table(&mut self) -> ArconResult<Option<ImmutableTable>> {
        Ok(None)
    }
}