akar_common/
arrow_vector.rs

1use crate::selection::SelectionVector;
2use crate::types::{PhysicalTypeID, Value};
3use crate::vector::ValueVector;
4use arrow::array::{
5    Array, ArrayRef, BooleanArray, BooleanBuilder, Float32Array, Float32Builder, Float64Array, Float64Builder,
6    Int32Array, Int32Builder, Int64Array, Int64Builder, ListArray, StringArray, StringBuilder, StructArray,
7    UInt64Array, UInt64Builder,
8};
9use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
10use arrow::datatypes::{DataType, Field};
11use std::sync::Arc;
12
13pub trait VectorAccess {
14    fn size(&self) -> usize;
15    fn physical_type(&self) -> PhysicalTypeID;
16    fn is_null(&self, row: usize) -> bool;
17
18    fn get_i64(&self, row: usize) -> Option<i64>;
19    fn get_i32(&self, row: usize) -> Option<i32>;
20    fn get_f64(&self, row: usize) -> Option<f64>;
21    fn get_f32(&self, row: usize) -> Option<f32>;
22    fn get_bool(&self, row: usize) -> Option<bool>;
23    fn get_value(&self, row: usize) -> Option<Value>;
24
25    fn get_i64_sel(&self, pos: usize, sel: &SelectionVector) -> Option<i64> {
26        if pos < sel.size {
27            self.get_i64(sel.indices[pos] as usize)
28        } else {
29            None
30        }
31    }
32    fn get_i32_sel(&self, pos: usize, sel: &SelectionVector) -> Option<i32> {
33        if pos < sel.size {
34            self.get_i32(sel.indices[pos] as usize)
35        } else {
36            None
37        }
38    }
39    fn get_f64_sel(&self, pos: usize, sel: &SelectionVector) -> Option<f64> {
40        if pos < sel.size {
41            self.get_f64(sel.indices[pos] as usize)
42        } else {
43            None
44        }
45    }
46    fn get_f32_sel(&self, pos: usize, sel: &SelectionVector) -> Option<f32> {
47        if pos < sel.size {
48            self.get_f32(sel.indices[pos] as usize)
49        } else {
50            None
51        }
52    }
53    fn get_bool_sel(&self, pos: usize, sel: &SelectionVector) -> Option<bool> {
54        if pos < sel.size {
55            self.get_bool(sel.indices[pos] as usize)
56        } else {
57            None
58        }
59    }
60}
61
62#[derive(Debug, Clone)]
63pub struct ArrowVector {
64    pub array: ArrayRef,
65    pub physical_type: PhysicalTypeID,
66}
67
68impl ArrowVector {
69    pub fn new(array: ArrayRef, physical_type: PhysicalTypeID) -> Self {
70        Self { array, physical_type }
71    }
72
73    pub fn from_legacy(vec: &ValueVector) -> Self {
74        let phys_type = vec.physical_type();
75        let size = vec.size();
76
77        // Fast path for primitive arrays
78        let build_primitive_array = |data_type: arrow::datatypes::DataType, type_size: usize| -> ArrayRef {
79            let num_bytes = size.div_ceil(8);
80            let mut null_buffer = arrow::buffer::MutableBuffer::from_len_zeroed(num_bytes);
81            let slice = null_buffer.as_slice_mut();
82            for i in 0..size {
83                if !vec.is_null(i) {
84                    arrow::util::bit_util::set_bit(slice, i);
85                }
86            }
87            let null_buffer = null_buffer.into();
88            let data_buffer = arrow::buffer::Buffer::from_slice_ref(&vec.data()[..size * type_size]);
89
90            let array_data = arrow::array::ArrayData::builder(data_type.clone())
91                .len(size)
92                .add_buffer(data_buffer)
93                .null_bit_buffer(Some(null_buffer))
94                .build()
95                .unwrap();
96            arrow::array::make_array(array_data)
97        };
98
99        let array: ArrayRef = match phys_type {
100            PhysicalTypeID::Bool => {
101                let mut builder = arrow::array::BooleanBuilder::with_capacity(size);
102                for i in 0..size {
103                    if vec.is_null(i) {
104                        builder.append_null();
105                    } else {
106                        builder.append_value(vec.get_bool(i).unwrap_or(false));
107                    }
108                }
109                Arc::new(builder.finish())
110            }
111            PhysicalTypeID::Int64 => build_primitive_array(arrow::datatypes::DataType::Int64, 8),
112            PhysicalTypeID::UInt64 => build_primitive_array(arrow::datatypes::DataType::UInt64, 8),
113            PhysicalTypeID::Int32 => build_primitive_array(arrow::datatypes::DataType::Int32, 4),
114            PhysicalTypeID::Double => build_primitive_array(arrow::datatypes::DataType::Float64, 8),
115            PhysicalTypeID::Float => build_primitive_array(arrow::datatypes::DataType::Float32, 4),
116            PhysicalTypeID::String => {
117                let mut builder = arrow::array::StringBuilder::with_capacity(size, size * 16);
118                for i in 0..size {
119                    if vec.is_null(i) {
120                        builder.append_null();
121                    } else {
122                        if let Some(val) = vec.get_value(i) {
123                            if let Value::String(s) = val {
124                                builder.append_value(&s);
125                            } else {
126                                builder.append_null();
127                            }
128                        } else {
129                            builder.append_null();
130                        }
131                    }
132                }
133                Arc::new(builder.finish())
134            }
135            PhysicalTypeID::List | PhysicalTypeID::Array => {
136                // A legacy vector has no side-storage for variable-length list
137                // payloads; emit empty lists for non-null rows (callers should
138                // prefer `arrow_array_from_values` for real list data).
139                let values: Vec<Value> = (0..size)
140                    .map(|i| {
141                        if vec.is_null(i) {
142                            Value::Null
143                        } else {
144                            vec.get_value(i).unwrap_or(Value::Null)
145                        }
146                    })
147                    .collect();
148                arrow_array_from_values(&values)
149            }
150            PhysicalTypeID::Struct => {
151                // Same limitation as List: emit the legacy placeholder value.
152                let values: Vec<Value> = (0..size)
153                    .map(|i| {
154                        if vec.is_null(i) {
155                            Value::Null
156                        } else {
157                            vec.get_value(i).unwrap_or(Value::Null)
158                        }
159                    })
160                    .collect();
161                arrow_array_from_values(&values)
162            }
163            _ => {
164                let mut builder = arrow::array::Int64Builder::with_capacity(size);
165                for _ in 0..size {
166                    builder.append_null();
167                }
168                Arc::new(builder.finish())
169            }
170        };
171        Self::new(array, phys_type)
172    }
173
174    pub fn data_type(&self) -> DataType {
175        self.array.data_type().clone()
176    }
177}
178
179impl VectorAccess for ArrowVector {
180    #[inline(always)]
181    fn size(&self) -> usize {
182        self.array.len()
183    }
184
185    #[inline(always)]
186    fn physical_type(&self) -> PhysicalTypeID {
187        self.physical_type
188    }
189
190    #[inline(always)]
191    fn is_null(&self, row: usize) -> bool {
192        if row >= self.array.len() {
193            return true;
194        }
195        self.array.is_null(row)
196    }
197
198    #[inline]
199    fn get_i64(&self, row: usize) -> Option<i64> {
200        if row >= self.array.len() {
201            return None;
202        }
203        let array = self.array.as_any().downcast_ref::<Int64Array>()?;
204        if array.is_null(row) {
205            None
206        } else {
207            Some(array.value(row))
208        }
209    }
210
211    #[inline]
212    fn get_i32(&self, row: usize) -> Option<i32> {
213        if row >= self.array.len() {
214            return None;
215        }
216        let array = self.array.as_any().downcast_ref::<Int32Array>()?;
217        if array.is_null(row) {
218            None
219        } else {
220            Some(array.value(row))
221        }
222    }
223
224    #[inline]
225    fn get_f64(&self, row: usize) -> Option<f64> {
226        if row >= self.array.len() {
227            return None;
228        }
229        let array = self.array.as_any().downcast_ref::<Float64Array>()?;
230        if array.is_null(row) {
231            None
232        } else {
233            Some(array.value(row))
234        }
235    }
236
237    #[inline]
238    fn get_f32(&self, row: usize) -> Option<f32> {
239        if row >= self.array.len() {
240            return None;
241        }
242        let array = self.array.as_any().downcast_ref::<Float32Array>()?;
243        if array.is_null(row) {
244            None
245        } else {
246            Some(array.value(row))
247        }
248    }
249
250    #[inline]
251    fn get_bool(&self, row: usize) -> Option<bool> {
252        if row >= self.array.len() {
253            return None;
254        }
255        let array = self.array.as_any().downcast_ref::<BooleanArray>()?;
256        if array.is_null(row) {
257            None
258        } else {
259            Some(array.value(row))
260        }
261    }
262
263    #[inline]
264    fn get_value(&self, row: usize) -> Option<Value> {
265        if row >= self.array.len() || self.array.is_null(row) {
266            return None;
267        }
268        match self.physical_type {
269            PhysicalTypeID::Bool => self.get_bool(row).map(Value::Bool),
270            PhysicalTypeID::Int64 => self.get_i64(row).map(Value::Int64),
271            PhysicalTypeID::UInt64 => {
272                let array = self.array.as_any().downcast_ref::<UInt64Array>()?;
273                Some(Value::UInt64(array.value(row)))
274            }
275            PhysicalTypeID::Int32 => self.get_i32(row).map(Value::Int32),
276            PhysicalTypeID::Double => self.get_f64(row).map(Value::Double),
277            PhysicalTypeID::Float => self.get_f32(row).map(Value::Float),
278            PhysicalTypeID::String => {
279                let array = self.array.as_any().downcast_ref::<StringArray>()?;
280                Some(Value::String(array.value(row).to_string()))
281            }
282            PhysicalTypeID::List => convert_arrow_scalar(&self.array, row),
283            PhysicalTypeID::Struct => convert_arrow_scalar(&self.array, row),
284            _ => {
285                // For unsupported types, return null
286                None
287            }
288        }
289    }
290}
291
292pub fn convert_arrow_scalar(array: &ArrayRef, row: usize) -> Option<Value> {
293    if array.is_null(row) {
294        return None;
295    }
296    match array.data_type() {
297        DataType::Boolean => {
298            let arr = array.as_any().downcast_ref::<BooleanArray>()?;
299            Some(Value::Bool(arr.value(row)))
300        }
301        DataType::Int64 => {
302            let arr = array.as_any().downcast_ref::<Int64Array>()?;
303            Some(Value::Int64(arr.value(row)))
304        }
305        DataType::Int32 => {
306            let arr = array.as_any().downcast_ref::<Int32Array>()?;
307            Some(Value::Int32(arr.value(row)))
308        }
309        DataType::Float64 => {
310            let arr = array.as_any().downcast_ref::<Float64Array>()?;
311            Some(Value::Double(arr.value(row)))
312        }
313        DataType::Float32 => {
314            let arr = array.as_any().downcast_ref::<Float32Array>()?;
315            Some(Value::Float(arr.value(row)))
316        }
317        DataType::Utf8 | DataType::LargeUtf8 => {
318            let arr = array.as_any().downcast_ref::<StringArray>();
319            if let Some(s_arr) = arr {
320                return Some(Value::String(s_arr.value(row).to_string()));
321            }
322            // LargeUtf8 would need LargeStringArray, but ignoring for now
323            None
324        }
325        DataType::List(_) => {
326            let arr = array.as_any().downcast_ref::<ListArray>()?;
327            let list_array = arr.value(row); // This is an ArrayRef
328            let mut list_vals = Vec::new();
329            for i in 0..list_array.len() {
330                if let Some(v) = convert_arrow_scalar(&list_array, i) {
331                    list_vals.push(v);
332                } else {
333                    list_vals.push(Value::Null);
334                }
335            }
336            Some(Value::List(list_vals))
337        }
338        DataType::Struct(fields) => {
339            let arr = array.as_any().downcast_ref::<StructArray>()?;
340            let mut entries = Vec::new();
341            for (i, field) in fields.iter().enumerate() {
342                let col = arr.column(i);
343                let val = convert_arrow_scalar(col, row).unwrap_or(Value::Null);
344                entries.push((field.name().clone(), val));
345            }
346            Some(Value::Struct(entries))
347        }
348        _ => None,
349    }
350}
351
352/// Infer the Arrow `DataType` for a list of `Value`s.
353///
354/// The first non-null value is used as the type sample. All-null inputs (or
355/// unknown types) default to `Int64`, matching the engine's null fallback.
356/// `Value::List`/`Value::Struct`/`Value::Map` are inferred recursively so that
357/// nested complex types produce the right `List`/`Struct` Arrow types.
358pub fn infer_arrow_type(values: &[Value]) -> DataType {
359    let sample = values.iter().find(|v| !matches!(v, Value::Null));
360    match sample {
361        Some(Value::Bool(_)) => DataType::Boolean,
362        Some(Value::Int64(_))
363        | Some(Value::Int128(_))
364        | Some(Value::UInt128(_))
365        | Some(Value::Date(_))
366        | Some(Value::Timestamp(_))
367        | Some(Value::TimestampTz(_))
368        | Some(Value::TimestampNs(_))
369        | Some(Value::TimestampMs(_))
370        | Some(Value::TimestampSec(_))
371        | Some(Value::DTime(_)) => DataType::Int64,
372        Some(Value::UInt64(_)) => DataType::UInt64,
373        Some(Value::Int32(_)) | Some(Value::UInt32(_)) => DataType::Int32,
374        Some(Value::Double(_)) => DataType::Float64,
375        Some(Value::Float(_)) => DataType::Float32,
376        Some(Value::String(_)) | Some(Value::Json(_)) | Some(Value::Interval(_)) => DataType::Utf8,
377        Some(Value::List(inner)) => DataType::List(Arc::new(Field::new("item", infer_arrow_type(inner), true))),
378        Some(Value::Struct(entries)) => {
379            let fields = entries
380                .iter()
381                .map(|(k, v)| Field::new(k.clone(), infer_arrow_type(std::slice::from_ref(v)), true))
382                .collect();
383            DataType::Struct(fields)
384        }
385        Some(Value::Map(entries)) => {
386            let fields = entries
387                .iter()
388                .map(|(k, v)| {
389                    let name = match k {
390                        Value::String(s) => s.clone(),
391                        other => format!("{other:?}"),
392                    };
393                    Field::new(name, infer_arrow_type(std::slice::from_ref(v)), true)
394                })
395                .collect();
396            DataType::Struct(fields)
397        }
398        Some(Value::Union(_, v)) => infer_arrow_type(std::slice::from_ref(v.as_ref())),
399        Some(Value::InternalID(_)) => DataType::Int64,
400        Some(Value::Blob(_)) => DataType::Binary,
401        _ => DataType::Int64,
402    }
403}
404
405/// Build an Arrow `ArrayRef` from a `Vec<Value>`, recursively handling
406/// primitives, `List`, `Struct` and `Map` values.
407///
408/// This is the inverse of [`convert_arrow_scalar`] and is used by the scan /
409/// projection paths to emit complex-typed columns (e.g. `FLOAT[]`, `{..}`
410/// literals) without going through a `ValueVector` (which has no side-storage
411/// for variable-length List/Struct payloads).
412pub fn arrow_array_from_values(values: &[Value]) -> ArrayRef {
413    let data_type = infer_arrow_type(values);
414    match &data_type {
415        DataType::Boolean => {
416            let mut builder = BooleanBuilder::with_capacity(values.len());
417            for v in values {
418                match v {
419                    Value::Null => builder.append_null(),
420                    Value::Bool(b) => builder.append_value(*b),
421                    _ => builder.append_null(),
422                }
423            }
424            Arc::new(builder.finish())
425        }
426        DataType::Int64 => {
427            let mut builder = Int64Builder::with_capacity(values.len());
428            for v in values {
429                match v {
430                    Value::Null => builder.append_null(),
431                    Value::Int64(n) => builder.append_value(*n),
432                    Value::Int128(n) => builder.append_value(*n as i64),
433                    Value::UInt128(n) => builder.append_value(*n as i64),
434                    Value::Int32(n) => builder.append_value(*n as i64),
435                    Value::UInt64(n) => builder.append_value(*n as i64),
436                    Value::Date(n) => builder.append_value(n.0 as i64),
437                    Value::Timestamp(n) | Value::TimestampNs(n) | Value::TimestampMs(n) | Value::TimestampSec(n) => {
438                        builder.append_value(n.0)
439                    }
440                    Value::TimestampTz(n) => builder.append_value(n.0),
441                    Value::DTime(n) => builder.append_value(*n),
442                    Value::InternalID(id) => builder.append_value(id.offset as i64),
443                    _ => builder.append_null(),
444                }
445            }
446            Arc::new(builder.finish())
447        }
448        DataType::UInt64 => {
449            let mut builder = UInt64Builder::with_capacity(values.len());
450            for v in values {
451                match v {
452                    Value::Null => builder.append_null(),
453                    Value::UInt64(n) => builder.append_value(*n),
454                    Value::Int64(n) if *n >= 0 => builder.append_value(*n as u64),
455                    _ => builder.append_null(),
456                }
457            }
458            Arc::new(builder.finish())
459        }
460        DataType::Int32 => {
461            let mut builder = Int32Builder::with_capacity(values.len());
462            for v in values {
463                match v {
464                    Value::Null => builder.append_null(),
465                    Value::Int32(n) => builder.append_value(*n),
466                    Value::Int64(n) => builder.append_value(*n as i32),
467                    _ => builder.append_null(),
468                }
469            }
470            Arc::new(builder.finish())
471        }
472        DataType::Float64 => {
473            let mut builder = Float64Builder::with_capacity(values.len());
474            for v in values {
475                match v {
476                    Value::Null => builder.append_null(),
477                    Value::Double(n) => builder.append_value(*n),
478                    Value::Float(n) => builder.append_value(*n as f64),
479                    Value::Int64(n) => builder.append_value(*n as f64),
480                    _ => builder.append_null(),
481                }
482            }
483            Arc::new(builder.finish())
484        }
485        DataType::Float32 => {
486            let mut builder = Float32Builder::with_capacity(values.len());
487            for v in values {
488                match v {
489                    Value::Null => builder.append_null(),
490                    Value::Float(n) => builder.append_value(*n),
491                    Value::Double(n) => builder.append_value(*n as f32),
492                    _ => builder.append_null(),
493                }
494            }
495            Arc::new(builder.finish())
496        }
497        DataType::Utf8 => {
498            let mut builder = StringBuilder::with_capacity(values.len(), values.len() * 16);
499            for v in values {
500                match v {
501                    Value::Null => builder.append_null(),
502                    Value::String(s) => builder.append_value(s),
503                    Value::Interval(i) => {
504                        builder.append_value(format!("{} months {} days {} microseconds", i.months, i.days, i.micros))
505                    }
506                    Value::Json(j) => builder.append_value(j.to_string()),
507                    _ => builder.append_null(),
508                }
509            }
510            Arc::new(builder.finish())
511        }
512        DataType::List(_) => {
513            let mut child_values: Vec<Value> = Vec::new();
514            let mut offsets: Vec<i32> = Vec::with_capacity(values.len() + 1);
515            let mut valid: Vec<bool> = Vec::with_capacity(values.len());
516            offsets.push(0);
517            for v in values {
518                match v {
519                    Value::List(inner) => {
520                        child_values.extend(inner.iter().cloned());
521                        offsets.push(child_values.len() as i32);
522                        valid.push(true);
523                    }
524                    _ => {
525                        offsets.push(child_values.len() as i32);
526                        valid.push(false);
527                    }
528                }
529            }
530            let child = arrow_array_from_values(&child_values);
531            let field = Arc::new(Field::new("item", child.data_type().clone(), true));
532            let offsets_buf = OffsetBuffer::new(ScalarBuffer::from(offsets));
533            let nulls = NullBuffer::from(valid);
534            match ListArray::try_new(field, offsets_buf, child, Some(nulls)) {
535                Ok(a) => Arc::new(a),
536                Err(_) => {
537                    let mut builder = Int64Builder::with_capacity(values.len());
538                    builder.append_nulls(values.len());
539                    Arc::new(builder.finish())
540                }
541            }
542        }
543        DataType::Struct(fields) => {
544            let mut columns: Vec<ArrayRef> = Vec::with_capacity(fields.len());
545            for field in fields.iter() {
546                let col: Vec<Value> = values
547                    .iter()
548                    .map(|v| match v {
549                        Value::Struct(entries) => entries
550                            .iter()
551                            .find(|(k, _)| k == field.name())
552                            .map(|(_, val)| val.clone())
553                            .unwrap_or(Value::Null),
554                        Value::Map(entries) => entries
555                            .iter()
556                            .find(|(k, _)| matches!(k, Value::String(s) if s == field.name()))
557                            .map(|(_, val)| val.clone())
558                            .unwrap_or(Value::Null),
559                        _ => Value::Null,
560                    })
561                    .collect();
562                columns.push(arrow_array_from_values(&col));
563            }
564            let valid: Vec<bool> = values
565                .iter()
566                .map(|v| matches!(v, Value::Struct(_) | Value::Map(_)))
567                .collect();
568            let nulls = NullBuffer::from(valid);
569            match StructArray::try_new(fields.clone(), columns, Some(nulls)) {
570                Ok(a) => Arc::new(a),
571                Err(_) => {
572                    let mut builder = Int64Builder::with_capacity(values.len());
573                    builder.append_nulls(values.len());
574                    Arc::new(builder.finish())
575                }
576            }
577        }
578        DataType::Binary => {
579            let mut builder = arrow::array::BinaryBuilder::with_capacity(values.len(), values.len() * 8);
580            for v in values {
581                match v {
582                    Value::Null => builder.append_null(),
583                    Value::Blob(b) => builder.append_value(b),
584                    _ => builder.append_null(),
585                }
586            }
587            Arc::new(builder.finish())
588        }
589        _ => {
590            let mut builder = Int64Builder::with_capacity(values.len());
591            builder.append_nulls(values.len());
592            Arc::new(builder.finish())
593        }
594    }
595}
596
597#[derive(Debug, Clone)]
598pub enum Vector {
599    Arrow(ArrowVector),
600    Legacy(ValueVector),
601}
602
603impl Vector {
604    #[inline]
605    pub fn as_arrow(&self) -> Option<&ArrowVector> {
606        match self {
607            Vector::Arrow(a) => Some(a),
608            _ => None,
609        }
610    }
611
612    #[inline]
613    pub fn as_legacy(&self) -> Option<&ValueVector> {
614        match self {
615            Vector::Legacy(l) => Some(l),
616            _ => None,
617        }
618    }
619
620    #[inline]
621    pub fn as_legacy_mut(&mut self) -> Option<&mut ValueVector> {
622        match self {
623            Vector::Legacy(l) => Some(l),
624            _ => None,
625        }
626    }
627}
628
629impl VectorAccess for Vector {
630    fn size(&self) -> usize {
631        match self {
632            Vector::Arrow(a) => a.size(),
633            Vector::Legacy(l) => l.size(),
634        }
635    }
636
637    fn physical_type(&self) -> PhysicalTypeID {
638        match self {
639            Vector::Arrow(a) => a.physical_type(),
640            Vector::Legacy(l) => l.physical_type(),
641        }
642    }
643
644    fn is_null(&self, row: usize) -> bool {
645        match self {
646            Vector::Arrow(a) => a.is_null(row),
647            Vector::Legacy(l) => l.is_null(row),
648        }
649    }
650
651    fn get_i64(&self, row: usize) -> Option<i64> {
652        match self {
653            Vector::Arrow(a) => a.get_i64(row),
654            Vector::Legacy(l) => l.get_i64(row),
655        }
656    }
657
658    fn get_i32(&self, row: usize) -> Option<i32> {
659        match self {
660            Vector::Arrow(a) => a.get_i32(row),
661            Vector::Legacy(l) => l.get_i32(row),
662        }
663    }
664
665    fn get_f64(&self, row: usize) -> Option<f64> {
666        match self {
667            Vector::Arrow(a) => a.get_f64(row),
668            Vector::Legacy(l) => l.get_double(row),
669        }
670    }
671
672    fn get_f32(&self, row: usize) -> Option<f32> {
673        match self {
674            Vector::Arrow(a) => a.get_f32(row),
675            Vector::Legacy(l) => {
676                if l.is_null(row) {
677                    return None;
678                }
679                let v = l.get_value(row)?;
680                match v {
681                    Value::Float(f) => Some(f),
682                    Value::Double(d) => Some(d as f32),
683                    _ => None,
684                }
685            }
686        }
687    }
688
689    fn get_bool(&self, row: usize) -> Option<bool> {
690        match self {
691            Vector::Arrow(a) => a.get_bool(row),
692            Vector::Legacy(l) => l.get_bool(row),
693        }
694    }
695
696    fn get_value(&self, row: usize) -> Option<Value> {
697        match self {
698            Vector::Arrow(a) => a.get_value(row),
699            Vector::Legacy(l) => l.get_value(row),
700        }
701    }
702}
703
704impl From<ValueVector> for Vector {
705    fn from(v: ValueVector) -> Self {
706        Vector::Legacy(v)
707    }
708}
709
710impl From<ArrowVector> for Vector {
711    fn from(a: ArrowVector) -> Self {
712        Vector::Arrow(a)
713    }
714}
715
716#[cfg(test)]
717mod tests {
718    use super::*;
719    use crate::types::PhysicalTypeID;
720    use crate::vector::ValueVector;
721
722    #[test]
723    fn test_arrow_vector_from_legacy_i64() {
724        let mut legacy = ValueVector::new(PhysicalTypeID::Int64, 5);
725        legacy.set_i64(0, 10);
726        legacy.set_i64(1, 20);
727        legacy.set_i64(2, 30);
728        legacy.set_null(3, true);
729        legacy.set_i64(4, 50);
730        legacy.resize(5);
731
732        let arrow = ArrowVector::from_legacy(&legacy);
733        assert_eq!(arrow.size(), 5);
734        assert_eq!(arrow.get_i64(0), Some(10));
735        assert_eq!(arrow.get_i64(1), Some(20));
736        assert_eq!(arrow.get_i64(2), Some(30));
737        assert_eq!(arrow.get_i64(3), None);
738        assert_eq!(arrow.get_i64(4), Some(50));
739    }
740
741    #[test]
742    fn test_arrow_vector_bool() {
743        let mut legacy = ValueVector::new(PhysicalTypeID::Bool, 4);
744        legacy.push_bool(true);
745        legacy.push_bool(false);
746        legacy.push_bool(true);
747        legacy.set_null(3, true);
748        legacy.resize(4);
749
750        let arrow = ArrowVector::from_legacy(&legacy);
751        assert_eq!(arrow.size(), 4);
752        assert_eq!(arrow.get_bool(0), Some(true));
753        assert_eq!(arrow.get_bool(1), Some(false));
754        assert_eq!(arrow.get_bool(2), Some(true));
755        assert_eq!(arrow.get_bool(3), None);
756    }
757
758    #[test]
759    fn test_vector_enum_dispatch() {
760        let mut legacy = ValueVector::new(PhysicalTypeID::Int64, 3);
761        legacy.set_i64(0, 100);
762        legacy.set_i64(1, 200);
763        legacy.resize(2);
764
765        let vec = Vector::Legacy(legacy);
766        assert_eq!(vec.get_i64(0), Some(100));
767        assert_eq!(vec.get_i64(1), Some(200));
768        assert_eq!(vec.size(), 2);
769    }
770
771    #[test]
772    fn test_selection_vector_access() {
773        let mut legacy = ValueVector::new(PhysicalTypeID::Int64, 5);
774        legacy.set_i64(0, 10);
775        legacy.set_i64(1, 20);
776        legacy.set_i64(2, 30);
777        legacy.set_i64(3, 40);
778        legacy.set_i64(4, 50);
779        legacy.resize(5);
780
781        let arrow = ArrowVector::from_legacy(&legacy);
782        let sel = SelectionVector::from_slice(&[0, 2, 4]);
783
784        assert_eq!(arrow.get_i64_sel(0, &sel), Some(10));
785        assert_eq!(arrow.get_i64_sel(1, &sel), Some(30));
786        assert_eq!(arrow.get_i64_sel(2, &sel), Some(50));
787        assert_eq!(arrow.get_i64_sel(3, &sel), None);
788    }
789
790    #[test]
791    fn test_arrow_array_from_list_values() {
792        let values = vec![
793            Value::List(vec![Value::Float(0.1), Value::Float(0.2)]),
794            Value::List(vec![Value::Float(0.3)]),
795            Value::Null,
796        ];
797        let arr = arrow_array_from_values(&values);
798        assert_eq!(
799            arr.data_type(),
800            &DataType::List(Arc::new(Field::new("item", DataType::Float32, true)))
801        );
802        assert_eq!(arr.len(), 3);
803        assert!(!arr.is_null(0));
804        assert!(!arr.is_null(1));
805        assert!(arr.is_null(2));
806        let list_arr = arr.as_any().downcast_ref::<ListArray>().unwrap();
807        let v0 = convert_arrow_scalar(&arr, 0).unwrap();
808        assert_eq!(v0, Value::List(vec![Value::Float(0.1), Value::Float(0.2)]));
809        let v1 = convert_arrow_scalar(&arr, 1).unwrap();
810        assert_eq!(v1, Value::List(vec![Value::Float(0.3)]));
811        assert_eq!(list_arr.value(0).len(), 2);
812    }
813
814    #[test]
815    fn test_arrow_array_from_struct_values() {
816        let values = vec![
817            Value::Map(vec![(Value::String("id".into()), Value::Int64(1))]),
818            Value::Null,
819        ];
820        let arr = arrow_array_from_values(&values);
821        assert!(matches!(arr.data_type(), DataType::Struct(_)));
822        assert_eq!(arr.len(), 2);
823        assert!(!arr.is_null(0));
824        assert!(arr.is_null(1));
825        let v0 = convert_arrow_scalar(&arr, 0).unwrap();
826        assert_eq!(v0, Value::Struct(vec![("id".to_string(), Value::Int64(1))]));
827    }
828
829    #[test]
830    fn test_arrow_array_from_empty_list() {
831        let values = vec![Value::List(vec![])];
832        let arr = arrow_array_from_values(&values);
833        assert!(matches!(arr.data_type(), DataType::List(_)));
834        assert_eq!(convert_arrow_scalar(&arr, 0).unwrap(), Value::List(vec![]));
835    }
836
837    #[test]
838    fn test_arrow_array_nested_list() {
839        let values = vec![Value::List(vec![Value::List(vec![Value::Int64(1), Value::Int64(2)])])];
840        let arr = arrow_array_from_values(&values);
841        assert!(matches!(arr.data_type(), DataType::List(_)));
842        let v0 = convert_arrow_scalar(&arr, 0).unwrap();
843        assert_eq!(
844            v0,
845            Value::List(vec![Value::List(vec![Value::Int64(1), Value::Int64(2)])])
846        );
847    }
848}
akar_common/arrow_vector.rs

akar_common/
arrow_vector.rs
