1use polars::prelude::*;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum Kind {
11 Time,
13 F64,
14 U32,
15 U64,
16 I32,
17 Str,
18 Bool,
19 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#[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 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 list.with_name(name).into_series().cast(&kind.dtype())?
141 }
142 };
143 Ok(series.into_column())
144 }
145}
146
147#[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 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 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 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 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 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 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 b.push([]);
279 let df = b.take().unwrap();
280 assert_eq!(df.column("s").unwrap().dtype(), &Kind::ListU32.dtype());
281 }
282}