Skip to main content

datui_lib/gps/
table.rs

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