somatize_data/
frame.rs

1//! A batch of columns, which is what a source answers with.
2
3use arrow_array::RecordBatch;
4use arrow_ipc::reader::StreamReader;
5use arrow_ipc::writer::StreamWriter;
6use arrow_schema::SchemaRef;
7use arrow_select::concat::concat_batches;
8use somatize_core::Value;
9use std::fmt;
10
11/// Rows and columns, together with what each column is called and holds.
12///
13/// A type of ours rather than a `RecordBatch` in an `Opaque` for two reasons
14/// that are not style: a codec is an `impl` of somebody else's trait and the
15/// orphan rule wants one of the two types to be ours, and this is the word the
16/// design is written in — *the difference between training and deploying is how
17/// many rows the frame brings*.
18///
19/// It is not a tensor and does not want to be. Numbers and fixed-size lists of
20/// them are contiguous Arrow buffers, so converting is a reshape and not a copy;
21/// text is not a tensor until something tokenizes it. **The conversion is a
22/// node**, and the same node whether the frame came from a file or a topic.
23#[derive(Debug, Clone)]
24pub struct Frame(RecordBatch);
25
26impl Frame {
27    /// What gets written beside the bytes. Named after the **format** and not
28    /// the language: what is on disk is an Arrow IPC stream, and whoever reads
29    /// it back may be a `polars` on the other side of the wall.
30    pub const KIND: &'static str = "arrow.RecordBatch";
31
32    /// This batch, as a frame.
33    pub fn new(batch: RecordBatch) -> Self {
34        Self(batch)
35    }
36
37    /// How many rows it brings. The only number a caller usually wants: it is
38    /// what says whether a span was short because the dataset ended.
39    pub fn rows(&self) -> usize {
40        self.0.num_rows()
41    }
42
43    /// What the columns are called and what they hold — **without reading a
44    /// value**, which is the half of a virtual table worth having.
45    pub fn schema(&self) -> &SchemaRef {
46        self.0.schema_ref()
47    }
48
49    /// The batch itself, for whoever brought their own engine.
50    pub fn batch(&self) -> &RecordBatch {
51        &self.0
52    }
53
54    /// As a value, which is how it crosses an edge. `Opaque` and not a variant
55    /// of its own: the core has no dependencies and is not going to learn what a
56    /// column is.
57    pub fn value(self) -> Value {
58        Value::opaque(self)
59    }
60
61    /// The frame this value carries, if it carries one.
62    pub fn of(value: &Value) -> Option<&Self> {
63        value.downcast::<Self>()
64    }
65
66    /// What it weighs in bytes: Arrow IPC, buffers and all. No encoding pass and
67    /// no per-value work, which is the reason Arrow is the type that crosses an
68    /// edge rather than something converted to at the edges.
69    pub fn written(&self) -> Result<Vec<u8>, FrameError> {
70        let mut out = Vec::new();
71        let mut writer = StreamWriter::try_new(&mut out, self.schema())
72            .map_err(|e| FrameError(format!("a frame could not be written down: {e}")))?;
73        writer
74            .write(&self.0)
75            .map_err(|e| FrameError(format!("a frame could not be written down: {e}")))?;
76        writer
77            .finish()
78            .map_err(|e| FrameError(format!("a frame could not be written down: {e}")))?;
79        drop(writer);
80        Ok(out)
81    }
82
83    /// And back.
84    ///
85    /// One frame, however many batches the stream holds: what was written down
86    /// was a frame, and a frame is what comes back.
87    pub fn read(bytes: &[u8]) -> Result<Self, FrameError> {
88        let reader = StreamReader::try_new(bytes, None)
89            .map_err(|e| FrameError(format!("those bytes are not a frame: {e}")))?;
90        let schema = reader.schema();
91        let batches: Vec<RecordBatch> = reader
92            .collect::<Result<_, _>>()
93            .map_err(|e| FrameError(format!("those bytes are not a frame: {e}")))?;
94        let whole = concat_batches(&schema, batches.iter())
95            .map_err(|e| FrameError(format!("those rows would not join up: {e}")))?;
96        Ok(Self(whole))
97    }
98}
99
100/// Why that could not be written down, or read back.
101#[derive(Debug, Clone, PartialEq, Eq)]
102pub struct FrameError(String);
103
104impl FrameError {
105    /// The message.
106    pub fn message(&self) -> &str {
107        &self.0
108    }
109}
110
111impl fmt::Display for FrameError {
112    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
113        f.write_str(&self.0)
114    }
115}
116
117impl std::error::Error for FrameError {}
somatize_data/frame.rs

somatize_data/
frame.rs
