vortex_alp/alp/
array.rs

1use std::fmt::Debug;
2
3use vortex_array::patches::Patches;
4use vortex_array::stats::{ArrayStats, StatsSetRef};
5use vortex_array::variants::PrimitiveArrayTrait;
6use vortex_array::vtable::VTableRef;
7use vortex_array::{
8    Array, ArrayCanonicalImpl, ArrayImpl, ArrayRef, ArrayStatisticsImpl, ArrayValidityImpl,
9    ArrayVariantsImpl, Canonical, Encoding, SerdeMetadata,
10};
11use vortex_dtype::{DType, PType};
12use vortex_error::{VortexResult, vortex_bail};
13use vortex_mask::Mask;
14
15use crate::alp::serde::ALPMetadata;
16use crate::alp::{Exponents, decompress};
17
18#[derive(Clone, Debug)]
19pub struct ALPArray {
20    encoded: ArrayRef,
21    patches: Option<Patches>,
22    dtype: DType,
23    exponents: Exponents,
24    stats_set: ArrayStats,
25}
26
27pub struct ALPEncoding;
28impl Encoding for ALPEncoding {
29    type Array = ALPArray;
30    type Metadata = SerdeMetadata<ALPMetadata>;
31}
32
33impl ALPArray {
34    // TODO(ngates): remove try_new and panic on wrong DType?
35    pub fn try_new(
36        encoded: ArrayRef,
37        exponents: Exponents,
38        patches: Option<Patches>,
39    ) -> VortexResult<Self> {
40        let dtype = match encoded.dtype() {
41            DType::Primitive(PType::I32, nullability) => DType::Primitive(PType::F32, *nullability),
42            DType::Primitive(PType::I64, nullability) => DType::Primitive(PType::F64, *nullability),
43            d => vortex_bail!(MismatchedTypes: "int32 or int64", d),
44        };
45        Ok(Self {
46            dtype,
47            encoded,
48            exponents,
49            patches,
50            stats_set: Default::default(),
51        })
52    }
53
54    pub fn encoded(&self) -> &ArrayRef {
55        &self.encoded
56    }
57
58    #[inline]
59    pub fn exponents(&self) -> Exponents {
60        self.exponents
61    }
62
63    pub fn patches(&self) -> Option<&Patches> {
64        self.patches.as_ref()
65    }
66}
67
68impl ArrayImpl for ALPArray {
69    type Encoding = ALPEncoding;
70
71    fn _len(&self) -> usize {
72        self.encoded.len()
73    }
74
75    fn _dtype(&self) -> &DType {
76        &self.dtype
77    }
78
79    fn _vtable(&self) -> VTableRef {
80        VTableRef::new_ref(&ALPEncoding)
81    }
82
83    fn _with_children(&self, children: &[ArrayRef]) -> VortexResult<Self> {
84        let encoded = children[0].clone();
85
86        let patches = self.patches().map(|existing| {
87            let indices = children[1].clone();
88            let values = children[2].clone();
89            Patches::new(existing.array_len(), existing.offset(), indices, values)
90        });
91
92        ALPArray::try_new(encoded, self.exponents(), patches)
93    }
94}
95
96impl ArrayCanonicalImpl for ALPArray {
97    fn _to_canonical(&self) -> VortexResult<Canonical> {
98        decompress(self).map(Canonical::Primitive)
99    }
100}
101
102impl ArrayStatisticsImpl for ALPArray {
103    fn _stats_ref(&self) -> StatsSetRef<'_> {
104        self.stats_set.to_ref(self)
105    }
106}
107
108impl ArrayValidityImpl for ALPArray {
109    fn _is_valid(&self, index: usize) -> VortexResult<bool> {
110        self.encoded.is_valid(index)
111    }
112
113    fn _all_valid(&self) -> VortexResult<bool> {
114        self.encoded.all_valid()
115    }
116
117    fn _all_invalid(&self) -> VortexResult<bool> {
118        self.encoded.all_invalid()
119    }
120
121    fn _valid_count(&self) -> VortexResult<usize> {
122        self.encoded.valid_count()
123    }
124
125    fn _invalid_count(&self) -> VortexResult<usize> {
126        self.encoded.invalid_count()
127    }
128
129    fn _validity_mask(&self) -> VortexResult<Mask> {
130        self.encoded.validity_mask()
131    }
132}
133
134impl ArrayVariantsImpl for ALPArray {
135    fn _as_primitive_typed(&self) -> Option<&dyn PrimitiveArrayTrait> {
136        Some(self)
137    }
138}
139
140impl PrimitiveArrayTrait for ALPArray {}
vortex_alp/alp/array.rs

vortex_alp/alp/
array.rs
