vortex_fastlanes/for/
mod.rs

1use std::fmt::Debug;
2
3pub use compress::*;
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,
10};
11use vortex_dtype::DType;
12use vortex_error::{VortexResult, vortex_bail};
13use vortex_mask::Mask;
14use vortex_scalar::Scalar;
15
16use crate::r#for::serde::ScalarValueMetadata;
17
18mod compress;
19mod compute;
20mod serde;
21
22#[derive(Clone, Debug)]
23pub struct FoRArray {
24    encoded: ArrayRef,
25    reference: Scalar,
26    stats_set: ArrayStats,
27}
28
29pub struct FoREncoding;
30impl Encoding for FoREncoding {
31    type Array = FoRArray;
32    type Metadata = ScalarValueMetadata;
33}
34
35impl FoRArray {
36    pub fn try_new(encoded: ArrayRef, reference: Scalar) -> VortexResult<Self> {
37        if reference.is_null() {
38            vortex_bail!("Reference value cannot be null");
39        }
40        let reference = reference.cast(
41            &reference
42                .dtype()
43                .with_nullability(encoded.dtype().nullability()),
44        )?;
45        Ok(Self {
46            encoded,
47            reference,
48            stats_set: Default::default(),
49        })
50    }
51
52    #[inline]
53    pub fn encoded(&self) -> &ArrayRef {
54        &self.encoded
55    }
56
57    #[inline]
58    pub fn reference_scalar(&self) -> &Scalar {
59        &self.reference
60    }
61}
62
63impl ArrayImpl for FoRArray {
64    type Encoding = FoREncoding;
65
66    fn _len(&self) -> usize {
67        self.encoded().len()
68    }
69
70    fn _dtype(&self) -> &DType {
71        self.reference_scalar().dtype()
72    }
73
74    fn _vtable(&self) -> VTableRef {
75        VTableRef::new_ref(&FoREncoding)
76    }
77
78    fn _with_children(&self, children: &[ArrayRef]) -> VortexResult<Self> {
79        Self::try_new(children[0].clone(), self.reference.clone())
80    }
81}
82
83impl ArrayCanonicalImpl for FoRArray {
84    fn _to_canonical(&self) -> VortexResult<Canonical> {
85        decompress(self).map(Canonical::Primitive)
86    }
87}
88
89impl ArrayStatisticsImpl for FoRArray {
90    fn _stats_ref(&self) -> StatsSetRef<'_> {
91        self.stats_set.to_ref(self)
92    }
93}
94
95impl ArrayValidityImpl for FoRArray {
96    fn _is_valid(&self, index: usize) -> VortexResult<bool> {
97        self.encoded().is_valid(index)
98    }
99
100    fn _all_valid(&self) -> VortexResult<bool> {
101        self.encoded().all_valid()
102    }
103
104    fn _all_invalid(&self) -> VortexResult<bool> {
105        self.encoded().all_invalid()
106    }
107
108    fn _validity_mask(&self) -> VortexResult<Mask> {
109        self.encoded().validity_mask()
110    }
111}
112
113impl ArrayVariantsImpl for FoRArray {
114    fn _as_primitive_typed(&self) -> Option<&dyn PrimitiveArrayTrait> {
115        Some(self)
116    }
117}
118
119impl PrimitiveArrayTrait for FoRArray {}
vortex_fastlanes/for/mod.rs

vortex_fastlanes/for/
mod.rs
