murr 0.2.1

Columnar in-memory cache for AI/ML inference workloads
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

use arrow::array::{Array, StringArray};

use crate::core::DTypeName;
use crate::io::row::read::ReadBatchBuilder;
use crate::io::row::write::WriteRow;
use crate::io::schema::{SegmentColumnSchema, SegmentSchema};
use crate::io::store::{KeyValue, Store};

pub fn payload_segment() -> SegmentSchema {
    SegmentSchema::new(&[SegmentColumnSchema {
        index: 0,
        dtype: DTypeName::Utf8,
        name: "payload".into(),
        offset: 0,
    }])
}

pub fn put<S: Store>(store: &mut S, table: &str, rows: &[(&str, &[u8])]) {
    let segment = payload_segment();
    let col = &segment.columns[0];
    let kvs: Vec<KeyValue> = rows
        .iter()
        .map(|(k, v)| {
            let mut row = WriteRow::new(&segment, k);
            row.write_dynamic(col, v);
            row.into()
        })
        .collect();
    store.write(table, kvs).unwrap();
}

pub fn fetch<S: Store>(store: &S, table: &str, keys: &[&[u8]]) -> Vec<Option<Vec<u8>>> {
    let segment = payload_segment();
    let cols: Vec<&SegmentColumnSchema> = segment.columns.iter().collect();
    let builder = ReadBatchBuilder::new(&segment, cols, keys.len());
    let batch = store.read(table, keys, builder).unwrap();
    let arr = batch
        .column(0)
        .as_any()
        .downcast_ref::<StringArray>()
        .expect("payload column is Utf8");
    (0..arr.len())
        .map(|i| {
            if arr.is_null(i) {
                None
            } else {
                Some(arr.value(i).as_bytes().to_vec())
            }
        })
        .collect()
}