datui_lib/formats/
columns.rs

1//! Rows built a row at a time into typed columns, handed over as a `DataFrame`.
2//!
3//! A reader that knows each column's type up front (the GPS logs, a candump's frames,
4//! an ELF file's symbols) fills plain vectors here and never asks Polars to infer.
5
6use polars::prelude::*;
7
8/// The kinds of column, each with the value a cell holds, its type, and how a column of
9/// them becomes a series.
10macro_rules! kinds {
11    ($($(#[$doc:meta])* $kind:ident($value:ty) => $dtype:expr, $build:expr;)*) => {
12        /// The type of one column.
13        #[derive(Debug, Clone, Copy, PartialEq, Eq)]
14        pub enum Kind {
15            $($(#[$doc])* $kind,)*
16        }
17
18        impl Kind {
19            pub fn dtype(self) -> DataType {
20                match self {
21                    $(Kind::$kind => $dtype,)*
22                }
23            }
24        }
25
26        /// One value of a row, of the column's kind; `None` is null.
27        #[derive(Debug, Clone, PartialEq)]
28        pub enum Cell {
29            $($kind(Option<$value>),)*
30        }
31
32        #[derive(Debug)]
33        enum Values {
34            $($kind(Vec<Option<$value>>),)*
35        }
36
37        impl Values {
38            fn new(kind: Kind) -> Self {
39                match kind {
40                    $(Kind::$kind => Values::$kind(Vec::new()),)*
41                }
42            }
43
44            /// Append `cell`, or a null when it is of another kind: a reader's mistake,
45            /// which costs a value rather than misaligning the columns.
46            fn push(&mut self, cell: Cell) {
47                match (self, cell) {
48                    $((Values::$kind(v), Cell::$kind(x)) => v.push(x),)*
49                    (values, _) => {
50                        debug_assert!(false, "a cell of the wrong kind for {values:?}");
51                        values.push_null();
52                    }
53                }
54            }
55
56            fn push_null(&mut self) {
57                match self {
58                    $(Values::$kind(v) => v.push(None),)*
59                }
60            }
61
62            fn take(&mut self, name: PlSmallStr) -> PolarsResult<Series> {
63                match self {
64                    $(Values::$kind(v) => {
65                        let build: fn(PlSmallStr, Vec<Option<$value>>) -> PolarsResult<Series> =
66                            $build;
67                        build(name, std::mem::take(v))
68                    })*
69                }
70            }
71        }
72    };
73}
74
75kinds! {
76    /// Milliseconds since the epoch, UTC.
77    Time(i64) => DataType::Datetime(TimeUnit::Milliseconds, Some(TimeZone::UTC)), |name, v| {
78        Ok(Int64Chunked::from_iter_options(name, v.into_iter())
79            .into_datetime(TimeUnit::Milliseconds, Some(TimeZone::UTC))
80            .into_series())
81    };
82    /// Microseconds since the epoch, with no time zone.
83    DatetimeUs(i64) => DataType::Datetime(TimeUnit::Microseconds, None), |name, v| {
84        Ok(Int64Chunked::from_iter_options(name, v.into_iter())
85            .into_datetime(TimeUnit::Microseconds, None)
86            .into_series())
87    };
88    /// Microseconds since a start the file gives.
89    DurationUs(i64) => DataType::Duration(TimeUnit::Microseconds), |name, v| {
90        Ok(Int64Chunked::from_iter_options(name, v.into_iter())
91            .into_duration(TimeUnit::Microseconds)
92            .into_series())
93    };
94    F64(f64) => DataType::Float64, |name, v| Ok(Series::new(name, v));
95    U8(u8) => DataType::UInt8, |name, v| Ok(Series::new(name, v));
96    U16(u16) => DataType::UInt16, |name, v| Ok(Series::new(name, v));
97    U32(u32) => DataType::UInt32, |name, v| Ok(Series::new(name, v));
98    U64(u64) => DataType::UInt64, |name, v| Ok(Series::new(name, v));
99    I32(i32) => DataType::Int32, |name, v| Ok(Series::new(name, v));
100    Str(String) => DataType::String, |name, v| Ok(Series::new(name, v));
101    /// Text from a fixed set the reader names: no allocation a value.
102    Label(&'static str) => DataType::String, |name, v| Ok(Series::new(name, v));
103    /// Text many rows repeat (an ELF symbol's section), shared rather than copied.
104    Shared(std::sync::Arc<str>) => DataType::String, |name, v| {
105        Ok(StringChunked::from_iter_options(name, v.iter().map(|s| s.as_deref())).into_series())
106    };
107    Bool(bool) => DataType::Boolean, |name, v| Ok(Series::new(name, v));
108    Binary(Vec<u8>) => DataType::Binary, |name, v| {
109        Ok(BinaryChunked::from_iter_options(name, v.into_iter()).into_series())
110    };
111    /// A list of unsigned integers: the satellites a GSA sentence names.
112    ListU32(Vec<u32>) => DataType::List(Box::new(DataType::UInt32)), |name, v| {
113        let list: ListChunked = v
114            .into_iter()
115            .map(|items| items.map(|items| Series::new(PlSmallStr::EMPTY, items)))
116            .collect();
117        // A batch of nulls only would otherwise be a list of nulls, and one segment's
118        // type must be every batch's.
119        list.with_name(name)
120            .into_series()
121            .cast(&DataType::List(Box::new(DataType::UInt32)))
122    };
123}
124
125/// One column of `kind` named `name`, of `cells`.
126pub fn series(
127    name: &str,
128    kind: Kind,
129    cells: impl IntoIterator<Item = Cell>,
130) -> PolarsResult<Series> {
131    let mut values = Values::new(kind);
132    for cell in cells {
133        values.push(cell);
134    }
135    values.take(name.into())
136}
137
138/// Columns being filled, a row at a time.
139#[derive(Debug)]
140pub struct Builder {
141    names: Vec<String>,
142    values: Vec<Values>,
143    len: usize,
144}
145
146impl Builder {
147    pub fn new(columns: &[(&str, Kind)]) -> Self {
148        let mut builder = Self {
149            names: Vec::new(),
150            values: Vec::new(),
151            len: 0,
152        };
153        for (name, kind) in columns {
154            builder.add_column(name, *kind);
155        }
156        builder
157    }
158
159    /// Rows held, not yet taken.
160    pub fn len(&self) -> usize {
161        self.len
162    }
163
164    pub fn is_empty(&self) -> bool {
165        self.len == 0
166    }
167
168    /// A column added after rows were: null in each of them.
169    pub fn add_column(&mut self, name: &str, kind: Kind) -> usize {
170        let mut values = Values::new(kind);
171        for _ in 0..self.len {
172            values.push_null();
173        }
174        self.names.push(name.to_string());
175        self.values.push(values);
176        self.values.len() - 1
177    }
178
179    /// One row, a cell per column in order. Missing cells are null; extra ones are
180    /// dropped.
181    pub fn push(&mut self, row: impl IntoIterator<Item = Cell>) {
182        let mut row = row.into_iter();
183        for values in &mut self.values {
184            match row.next() {
185                Some(cell) => values.push(cell),
186                None => values.push_null(),
187            }
188        }
189        self.len += 1;
190    }
191
192    /// One row given as the column each value goes to; every other column is null.
193    pub fn push_sparse(&mut self, cells: impl IntoIterator<Item = (usize, Cell)>) {
194        let mut row: Vec<Option<Cell>> = vec![None; self.values.len()];
195        for (at, cell) in cells {
196            if let Some(slot) = row.get_mut(at) {
197                *slot = Some(cell);
198            }
199        }
200        for (values, cell) in self.values.iter_mut().zip(row) {
201            match cell {
202                Some(cell) => values.push(cell),
203                None => values.push_null(),
204            }
205        }
206        self.len += 1;
207    }
208
209    /// Set the time in column `column` of a row not yet taken. Used to date the rows
210    /// read before the log said what day it is.
211    pub fn set_time(&mut self, column: usize, row: usize, ms: i64) {
212        if let Some(Values::Time(v)) = self.values.get_mut(column)
213            && let Some(slot) = v.get_mut(row)
214        {
215            *slot = Some(ms);
216        }
217    }
218
219    /// The rows held, as a frame; the builder is left empty with the same columns.
220    pub fn take(&mut self) -> PolarsResult<DataFrame> {
221        let height = self.len;
222        let columns = self
223            .values
224            .iter_mut()
225            .zip(&self.names)
226            .map(|(values, name)| Ok(values.take(name.as_str().into())?.into_column()))
227            .collect::<PolarsResult<Vec<_>>>()?;
228        self.len = 0;
229        DataFrame::new(height, columns)
230    }
231}
232
233#[cfg(test)]
234mod tests {
235    use super::*;
236
237    #[test]
238    fn rows_become_typed_columns() {
239        let mut b = Builder::new(&[("t", Kind::Time), ("x", Kind::F64), ("s", Kind::ListU32)]);
240        b.push([
241            Cell::Time(Some(1_000)),
242            Cell::F64(Some(1.5)),
243            Cell::ListU32(Some(vec![3, 7])),
244        ]);
245        b.push([Cell::Time(None)]);
246        let late = b.add_column("late", Kind::Str);
247        b.push_sparse([(late, Cell::Str(Some("y".into())))]);
248        b.set_time(0, 1, 2_000);
249        let df = b.take().unwrap();
250        assert_eq!(df.height(), 3);
251        assert_eq!(df.column("t").unwrap().dtype(), &Kind::Time.dtype());
252        assert_eq!(df.column("s").unwrap().dtype(), &Kind::ListU32.dtype());
253        assert_eq!(df.column("late").unwrap().null_count(), 2);
254        assert_eq!(df.column("t").unwrap().null_count(), 1);
255        assert!(b.is_empty());
256        // Nulls only still give the declared type.
257        b.push([]);
258        let df = b.take().unwrap();
259        assert_eq!(df.column("s").unwrap().dtype(), &Kind::ListU32.dtype());
260    }
261}
datui_lib/formats/columns.rs

datui_lib/formats/
columns.rs
