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, Float32Array, Float64Array, Int32Array, Int64Array, ListArray, StringArray,
6    StructArray, UInt64Array,
7};
8use arrow::datatypes::DataType;
9use std::sync::Arc;
10
11pub trait VectorAccess {
12    fn size(&self) -> usize;
13    fn physical_type(&self) -> PhysicalTypeID;
14    fn is_null(&self, row: usize) -> bool;
15
16    fn get_i64(&self, row: usize) -> Option<i64>;
17    fn get_i32(&self, row: usize) -> Option<i32>;
18    fn get_f64(&self, row: usize) -> Option<f64>;
19    fn get_f32(&self, row: usize) -> Option<f32>;
20    fn get_bool(&self, row: usize) -> Option<bool>;
21    fn get_value(&self, row: usize) -> Option<Value>;
22
23    fn get_i64_sel(&self, pos: usize, sel: &SelectionVector) -> Option<i64> {
24        if pos < sel.size {
25            self.get_i64(sel.indices[pos] as usize)
26        } else {
27            None
28        }
29    }
30    fn get_i32_sel(&self, pos: usize, sel: &SelectionVector) -> Option<i32> {
31        if pos < sel.size {
32            self.get_i32(sel.indices[pos] as usize)
33        } else {
34            None
35        }
36    }
37    fn get_f64_sel(&self, pos: usize, sel: &SelectionVector) -> Option<f64> {
38        if pos < sel.size {
39            self.get_f64(sel.indices[pos] as usize)
40        } else {
41            None
42        }
43    }
44    fn get_f32_sel(&self, pos: usize, sel: &SelectionVector) -> Option<f32> {
45        if pos < sel.size {
46            self.get_f32(sel.indices[pos] as usize)
47        } else {
48            None
49        }
50    }
51    fn get_bool_sel(&self, pos: usize, sel: &SelectionVector) -> Option<bool> {
52        if pos < sel.size {
53            self.get_bool(sel.indices[pos] as usize)
54        } else {
55            None
56        }
57    }
58}
59
60#[derive(Debug, Clone)]
61pub struct ArrowVector {
62    pub array: ArrayRef,
63    pub physical_type: PhysicalTypeID,
64}
65
66impl ArrowVector {
67    pub fn new(array: ArrayRef, physical_type: PhysicalTypeID) -> Self {
68        Self { array, physical_type }
69    }
70
71    pub fn from_legacy(vec: &ValueVector) -> Self {
72        let phys_type = vec.physical_type();
73        let size = vec.size();
74
75        // Fast path for primitive arrays
76        let build_primitive_array = |data_type: arrow::datatypes::DataType, type_size: usize| -> ArrayRef {
77            let num_bytes = size.div_ceil(8);
78            let mut null_buffer = arrow::buffer::MutableBuffer::from_len_zeroed(num_bytes);
79            let slice = null_buffer.as_slice_mut();
80            for i in 0..size {
81                if !vec.is_null(i) {
82                    arrow::util::bit_util::set_bit(slice, i);
83                }
84            }
85            let null_buffer = null_buffer.into();
86            let data_buffer = arrow::buffer::Buffer::from_slice_ref(&vec.data()[..size * type_size]);
87
88            let array_data = arrow::array::ArrayData::builder(data_type.clone())
89                .len(size)
90                .add_buffer(data_buffer)
91                .null_bit_buffer(Some(null_buffer))
92                .build()
93                .unwrap();
94            arrow::array::make_array(array_data)
95        };
96
97        let array: ArrayRef = match phys_type {
98            PhysicalTypeID::Bool => {
99                let mut builder = arrow::array::BooleanBuilder::with_capacity(size);
100                for i in 0..size {
101                    if vec.is_null(i) {
102                        builder.append_null();
103                    } else {
104                        builder.append_value(vec.get_bool(i).unwrap_or(false));
105                    }
106                }
107                Arc::new(builder.finish())
108            }
109            PhysicalTypeID::Int64 => build_primitive_array(arrow::datatypes::DataType::Int64, 8),
110            PhysicalTypeID::UInt64 => build_primitive_array(arrow::datatypes::DataType::UInt64, 8),
111            PhysicalTypeID::Int32 => build_primitive_array(arrow::datatypes::DataType::Int32, 4),
112            PhysicalTypeID::Double => build_primitive_array(arrow::datatypes::DataType::Float64, 8),
113            PhysicalTypeID::Float => build_primitive_array(arrow::datatypes::DataType::Float32, 4),
114            PhysicalTypeID::String => {
115                let mut builder = arrow::array::StringBuilder::with_capacity(size, size * 16);
116                for i in 0..size {
117                    if vec.is_null(i) {
118                        builder.append_null();
119                    } else {
120                        if let Some(val) = vec.get_value(i) {
121                            if let Value::String(s) = val {
122                                builder.append_value(&s);
123                            } else {
124                                builder.append_null();
125                            }
126                        } else {
127                            builder.append_null();
128                        }
129                    }
130                }
131                Arc::new(builder.finish())
132            }
133            _ => {
134                let mut builder = arrow::array::Int64Builder::with_capacity(size);
135                for _ in 0..size {
136                    builder.append_null();
137                }
138                Arc::new(builder.finish())
139            }
140        };
141        Self::new(array, phys_type)
142    }
143
144    pub fn data_type(&self) -> DataType {
145        self.array.data_type().clone()
146    }
147}
148
149impl VectorAccess for ArrowVector {
150    #[inline(always)]
151    fn size(&self) -> usize {
152        self.array.len()
153    }
154
155    #[inline(always)]
156    fn physical_type(&self) -> PhysicalTypeID {
157        self.physical_type
158    }
159
160    #[inline(always)]
161    fn is_null(&self, row: usize) -> bool {
162        if row >= self.array.len() {
163            return true;
164        }
165        self.array.is_null(row)
166    }
167
168    #[inline]
169    fn get_i64(&self, row: usize) -> Option<i64> {
170        if row >= self.array.len() {
171            return None;
172        }
173        let array = self.array.as_any().downcast_ref::<Int64Array>()?;
174        if array.is_null(row) {
175            None
176        } else {
177            Some(array.value(row))
178        }
179    }
180
181    #[inline]
182    fn get_i32(&self, row: usize) -> Option<i32> {
183        if row >= self.array.len() {
184            return None;
185        }
186        let array = self.array.as_any().downcast_ref::<Int32Array>()?;
187        if array.is_null(row) {
188            None
189        } else {
190            Some(array.value(row))
191        }
192    }
193
194    #[inline]
195    fn get_f64(&self, row: usize) -> Option<f64> {
196        if row >= self.array.len() {
197            return None;
198        }
199        let array = self.array.as_any().downcast_ref::<Float64Array>()?;
200        if array.is_null(row) {
201            None
202        } else {
203            Some(array.value(row))
204        }
205    }
206
207    #[inline]
208    fn get_f32(&self, row: usize) -> Option<f32> {
209        if row >= self.array.len() {
210            return None;
211        }
212        let array = self.array.as_any().downcast_ref::<Float32Array>()?;
213        if array.is_null(row) {
214            None
215        } else {
216            Some(array.value(row))
217        }
218    }
219
220    #[inline]
221    fn get_bool(&self, row: usize) -> Option<bool> {
222        if row >= self.array.len() {
223            return None;
224        }
225        let array = self.array.as_any().downcast_ref::<BooleanArray>()?;
226        if array.is_null(row) {
227            None
228        } else {
229            Some(array.value(row))
230        }
231    }
232
233    #[inline]
234    fn get_value(&self, row: usize) -> Option<Value> {
235        if row >= self.array.len() || self.array.is_null(row) {
236            return None;
237        }
238        match self.physical_type {
239            PhysicalTypeID::Bool => self.get_bool(row).map(Value::Bool),
240            PhysicalTypeID::Int64 => self.get_i64(row).map(Value::Int64),
241            PhysicalTypeID::UInt64 => {
242                let array = self.array.as_any().downcast_ref::<UInt64Array>()?;
243                Some(Value::UInt64(array.value(row)))
244            }
245            PhysicalTypeID::Int32 => self.get_i32(row).map(Value::Int32),
246            PhysicalTypeID::Double => self.get_f64(row).map(Value::Double),
247            PhysicalTypeID::Float => self.get_f32(row).map(Value::Float),
248            PhysicalTypeID::String => {
249                let array = self.array.as_any().downcast_ref::<StringArray>()?;
250                Some(Value::String(array.value(row).to_string()))
251            }
252            PhysicalTypeID::List => convert_arrow_scalar(&self.array, row),
253            PhysicalTypeID::Struct => convert_arrow_scalar(&self.array, row),
254            _ => {
255                // For unsupported types, return null
256                None
257            }
258        }
259    }
260}
261
262pub fn convert_arrow_scalar(array: &ArrayRef, row: usize) -> Option<Value> {
263    if array.is_null(row) {
264        return None;
265    }
266    match array.data_type() {
267        DataType::Boolean => {
268            let arr = array.as_any().downcast_ref::<BooleanArray>()?;
269            Some(Value::Bool(arr.value(row)))
270        }
271        DataType::Int64 => {
272            let arr = array.as_any().downcast_ref::<Int64Array>()?;
273            Some(Value::Int64(arr.value(row)))
274        }
275        DataType::Int32 => {
276            let arr = array.as_any().downcast_ref::<Int32Array>()?;
277            Some(Value::Int32(arr.value(row)))
278        }
279        DataType::Float64 => {
280            let arr = array.as_any().downcast_ref::<Float64Array>()?;
281            Some(Value::Double(arr.value(row)))
282        }
283        DataType::Float32 => {
284            let arr = array.as_any().downcast_ref::<Float32Array>()?;
285            Some(Value::Float(arr.value(row)))
286        }
287        DataType::Utf8 | DataType::LargeUtf8 => {
288            let arr = array.as_any().downcast_ref::<StringArray>();
289            if let Some(s_arr) = arr {
290                return Some(Value::String(s_arr.value(row).to_string()));
291            }
292            // LargeUtf8 would need LargeStringArray, but ignoring for now
293            None
294        }
295        DataType::List(_) => {
296            let arr = array.as_any().downcast_ref::<ListArray>()?;
297            let list_array = arr.value(row); // This is an ArrayRef
298            let mut list_vals = Vec::new();
299            for i in 0..list_array.len() {
300                if let Some(v) = convert_arrow_scalar(&list_array, i) {
301                    list_vals.push(v);
302                } else {
303                    list_vals.push(Value::Null);
304                }
305            }
306            Some(Value::List(list_vals))
307        }
308        DataType::Struct(fields) => {
309            let arr = array.as_any().downcast_ref::<StructArray>()?;
310            let mut entries = Vec::new();
311            for (i, field) in fields.iter().enumerate() {
312                let col = arr.column(i);
313                let val = convert_arrow_scalar(col, row).unwrap_or(Value::Null);
314                entries.push((field.name().clone(), val));
315            }
316            Some(Value::Struct(entries))
317        }
318        _ => None,
319    }
320}
321
322#[derive(Debug, Clone)]
323pub enum Vector {
324    Arrow(ArrowVector),
325    Legacy(ValueVector),
326}
327
328impl Vector {
329    #[inline]
330    pub fn as_arrow(&self) -> Option<&ArrowVector> {
331        match self {
332            Vector::Arrow(a) => Some(a),
333            _ => None,
334        }
335    }
336
337    #[inline]
338    pub fn as_legacy(&self) -> Option<&ValueVector> {
339        match self {
340            Vector::Legacy(l) => Some(l),
341            _ => None,
342        }
343    }
344
345    #[inline]
346    pub fn as_legacy_mut(&mut self) -> Option<&mut ValueVector> {
347        match self {
348            Vector::Legacy(l) => Some(l),
349            _ => None,
350        }
351    }
352}
353
354impl VectorAccess for Vector {
355    fn size(&self) -> usize {
356        match self {
357            Vector::Arrow(a) => a.size(),
358            Vector::Legacy(l) => l.size(),
359        }
360    }
361
362    fn physical_type(&self) -> PhysicalTypeID {
363        match self {
364            Vector::Arrow(a) => a.physical_type(),
365            Vector::Legacy(l) => l.physical_type(),
366        }
367    }
368
369    fn is_null(&self, row: usize) -> bool {
370        match self {
371            Vector::Arrow(a) => a.is_null(row),
372            Vector::Legacy(l) => l.is_null(row),
373        }
374    }
375
376    fn get_i64(&self, row: usize) -> Option<i64> {
377        match self {
378            Vector::Arrow(a) => a.get_i64(row),
379            Vector::Legacy(l) => l.get_i64(row),
380        }
381    }
382
383    fn get_i32(&self, row: usize) -> Option<i32> {
384        match self {
385            Vector::Arrow(a) => a.get_i32(row),
386            Vector::Legacy(l) => l.get_i32(row),
387        }
388    }
389
390    fn get_f64(&self, row: usize) -> Option<f64> {
391        match self {
392            Vector::Arrow(a) => a.get_f64(row),
393            Vector::Legacy(l) => l.get_double(row),
394        }
395    }
396
397    fn get_f32(&self, row: usize) -> Option<f32> {
398        match self {
399            Vector::Arrow(a) => a.get_f32(row),
400            Vector::Legacy(l) => {
401                if l.is_null(row) {
402                    return None;
403                }
404                let v = l.get_value(row)?;
405                match v {
406                    Value::Float(f) => Some(f),
407                    Value::Double(d) => Some(d as f32),
408                    _ => None,
409                }
410            }
411        }
412    }
413
414    fn get_bool(&self, row: usize) -> Option<bool> {
415        match self {
416            Vector::Arrow(a) => a.get_bool(row),
417            Vector::Legacy(l) => l.get_bool(row),
418        }
419    }
420
421    fn get_value(&self, row: usize) -> Option<Value> {
422        match self {
423            Vector::Arrow(a) => a.get_value(row),
424            Vector::Legacy(l) => l.get_value(row),
425        }
426    }
427}
428
429impl From<ValueVector> for Vector {
430    fn from(v: ValueVector) -> Self {
431        Vector::Legacy(v)
432    }
433}
434
435impl From<ArrowVector> for Vector {
436    fn from(a: ArrowVector) -> Self {
437        Vector::Arrow(a)
438    }
439}
440
441#[cfg(test)]
442mod tests {
443    use super::*;
444    use crate::types::PhysicalTypeID;
445    use crate::vector::ValueVector;
446
447    #[test]
448    fn test_arrow_vector_from_legacy_i64() {
449        let mut legacy = ValueVector::new(PhysicalTypeID::Int64, 5);
450        legacy.set_i64(0, 10);
451        legacy.set_i64(1, 20);
452        legacy.set_i64(2, 30);
453        legacy.set_null(3, true);
454        legacy.set_i64(4, 50);
455        legacy.resize(5);
456
457        let arrow = ArrowVector::from_legacy(&legacy);
458        assert_eq!(arrow.size(), 5);
459        assert_eq!(arrow.get_i64(0), Some(10));
460        assert_eq!(arrow.get_i64(1), Some(20));
461        assert_eq!(arrow.get_i64(2), Some(30));
462        assert_eq!(arrow.get_i64(3), None);
463        assert_eq!(arrow.get_i64(4), Some(50));
464    }
465
466    #[test]
467    fn test_arrow_vector_bool() {
468        let mut legacy = ValueVector::new(PhysicalTypeID::Bool, 4);
469        legacy.push_bool(true);
470        legacy.push_bool(false);
471        legacy.push_bool(true);
472        legacy.set_null(3, true);
473        legacy.resize(4);
474
475        let arrow = ArrowVector::from_legacy(&legacy);
476        assert_eq!(arrow.size(), 4);
477        assert_eq!(arrow.get_bool(0), Some(true));
478        assert_eq!(arrow.get_bool(1), Some(false));
479        assert_eq!(arrow.get_bool(2), Some(true));
480        assert_eq!(arrow.get_bool(3), None);
481    }
482
483    #[test]
484    fn test_vector_enum_dispatch() {
485        let mut legacy = ValueVector::new(PhysicalTypeID::Int64, 3);
486        legacy.set_i64(0, 100);
487        legacy.set_i64(1, 200);
488        legacy.resize(2);
489
490        let vec = Vector::Legacy(legacy);
491        assert_eq!(vec.get_i64(0), Some(100));
492        assert_eq!(vec.get_i64(1), Some(200));
493        assert_eq!(vec.size(), 2);
494    }
495
496    #[test]
497    fn test_selection_vector_access() {
498        let mut legacy = ValueVector::new(PhysicalTypeID::Int64, 5);
499        legacy.set_i64(0, 10);
500        legacy.set_i64(1, 20);
501        legacy.set_i64(2, 30);
502        legacy.set_i64(3, 40);
503        legacy.set_i64(4, 50);
504        legacy.resize(5);
505
506        let arrow = ArrowVector::from_legacy(&legacy);
507        let sel = SelectionVector::from_slice(&[0, 2, 4]);
508
509        assert_eq!(arrow.get_i64_sel(0, &sel), Some(10));
510        assert_eq!(arrow.get_i64_sel(1, &sel), Some(30));
511        assert_eq!(arrow.get_i64_sel(2, &sel), Some(50));
512        assert_eq!(arrow.get_i64_sel(3, &sel), None);
513    }
514}
akar_common/arrow_vector.rs

akar_common/
arrow_vector.rs
