1use crate::{
2 entities::properties::{props::TPropError, tprop::IllegalPropType},
3 storage::lazy_vec::IllegalSet,
4};
5use arrow_schema::ArrowError;
6use bigdecimal::{num_bigint::BigInt, BigDecimal};
7use lazy_vec::LazyVec;
8use raphtory_api::core::{
9 entities::properties::prop::{prop_col::PropCol, AsPropRef, Prop, PropRef, PropType},
10 storage::arc_str::ArcStr,
11};
12use rustc_hash::FxHashMap;
13use std::{borrow::Cow, collections::HashMap, fmt::Debug, sync::Arc};
14use thiserror::Error;
15
16use crate::storage::string_col::StringCol;
17use raphtory_api::core::entities::properties::prop::{
18 IntoProp, PropArray, PropMapRef, PropNum, PropUnwrap,
19};
20
21pub mod lazy_vec;
22pub mod locked_view;
23mod string_col;
24pub mod timeindex;
25
26#[derive(Debug, Default)]
27pub struct TColumns {
28 t_props_log: Vec<PropColumn>,
29 num_rows: usize,
30}
31
32impl TColumns {
33 pub fn push<P: AsPropRef>(
34 &mut self,
35 props: impl IntoIterator<Item = (usize, P)>,
36 ) -> Result<Option<usize>, TPropError> {
37 let id = self.num_rows;
38 let mut has_props = false;
39
40 for (prop_id, prop) in props {
41 match self.t_props_log.get_mut(prop_id) {
42 Some(col) => col.push(prop.as_prop_ref())?,
43 None => {
44 let col = PropColumn::new(self.num_rows, prop.as_prop_ref());
45
46 self.t_props_log
47 .resize_with(prop_id + 1, || PropColumn::Empty(id));
48 self.t_props_log[prop_id] = col;
49 }
50 }
51
52 has_props = true;
53 }
54
55 if has_props {
56 self.num_rows += 1;
57
58 for col in self.t_props_log.iter_mut() {
59 col.grow(self.num_rows);
60 }
61
62 Ok(Some(id))
63 } else {
64 Ok(None)
65 }
66 }
67
68 pub fn ensure_column(&mut self, prop_id: usize) {
69 if self.t_props_log.len() <= prop_id {
70 self.t_props_log
71 .resize_with(prop_id + 1, || PropColumn::Empty(self.num_rows));
72 }
73 }
74
75 pub fn push_null(&mut self) -> usize {
76 let id = self.num_rows;
77 for col in self.t_props_log.iter_mut() {
78 col.push_null();
79 }
80 self.num_rows += 1;
81 id
82 }
83
84 pub fn get(&self, prop_id: usize) -> Option<&PropColumn> {
85 self.t_props_log.get(prop_id)
86 }
87
88 pub fn get_mut(&mut self, prop_id: usize) -> Option<&mut PropColumn> {
89 self.t_props_log.get_mut(prop_id)
90 }
91
92 pub fn getx(&self, prop_id: usize) -> Option<&PropColumn> {
93 self.t_props_log.get(prop_id)
94 }
95
96 pub fn len(&self) -> usize {
97 self.num_rows
98 }
99
100 pub fn is_empty(&self) -> bool {
101 self.num_rows == 0
102 }
103
104 pub fn iter(&self) -> impl Iterator<Item = &PropColumn> {
105 self.t_props_log.iter()
106 }
107
108 pub fn num_columns(&self) -> usize {
109 self.t_props_log.len()
110 }
111}
112
113#[derive(Debug)]
114pub enum PropColumn {
115 Empty(usize),
116 Bool(LazyVec<bool>),
117 U8(LazyVec<u8>),
118 U16(LazyVec<u16>),
119 U32(LazyVec<u32>),
120 U64(LazyVec<u64>),
121 I32(LazyVec<i32>),
122 I64(LazyVec<i64>),
123 F32(LazyVec<f32>),
124 F64(LazyVec<f64>),
125 Str(StringCol),
126 List(LazyVec<PropArray>),
127 Map(LazyVec<Arc<FxHashMap<ArcStr, Prop>>>),
128 NDTime(LazyVec<chrono::NaiveDateTime>),
129 DTime(LazyVec<chrono::DateTime<chrono::Utc>>),
130 Decimal(LazyVec<BigDecimal>),
131}
132
133#[derive(Error, Debug)]
134pub enum TPropColumnError {
135 #[error(transparent)]
136 IllegalSet(IllegalSet<Prop>),
137 #[error(transparent)]
138 IllegalType(#[from] IllegalPropType),
139 #[error(transparent)]
140 Arrow(#[from] ArrowError),
141}
142
143impl<A: Into<Prop> + Debug> From<IllegalSet<A>> for TPropColumnError {
144 fn from(value: IllegalSet<A>) -> Self {
145 let previous_value = value.previous_value.into();
146 let new_value = value.new_value.into();
147 TPropColumnError::IllegalSet(IllegalSet {
148 index: value.index,
149 previous_value,
150 new_value,
151 })
152 }
153}
154
155impl Default for PropColumn {
156 fn default() -> Self {
157 PropColumn::Empty(0)
158 }
159}
160
161impl PropColumn {
162 pub(crate) fn new(idx: usize, prop: PropRef<'_>) -> Self {
163 let mut col = PropColumn::default();
164 col.upsert(idx, prop).unwrap();
165 col
166 }
167
168 pub(crate) fn dtype(&self) -> PropType {
169 match self {
170 PropColumn::Empty(_) => PropType::Empty,
171 PropColumn::Bool(_) => PropType::Bool,
172 PropColumn::U8(_) => PropType::U8,
173 PropColumn::U16(_) => PropType::U16,
174 PropColumn::U32(_) => PropType::U32,
175 PropColumn::U64(_) => PropType::U64,
176 PropColumn::I32(_) => PropType::I32,
177 PropColumn::I64(_) => PropType::I64,
178 PropColumn::F32(_) => PropType::F32,
179 PropColumn::F64(_) => PropType::F64,
180 PropColumn::Str(_) => PropType::Str,
181 PropColumn::List(_) => PropType::List(Box::new(PropType::Empty)),
182 PropColumn::Map(_) => PropType::Map(HashMap::new().into()),
183 PropColumn::NDTime(_) => PropType::NDTime,
184 PropColumn::DTime(_) => PropType::DTime,
185 PropColumn::Decimal(_) => PropType::Decimal { scale: 0 },
186 }
187 }
188
189 pub(crate) fn grow(&mut self, new_len: usize) {
190 while self.len() < new_len {
191 self.push_null();
192 }
193 }
194
195 pub fn upsert(&mut self, index: usize, prop: PropRef<'_>) -> Result<(), TPropColumnError> {
196 self.init_empty_col(&prop);
197 match (self, prop) {
198 (PropColumn::Bool(col), PropRef::Bool(v)) => col.upsert(index, v),
199 (PropColumn::I64(col), PropRef::Num(PropNum::I64(v))) => col.upsert(index, v),
200 (PropColumn::U32(col), PropRef::Num(PropNum::U32(v))) => col.upsert(index, v),
201 (PropColumn::U64(col), PropRef::Num(PropNum::U64(v))) => col.upsert(index, v),
202 (PropColumn::F32(col), PropRef::Num(PropNum::F32(v))) => col.upsert(index, v),
203 (PropColumn::F64(col), PropRef::Num(PropNum::F64(v))) => col.upsert(index, v),
204 (PropColumn::Str(col), PropRef::Str(v)) => col.upsert(index, v)?,
205 (PropColumn::U8(col), PropRef::Num(PropNum::U8(v))) => col.upsert(index, v),
206 (PropColumn::U16(col), PropRef::Num(PropNum::U16(v))) => col.upsert(index, v),
207 (PropColumn::I32(col), PropRef::Num(PropNum::I32(v))) => col.upsert(index, v),
208 (PropColumn::List(col), PropRef::List(v)) => col.upsert(index, v.into_owned()),
209 (PropColumn::Map(col), PropRef::Map(v)) => match v {
210 PropMapRef::Mem(map) => col.upsert(index, map.clone()),
211 PropMapRef::PropCol { map, i } => {
212 if let Some(entry) = map.get(i).and_then(|prop| prop.into_map()) {
213 col.upsert(index, entry);
214 }
215 }
216 PropMapRef::Arrow(arc_map) => {
217 if let Some(prop) = arc_map.into_prop() {
218 if let Some(map_ref) = prop.as_prop_ref().as_map_ref() {
219 if let Some(map) = map_ref.as_map() {
220 col.upsert(index, map.clone());
221 }
222 }
223 }
224 }
225 },
226 (PropColumn::NDTime(col), PropRef::NDTime(v)) => col.upsert(index, v),
227 (PropColumn::DTime(col), PropRef::DTime(v)) => col.upsert(index, v),
228 (PropColumn::Decimal(col), PropRef::Decimal { num, scale }) => {
229 col.upsert(index, BigDecimal::from_bigint(num.into(), scale as i64))
230 }
231 (col, prop) => {
232 Err(IllegalPropType {
233 expected: col.dtype(),
234 actual: prop.into_prop().dtype(),
235 })?;
236 }
237 }
238 Ok(())
239 }
240
241 pub fn check(&self, index: usize, prop: &PropRef<'_>) -> Result<(), TPropColumnError> {
242 match (self, prop) {
243 (PropColumn::Empty(_), _) => {}
244 (PropColumn::Bool(col), PropRef::Bool(v)) => col.check(index, v)?,
245 (PropColumn::I64(col), PropRef::Num(PropNum::I64(v))) => col.check(index, v)?,
246 (PropColumn::U32(col), PropRef::Num(PropNum::U32(v))) => col.check(index, v)?,
247 (PropColumn::U64(col), PropRef::Num(PropNum::U64(v))) => col.check(index, v)?,
248 (PropColumn::F32(col), PropRef::Num(PropNum::F32(v))) => col.check(index, v)?,
249 (PropColumn::F64(col), PropRef::Num(PropNum::F64(v))) => col.check(index, v)?,
250 (PropColumn::Str(col), PropRef::Str(v)) => col.check(index, v)?,
251 (PropColumn::U8(col), PropRef::Num(PropNum::U8(v))) => col.check(index, v)?,
252 (PropColumn::U16(col), PropRef::Num(PropNum::U16(v))) => col.check(index, v)?,
253 (PropColumn::I32(col), PropRef::Num(PropNum::I32(v))) => col.check(index, v)?,
254 (PropColumn::List(col), PropRef::List(v)) => col.check(index, v)?,
255 (PropColumn::Map(col), PropRef::Map(v)) => col.check(index, &v.as_mem())?,
256 (PropColumn::NDTime(col), PropRef::NDTime(v)) => col.check(index, v)?,
257 (PropColumn::DTime(col), PropRef::DTime(v)) => col.check(index, v)?,
258 (PropColumn::Decimal(col), PropRef::Decimal { num, scale }) => col.check(
259 index,
260 &BigDecimal::from_bigint(BigInt::from(*num), *scale as i64),
261 )?,
262 (col, prop) => {
263 Err(IllegalPropType {
264 expected: col.dtype(),
265 actual: prop.clone().into_prop().dtype(),
266 })?;
267 }
268 }
269 Ok(())
270 }
271
272 pub(crate) fn push(&mut self, prop: PropRef<'_>) -> Result<(), TPropColumnError> {
273 self.init_empty_col(&prop);
274 match (self, prop) {
275 (PropColumn::Bool(col), PropRef::Bool(v)) => col.push(Some(v)),
276 (PropColumn::U8(col), PropRef::Num(PropNum::U8(v))) => col.push(Some(v)),
277 (PropColumn::I64(col), PropRef::Num(PropNum::I64(v))) => col.push(Some(v)),
278 (PropColumn::U32(col), PropRef::Num(PropNum::U32(v))) => col.push(Some(v)),
279 (PropColumn::U64(col), PropRef::Num(PropNum::U64(v))) => col.push(Some(v)),
280 (PropColumn::F32(col), PropRef::Num(PropNum::F32(v))) => col.push(Some(v)),
281 (PropColumn::F64(col), PropRef::Num(PropNum::F64(v))) => col.push(Some(v)),
282 (PropColumn::Str(col), PropRef::Str(v)) => col.push_value(v)?,
283 (PropColumn::U16(col), PropRef::Num(PropNum::U16(v))) => col.push(Some(v)),
284 (PropColumn::I32(col), PropRef::Num(PropNum::I32(v))) => col.push(Some(v)),
285 (PropColumn::List(col), PropRef::List(v)) => col.push(Some(v.into_owned())),
286 (PropColumn::Map(col), PropRef::Map(v)) => {
287 match v {
289 PropMapRef::Mem(map) => col.push(Some(map.clone())),
290 PropMapRef::PropCol { map, i } => {
291 col.push(map.get(i).and_then(|prop| prop.into_map()))
292 }
293 PropMapRef::Arrow(arc_map) => {
294 if let Some(prop) = arc_map.into_prop() {
295 if let Some(map_ref) = prop.as_prop_ref().as_map_ref() {
296 if let Some(map) = map_ref.as_map() {
297 col.push(Some(map.clone()));
298 }
299 }
300 }
301 }
302 }
303 }
304 (PropColumn::NDTime(col), PropRef::NDTime(v)) => col.push(Some(v)),
305 (PropColumn::DTime(col), PropRef::DTime(v)) => col.push(Some(v)),
306 (PropColumn::Decimal(col), PropRef::Decimal { num, scale }) => {
307 col.push(Some(BigDecimal::from_bigint(num.into(), scale as i64)))
308 }
309 (col, prop) => {
310 Err(IllegalPropType {
311 expected: col.dtype(),
312 actual: prop.into_prop().dtype(),
313 })?;
314 }
315 }
316 Ok(())
317 }
318
319 fn init_empty_col(&mut self, prop: &PropRef<'_>) {
320 if let PropColumn::Empty(len) = self {
321 match prop {
322 PropRef::Bool(_) => *self = PropColumn::Bool(LazyVec::with_len(*len)),
323 PropRef::Num(PropNum::I64(_)) => *self = PropColumn::I64(LazyVec::with_len(*len)),
324 PropRef::Num(PropNum::U32(_)) => *self = PropColumn::U32(LazyVec::with_len(*len)),
325 PropRef::Num(PropNum::U64(_)) => *self = PropColumn::U64(LazyVec::with_len(*len)),
326 PropRef::Num(PropNum::F32(_)) => *self = PropColumn::F32(LazyVec::with_len(*len)),
327 PropRef::Num(PropNum::F64(_)) => *self = PropColumn::F64(LazyVec::with_len(*len)),
328 PropRef::Str(_) => *self = PropColumn::Str(StringCol::with_len(*len)),
329 PropRef::Num(PropNum::U8(_)) => *self = PropColumn::U8(LazyVec::with_len(*len)),
330 PropRef::Num(PropNum::U16(_)) => *self = PropColumn::U16(LazyVec::with_len(*len)),
331 PropRef::Num(PropNum::I32(_)) => *self = PropColumn::I32(LazyVec::with_len(*len)),
332 PropRef::List(_) => *self = PropColumn::List(LazyVec::with_len(*len)),
333 PropRef::Map(_) => *self = PropColumn::Map(LazyVec::with_len(*len)),
334 PropRef::NDTime(_) => *self = PropColumn::NDTime(LazyVec::with_len(*len)),
335 PropRef::DTime(_) => *self = PropColumn::DTime(LazyVec::with_len(*len)),
336 PropRef::Decimal { .. } => *self = PropColumn::Decimal(LazyVec::with_len(*len)),
337 }
338 }
339 }
340
341 pub fn is_empty(&self) -> bool {
342 matches!(self, PropColumn::Empty(_))
343 }
344
345 pub(crate) fn push_null(&mut self) {
346 match self {
347 PropColumn::Bool(col) => col.push(None),
348 PropColumn::I64(col) => col.push(None),
349 PropColumn::U32(col) => col.push(None),
350 PropColumn::U64(col) => col.push(None),
351 PropColumn::F32(col) => col.push(None),
352 PropColumn::F64(col) => col.push(None),
353 PropColumn::Str(col) => col.push_null(),
354 PropColumn::U8(col) => col.push(None),
355 PropColumn::U16(col) => col.push(None),
356 PropColumn::I32(col) => col.push(None),
357 PropColumn::List(col) => col.push(None),
358 PropColumn::Map(col) => col.push(None),
359 PropColumn::NDTime(col) => col.push(None),
360 PropColumn::DTime(col) => col.push(None),
361 PropColumn::Decimal(col) => col.push(None),
362 PropColumn::Empty(count) => {
363 *count += 1;
364 }
365 }
366 }
367
368 pub fn get(&self, index: usize) -> Option<Prop> {
369 match self {
370 PropColumn::Bool(col) => col.get_opt(index).map(|prop| (*prop).into()),
371 PropColumn::I64(col) => col.get_opt(index).map(|prop| (*prop).into()),
372 PropColumn::U32(col) => col.get_opt(index).map(|prop| (*prop).into()),
373 PropColumn::U64(col) => col.get_opt(index).map(|prop| (*prop).into()),
374 PropColumn::F32(col) => col.get_opt(index).map(|prop| (*prop).into()),
375 PropColumn::F64(col) => col.get_opt(index).map(|prop| (*prop).into()),
376 PropColumn::Str(col) => col.get_opt(index).map(|prop| prop.into()),
377 PropColumn::U8(col) => col.get_opt(index).map(|prop| (*prop).into()),
378 PropColumn::U16(col) => col.get_opt(index).map(|prop| (*prop).into()),
379 PropColumn::I32(col) => col.get_opt(index).map(|prop| (*prop).into()),
380 PropColumn::List(col) => col.get_opt(index).map(|prop| Prop::List(prop.clone())),
381 PropColumn::Map(col) => col.get_opt(index).map(|prop| Prop::Map(prop.clone())),
382 PropColumn::NDTime(col) => col.get_opt(index).map(|prop| Prop::NDTime(*prop)),
383 PropColumn::DTime(col) => col.get_opt(index).map(|prop| Prop::DTime(*prop)),
384 PropColumn::Decimal(col) => col.get_opt(index).map(|prop| Prop::Decimal(prop.clone())),
385 PropColumn::Empty(_) => None,
386 }
387 }
388
389 pub fn get_ref(&self, index: usize) -> Option<PropRef<'_>> {
390 match self {
391 PropColumn::Bool(col) => col.get_opt(index).map(|prop| PropRef::Bool(*prop)),
392 PropColumn::I64(col) => col.get_opt(index).map(|prop| PropRef::from(*prop)),
393 PropColumn::U32(col) => col.get_opt(index).map(|prop| PropRef::from(*prop)),
394 PropColumn::U64(col) => col.get_opt(index).map(|prop| PropRef::from(*prop)),
395 PropColumn::F32(col) => col.get_opt(index).map(|prop| PropRef::from(*prop)),
396 PropColumn::F64(col) => col.get_opt(index).map(|prop| PropRef::from(*prop)),
397 PropColumn::Str(col) => col.get_opt(index).map(|prop| PropRef::Str(prop.as_ref())),
398 PropColumn::U8(col) => col.get_opt(index).map(|prop| PropRef::from(*prop)),
399 PropColumn::U16(col) => col.get_opt(index).map(|prop| PropRef::from(*prop)),
400 PropColumn::I32(col) => col.get_opt(index).map(|prop| PropRef::from(*prop)),
401 PropColumn::List(col) => col
402 .get_opt(index)
403 .map(|prop| PropRef::List(Cow::Borrowed(prop))),
404 PropColumn::Map(col) => col.get_opt(index).map(PropRef::from),
405 PropColumn::NDTime(col) => col.get_opt(index).copied().map(PropRef::from),
406 PropColumn::DTime(col) => col.get_opt(index).copied().map(PropRef::from),
407 PropColumn::Decimal(col) => col.get_opt(index).map(PropRef::from),
408 PropColumn::Empty(_) => None,
409 }
410 }
411
412 pub(crate) fn len(&self) -> usize {
413 match self {
414 PropColumn::Bool(col) => col.len(),
415 PropColumn::I64(col) => col.len(),
416 PropColumn::U32(col) => col.len(),
417 PropColumn::U64(col) => col.len(),
418 PropColumn::F32(col) => col.len(),
419 PropColumn::F64(col) => col.len(),
420 PropColumn::Str(col) => col.len(),
421 PropColumn::U8(col) => col.len(),
422 PropColumn::U16(col) => col.len(),
423 PropColumn::I32(col) => col.len(),
424 PropColumn::List(col) => col.len(),
425 PropColumn::Map(col) => col.len(),
426 PropColumn::NDTime(col) => col.len(),
427 PropColumn::DTime(col) => col.len(),
428 PropColumn::Decimal(col) => col.len(),
429 PropColumn::Empty(count) => *count,
430 }
431 }
432}
433
434#[cfg(test)]
435mod test {
436 use super::TColumns;
437 use raphtory_api::core::entities::properties::prop::Prop;
438
439 #[test]
440 fn tcolumns_append_1() {
441 let mut t_cols = TColumns::default();
442
443 t_cols.push([(1, Prop::U64(1))]).unwrap();
444
445 let col0 = t_cols.get(0).unwrap();
446 let col1 = t_cols.get(1).unwrap();
447
448 assert_eq!(col0.len(), 1);
449 assert_eq!(col1.len(), 1);
450 }
451
452 #[test]
453 fn tcolumns_append_3_rows() {
454 let mut t_cols = TColumns::default();
455
456 t_cols
457 .push([(1, Prop::U64(1)), (0, Prop::Str("a".into()))])
458 .unwrap();
459 t_cols
460 .push([(0, Prop::Str("c".into())), (2, Prop::I64(9))])
461 .unwrap();
462 t_cols
463 .push([(1, Prop::U64(1)), (3, Prop::Str("c".into()))])
464 .unwrap();
465
466 assert_eq!(t_cols.len(), 3);
467
468 for col_id in 0..4 {
469 let col = t_cols.get(col_id).unwrap();
470 assert_eq!(col.len(), 3);
471 }
472
473 let col0 = (0..3)
474 .map(|row| t_cols.get(0).and_then(|col| col.get(row)))
475 .collect::<Vec<_>>();
476 assert_eq!(
477 col0,
478 vec![
479 Some(Prop::Str("a".into())),
480 Some(Prop::Str("c".into())),
481 None
482 ]
483 );
484
485 let col1 = (0..3)
486 .map(|row| t_cols.get(1).and_then(|col| col.get(row)))
487 .collect::<Vec<_>>();
488 assert_eq!(col1, vec![Some(Prop::U64(1)), None, Some(Prop::U64(1))]);
489
490 let col2 = (0..3)
491 .map(|row| t_cols.get(2).and_then(|col| col.get(row)))
492 .collect::<Vec<_>>();
493 assert_eq!(col2, vec![None, Some(Prop::I64(9)), None]);
494
495 let col3 = (0..3)
496 .map(|row| t_cols.get(3).and_then(|col| col.get(row)))
497 .collect::<Vec<_>>();
498 assert_eq!(col3, vec![None, None, Some(Prop::Str("c".into()))]);
499 }
500
501 #[test]
502 fn tcolumns_append_2_columns_12_items() {
503 let mut t_cols = TColumns::default();
504
505 for value in 0..12 {
506 if value % 2 == 0 {
507 t_cols
508 .push([
509 (1, Prop::U64(value)),
510 (0, Prop::Str(value.to_string().into())),
511 ])
512 .unwrap();
513 } else {
514 t_cols.push([(1, Prop::U64(value))]).unwrap();
515 }
516 }
517
518 assert_eq!(t_cols.len(), 12);
519
520 let col0 = (0..12)
521 .map(|row| t_cols.get(0).and_then(|col| col.get(row)))
522 .collect::<Vec<_>>();
523 assert_eq!(
524 col0,
525 vec![
526 Some(Prop::Str("0".into())),
527 None,
528 Some(Prop::Str("2".into())),
529 None,
530 Some(Prop::Str("4".into())),
531 None,
532 Some(Prop::Str("6".into())),
533 None,
534 Some(Prop::Str("8".into())),
535 None,
536 Some(Prop::Str("10".into())),
537 None
538 ]
539 );
540
541 let col1 = (0..12)
542 .map(|row| t_cols.get(1).and_then(|col| col.get(row)))
543 .collect::<Vec<_>>();
544 assert_eq!(
545 col1,
546 vec![
547 Some(Prop::U64(0)),
548 Some(Prop::U64(1)),
549 Some(Prop::U64(2)),
550 Some(Prop::U64(3)),
551 Some(Prop::U64(4)),
552 Some(Prop::U64(5)),
553 Some(Prop::U64(6)),
554 Some(Prop::U64(7)),
555 Some(Prop::U64(8)),
556 Some(Prop::U64(9)),
557 Some(Prop::U64(10)),
558 Some(Prop::U64(11))
559 ]
560 );
561 }
562}