1use arcstr::ArcStr;
4
5use grafeo_common::types::{EdgeId, LogicalType, NodeId, Value};
6
7pub const DEFAULT_VECTOR_CAPACITY: usize = 2048;
9
10#[derive(Debug, Clone)]
15pub struct ValueVector {
16 data_type: LogicalType,
18 data: VectorData,
20 len: usize,
22 validity: Option<Vec<bool>>,
24}
25
26#[derive(Debug, Clone)]
28enum VectorData {
29 Bool(Vec<bool>),
31 Int64(Vec<i64>),
33 Float64(Vec<f64>),
35 String(Vec<ArcStr>),
37 NodeId(Vec<NodeId>),
39 EdgeId(Vec<EdgeId>),
41 Generic(Vec<Value>),
43}
44
45impl ValueVector {
46 #[must_use]
48 pub fn new() -> Self {
49 Self::with_capacity(LogicalType::Any, DEFAULT_VECTOR_CAPACITY)
50 }
51
52 #[must_use]
54 pub fn with_type(data_type: LogicalType) -> Self {
55 Self::with_capacity(data_type, DEFAULT_VECTOR_CAPACITY)
56 }
57
58 pub fn from_values(values: &[Value]) -> Self {
60 let mut vec = Self::new();
61 for value in values {
62 vec.push_value(value.clone());
63 }
64 vec
65 }
66
67 #[must_use]
69 pub fn with_capacity(data_type: LogicalType, capacity: usize) -> Self {
70 let data = match &data_type {
71 LogicalType::Bool => VectorData::Bool(Vec::with_capacity(capacity)),
72 LogicalType::Int8 | LogicalType::Int16 | LogicalType::Int32 | LogicalType::Int64 => {
73 VectorData::Int64(Vec::with_capacity(capacity))
74 }
75 LogicalType::Float32 | LogicalType::Float64 => {
76 VectorData::Float64(Vec::with_capacity(capacity))
77 }
78 LogicalType::String => VectorData::String(Vec::with_capacity(capacity)),
79 LogicalType::Node => VectorData::NodeId(Vec::with_capacity(capacity)),
80 LogicalType::Edge => VectorData::EdgeId(Vec::with_capacity(capacity)),
81 _ => VectorData::Generic(Vec::with_capacity(capacity)),
82 };
83
84 Self {
85 data_type,
86 data,
87 len: 0,
88 validity: None,
89 }
90 }
91
92 #[must_use]
94 pub fn data_type(&self) -> &LogicalType {
95 &self.data_type
96 }
97
98 #[must_use]
100 pub fn len(&self) -> usize {
101 self.len
102 }
103
104 #[must_use]
106 pub fn is_empty(&self) -> bool {
107 self.len == 0
108 }
109
110 #[must_use]
112 pub fn is_null(&self, index: usize) -> bool {
113 self.validity
114 .as_ref()
115 .map_or(false, |v| !v.get(index).copied().unwrap_or(true))
116 }
117
118 pub fn set_null(&mut self, index: usize) {
120 if self.validity.is_none() {
121 self.validity = Some(vec![true; index + 1]);
122 }
123 if let Some(validity) = &mut self.validity {
124 if validity.len() <= index {
125 validity.resize(index + 1, true);
126 }
127 validity[index] = false;
128 }
129 }
130
131 pub fn push_bool(&mut self, value: bool) {
133 match &mut self.data {
134 VectorData::Bool(vec) => {
135 vec.push(value);
136 self.len += 1;
137 }
138 VectorData::Generic(vec) => {
139 vec.push(Value::Bool(value));
140 self.len += 1;
141 }
142 _ => {}
143 }
144 }
145
146 pub fn push_int64(&mut self, value: i64) {
148 match &mut self.data {
149 VectorData::Int64(vec) => {
150 vec.push(value);
151 self.len += 1;
152 }
153 VectorData::Generic(vec) => {
154 vec.push(Value::Int64(value));
155 self.len += 1;
156 }
157 _ => {}
158 }
159 }
160
161 pub fn push_float64(&mut self, value: f64) {
163 match &mut self.data {
164 VectorData::Float64(vec) => {
165 vec.push(value);
166 self.len += 1;
167 }
168 VectorData::Generic(vec) => {
169 vec.push(Value::Float64(value));
170 self.len += 1;
171 }
172 _ => {}
173 }
174 }
175
176 pub fn push_string(&mut self, value: impl Into<ArcStr>) {
178 match &mut self.data {
179 VectorData::String(vec) => {
180 vec.push(value.into());
181 self.len += 1;
182 }
183 VectorData::Generic(vec) => {
184 vec.push(Value::String(value.into()));
185 self.len += 1;
186 }
187 _ => {}
188 }
189 }
190
191 pub fn push_node_id(&mut self, value: NodeId) {
193 match &mut self.data {
194 VectorData::NodeId(vec) => {
195 vec.push(value);
196 self.len += 1;
197 }
198 VectorData::Generic(vec) => {
199 vec.push(Value::Int64(value.as_u64() as i64));
200 self.len += 1;
201 }
202 _ => {}
203 }
204 }
205
206 pub fn push_edge_id(&mut self, value: EdgeId) {
208 match &mut self.data {
209 VectorData::EdgeId(vec) => {
210 vec.push(value);
211 self.len += 1;
212 }
213 VectorData::Generic(vec) => {
214 vec.push(Value::Int64(value.as_u64() as i64));
215 self.len += 1;
216 }
217 _ => {}
218 }
219 }
220
221 pub fn push_value(&mut self, value: Value) {
223 if matches!(value, Value::Null) {
225 match &mut self.data {
226 VectorData::Bool(vec) => vec.push(false),
227 VectorData::Int64(vec) => vec.push(0),
228 VectorData::Float64(vec) => vec.push(0.0),
229 VectorData::String(vec) => vec.push("".into()),
230 VectorData::NodeId(vec) => vec.push(NodeId::new(0)),
231 VectorData::EdgeId(vec) => vec.push(EdgeId::new(0)),
232 VectorData::Generic(vec) => vec.push(Value::Null),
233 }
234 self.len += 1;
235 self.set_null(self.len - 1);
236 return;
237 }
238
239 match (&mut self.data, &value) {
240 (VectorData::Bool(vec), Value::Bool(b)) => vec.push(*b),
241 (VectorData::Int64(vec), Value::Int64(i)) => vec.push(*i),
242 (VectorData::Float64(vec), Value::Float64(f)) => vec.push(*f),
243 (VectorData::String(vec), Value::String(s)) => vec.push(s.clone()),
244 (VectorData::NodeId(vec), Value::Int64(i)) => vec.push(NodeId::new(*i as u64)),
246 (VectorData::EdgeId(vec), Value::Int64(i)) => vec.push(EdgeId::new(*i as u64)),
248 (VectorData::Generic(vec), _) => vec.push(value),
249 _ => {
250 match &mut self.data {
252 VectorData::Bool(vec) => vec.push(false),
253 VectorData::Int64(vec) => vec.push(0),
254 VectorData::Float64(vec) => vec.push(0.0),
255 VectorData::String(vec) => vec.push("".into()),
256 VectorData::NodeId(vec) => vec.push(NodeId::new(0)),
257 VectorData::EdgeId(vec) => vec.push(EdgeId::new(0)),
258 VectorData::Generic(vec) => vec.push(value),
259 }
260 }
261 }
262 self.len += 1;
263 }
264
265 #[must_use]
267 pub fn get_bool(&self, index: usize) -> Option<bool> {
268 if self.is_null(index) {
269 return None;
270 }
271 if let VectorData::Bool(vec) = &self.data {
272 vec.get(index).copied()
273 } else {
274 None
275 }
276 }
277
278 #[must_use]
280 pub fn get_int64(&self, index: usize) -> Option<i64> {
281 if self.is_null(index) {
282 return None;
283 }
284 if let VectorData::Int64(vec) = &self.data {
285 vec.get(index).copied()
286 } else {
287 None
288 }
289 }
290
291 #[must_use]
293 pub fn get_float64(&self, index: usize) -> Option<f64> {
294 if self.is_null(index) {
295 return None;
296 }
297 if let VectorData::Float64(vec) = &self.data {
298 vec.get(index).copied()
299 } else {
300 None
301 }
302 }
303
304 #[must_use]
306 pub fn get_string(&self, index: usize) -> Option<&str> {
307 if self.is_null(index) {
308 return None;
309 }
310 if let VectorData::String(vec) = &self.data {
311 vec.get(index).map(|s| s.as_ref())
312 } else {
313 None
314 }
315 }
316
317 #[must_use]
319 pub fn get_node_id(&self, index: usize) -> Option<NodeId> {
320 if self.is_null(index) {
321 return None;
322 }
323 match &self.data {
324 VectorData::NodeId(vec) => vec.get(index).copied(),
325 VectorData::Generic(vec) => match vec.get(index) {
327 Some(Value::Int64(i)) => Some(NodeId::new(*i as u64)),
328 _ => None,
329 },
330 _ => None,
331 }
332 }
333
334 #[must_use]
336 pub fn get_edge_id(&self, index: usize) -> Option<EdgeId> {
337 if self.is_null(index) {
338 return None;
339 }
340 match &self.data {
341 VectorData::EdgeId(vec) => vec.get(index).copied(),
342 VectorData::Generic(vec) => match vec.get(index) {
344 Some(Value::Int64(i)) => Some(EdgeId::new(*i as u64)),
345 _ => None,
346 },
347 _ => None,
348 }
349 }
350
351 #[must_use]
353 pub fn get_value(&self, index: usize) -> Option<Value> {
354 if self.is_null(index) {
355 return Some(Value::Null);
356 }
357
358 match &self.data {
359 VectorData::Bool(vec) => vec.get(index).map(|&v| Value::Bool(v)),
360 VectorData::Int64(vec) => vec.get(index).map(|&v| Value::Int64(v)),
361 VectorData::Float64(vec) => vec.get(index).map(|&v| Value::Float64(v)),
362 VectorData::String(vec) => vec.get(index).map(|v| Value::String(v.clone())),
363 VectorData::NodeId(vec) => vec.get(index).map(|&v| Value::Int64(v.as_u64() as i64)),
364 VectorData::EdgeId(vec) => vec.get(index).map(|&v| Value::Int64(v.as_u64() as i64)),
365 VectorData::Generic(vec) => vec.get(index).cloned(),
366 }
367 }
368
369 #[must_use]
371 pub fn get(&self, index: usize) -> Option<Value> {
372 self.get_value(index)
373 }
374
375 pub fn push(&mut self, value: Value) {
377 self.push_value(value);
378 }
379
380 #[must_use]
382 pub fn as_bool_slice(&self) -> Option<&[bool]> {
383 if let VectorData::Bool(vec) = &self.data {
384 Some(vec)
385 } else {
386 None
387 }
388 }
389
390 #[must_use]
392 pub fn as_int64_slice(&self) -> Option<&[i64]> {
393 if let VectorData::Int64(vec) = &self.data {
394 Some(vec)
395 } else {
396 None
397 }
398 }
399
400 #[must_use]
402 pub fn as_float64_slice(&self) -> Option<&[f64]> {
403 if let VectorData::Float64(vec) = &self.data {
404 Some(vec)
405 } else {
406 None
407 }
408 }
409
410 #[must_use]
412 pub fn as_node_id_slice(&self) -> Option<&[NodeId]> {
413 if let VectorData::NodeId(vec) = &self.data {
414 Some(vec)
415 } else {
416 None
417 }
418 }
419
420 #[must_use]
422 pub fn as_edge_id_slice(&self) -> Option<&[EdgeId]> {
423 if let VectorData::EdgeId(vec) = &self.data {
424 Some(vec)
425 } else {
426 None
427 }
428 }
429
430 #[must_use]
432 pub fn logical_type(&self) -> LogicalType {
433 self.data_type.clone()
434 }
435
436 pub fn copy_row_to(&self, row: usize, dest: &mut ValueVector) {
441 if self.is_null(row) {
442 dest.push_value(Value::Null);
443 return;
444 }
445
446 match &self.data {
447 VectorData::Bool(vec) => {
448 if let Some(&v) = vec.get(row) {
449 dest.push_bool(v);
450 }
451 }
452 VectorData::Int64(vec) => {
453 if let Some(&v) = vec.get(row) {
454 dest.push_int64(v);
455 }
456 }
457 VectorData::Float64(vec) => {
458 if let Some(&v) = vec.get(row) {
459 dest.push_float64(v);
460 }
461 }
462 VectorData::String(vec) => {
463 if let Some(v) = vec.get(row) {
464 dest.push_string(v.clone());
465 }
466 }
467 VectorData::NodeId(vec) => {
468 if let Some(&v) = vec.get(row) {
469 dest.push_node_id(v);
470 }
471 }
472 VectorData::EdgeId(vec) => {
473 if let Some(&v) = vec.get(row) {
474 dest.push_edge_id(v);
475 }
476 }
477 VectorData::Generic(vec) => {
478 if let Some(v) = vec.get(row) {
479 dest.push_value(v.clone());
480 }
481 }
482 }
483 }
484
485 pub fn clear(&mut self) {
487 match &mut self.data {
488 VectorData::Bool(vec) => vec.clear(),
489 VectorData::Int64(vec) => vec.clear(),
490 VectorData::Float64(vec) => vec.clear(),
491 VectorData::String(vec) => vec.clear(),
492 VectorData::NodeId(vec) => vec.clear(),
493 VectorData::EdgeId(vec) => vec.clear(),
494 VectorData::Generic(vec) => vec.clear(),
495 }
496 self.len = 0;
497 self.validity = None;
498 }
499}
500
501impl Default for ValueVector {
502 fn default() -> Self {
503 Self::new()
504 }
505}
506
507#[cfg(test)]
508mod tests {
509 use super::*;
510
511 #[test]
512 fn test_int64_vector() {
513 let mut vec = ValueVector::with_type(LogicalType::Int64);
514
515 vec.push_int64(1);
516 vec.push_int64(2);
517 vec.push_int64(3);
518
519 assert_eq!(vec.len(), 3);
520 assert_eq!(vec.get_int64(0), Some(1));
521 assert_eq!(vec.get_int64(1), Some(2));
522 assert_eq!(vec.get_int64(2), Some(3));
523 }
524
525 #[test]
526 fn test_string_vector() {
527 let mut vec = ValueVector::with_type(LogicalType::String);
528
529 vec.push_string("hello");
530 vec.push_string("world");
531
532 assert_eq!(vec.len(), 2);
533 assert_eq!(vec.get_string(0), Some("hello"));
534 assert_eq!(vec.get_string(1), Some("world"));
535 }
536
537 #[test]
538 fn test_null_values() {
539 let mut vec = ValueVector::with_type(LogicalType::Int64);
540
541 vec.push_int64(1);
542 vec.push_int64(2);
543 vec.push_int64(3);
544
545 assert!(!vec.is_null(1));
546 vec.set_null(1);
547 assert!(vec.is_null(1));
548
549 assert_eq!(vec.get_int64(0), Some(1));
550 assert_eq!(vec.get_int64(1), None); assert_eq!(vec.get_int64(2), Some(3));
552 }
553
554 #[test]
555 fn test_get_value() {
556 let mut vec = ValueVector::with_type(LogicalType::Int64);
557 vec.push_int64(42);
558
559 let value = vec.get_value(0);
560 assert_eq!(value, Some(Value::Int64(42)));
561 }
562
563 #[test]
564 fn test_slice_access() {
565 let mut vec = ValueVector::with_type(LogicalType::Int64);
566 vec.push_int64(1);
567 vec.push_int64(2);
568 vec.push_int64(3);
569
570 let slice = vec.as_int64_slice().unwrap();
571 assert_eq!(slice, &[1, 2, 3]);
572 }
573
574 #[test]
578 fn test_generic_fallback_push_int64() {
579 let mut vec = ValueVector::with_type(LogicalType::Any);
580 vec.push_int64(42);
581 vec.push_int64(-7);
582 assert_eq!(vec.len(), 2);
583 assert_eq!(vec.get_value(0), Some(Value::Int64(42)));
584 assert_eq!(vec.get_value(1), Some(Value::Int64(-7)));
585 }
586
587 #[test]
588 fn test_generic_fallback_push_bool() {
589 let mut vec = ValueVector::with_type(LogicalType::Any);
590 vec.push_bool(true);
591 vec.push_bool(false);
592 assert_eq!(vec.len(), 2);
593 assert_eq!(vec.get_value(0), Some(Value::Bool(true)));
594 assert_eq!(vec.get_value(1), Some(Value::Bool(false)));
595 }
596
597 #[test]
598 fn test_generic_fallback_push_float64() {
599 let mut vec = ValueVector::with_type(LogicalType::Any);
600 vec.push_float64(1.23);
601 vec.push_float64(-0.5);
602 assert_eq!(vec.len(), 2);
603 assert_eq!(vec.get_value(0), Some(Value::Float64(1.23)));
604 assert_eq!(vec.get_value(1), Some(Value::Float64(-0.5)));
605 }
606
607 #[test]
608 fn test_generic_fallback_push_string() {
609 let mut vec = ValueVector::with_type(LogicalType::Any);
610 vec.push_string("hello");
611 vec.push_string("world");
612 assert_eq!(vec.len(), 2);
613 assert_eq!(vec.get_value(0), Some(Value::String("hello".into())));
614 assert_eq!(vec.get_value(1), Some(Value::String("world".into())));
615 }
616
617 #[test]
619 fn test_generic_fallback_mixed_types() {
620 let mut vec = ValueVector::with_type(LogicalType::Any);
621 vec.push_int64(1);
622 vec.push_string("two");
623 vec.push_bool(true);
624 vec.push_float64(99.5);
625 assert_eq!(vec.len(), 4);
626 assert_eq!(vec.get_value(0), Some(Value::Int64(1)));
627 assert_eq!(vec.get_value(1), Some(Value::String("two".into())));
628 assert_eq!(vec.get_value(2), Some(Value::Bool(true)));
629 assert_eq!(vec.get_value(3), Some(Value::Float64(99.5)));
630 }
631
632 #[test]
634 fn test_type_mismatch_noop() {
635 let mut vec = ValueVector::with_type(LogicalType::Int64);
636 vec.push_string("wrong type");
637 assert_eq!(vec.len(), 0);
638
639 let mut vec = ValueVector::with_type(LogicalType::String);
640 vec.push_int64(42);
641 assert_eq!(vec.len(), 0);
642 }
643
644 #[test]
645 fn test_clear() {
646 let mut vec = ValueVector::with_type(LogicalType::Int64);
647 vec.push_int64(1);
648 vec.push_int64(2);
649
650 vec.clear();
651
652 assert!(vec.is_empty());
653 assert_eq!(vec.len(), 0);
654 }
655}