vortex_array/arrays/constant/
mod.rs

1use vortex_buffer::ByteBufferMut;
2use vortex_dtype::DType;
3use vortex_error::VortexResult;
4use vortex_mask::Mask;
5use vortex_scalar::Scalar;
6
7use crate::stats::{ArrayStats, StatsSet, StatsSetRef};
8use crate::vtable::{
9    ArrayVTable, NotSupported, OperationsVTable, VTable, ValidityVTable, VisitorVTable,
10};
11use crate::{
12    ArrayBufferVisitor, ArrayChildVisitor, ArrayRef, EncodingId, EncodingRef, IntoArray, vtable,
13};
14
15mod canonical;
16mod compute;
17mod encode;
18mod serde;
19
20vtable!(Constant);
21
22#[derive(Clone, Debug)]
23pub struct ConstantArray {
24    scalar: Scalar,
25    len: usize,
26    stats_set: ArrayStats,
27}
28
29#[derive(Clone, Debug)]
30pub struct ConstantEncoding;
31
32impl VTable for ConstantVTable {
33    type Array = ConstantArray;
34    type Encoding = ConstantEncoding;
35
36    type ArrayVTable = Self;
37    type CanonicalVTable = Self;
38    type OperationsVTable = Self;
39    type ValidityVTable = Self;
40    type VisitorVTable = Self;
41    // TODO(ngates): implement a compute kernel for elementwise operations
42    type ComputeVTable = NotSupported;
43    type EncodeVTable = Self;
44    type SerdeVTable = Self;
45
46    fn id(_encoding: &Self::Encoding) -> EncodingId {
47        EncodingId::new_ref("vortex.constant")
48    }
49
50    fn encoding(_array: &Self::Array) -> EncodingRef {
51        EncodingRef::new_ref(ConstantEncoding.as_ref())
52    }
53}
54
55impl ConstantArray {
56    pub fn new<S>(scalar: S, len: usize) -> Self
57    where
58        S: Into<Scalar>,
59    {
60        let scalar = scalar.into();
61        let stats = StatsSet::constant(scalar.clone(), len);
62        Self {
63            scalar,
64            len,
65            stats_set: ArrayStats::from(stats),
66        }
67    }
68
69    /// Returns the [`Scalar`] value of this constant array.
70    pub fn scalar(&self) -> &Scalar {
71        &self.scalar
72    }
73}
74
75impl ArrayVTable<ConstantVTable> for ConstantVTable {
76    fn len(array: &ConstantArray) -> usize {
77        array.len
78    }
79
80    fn dtype(array: &ConstantArray) -> &DType {
81        array.scalar.dtype()
82    }
83
84    fn stats(array: &ConstantArray) -> StatsSetRef<'_> {
85        array.stats_set.to_ref(array.as_ref())
86    }
87}
88
89impl OperationsVTable<ConstantVTable> for ConstantVTable {
90    fn slice(array: &ConstantArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {
91        Ok(ConstantArray::new(array.scalar.clone(), stop - start).into_array())
92    }
93
94    fn scalar_at(array: &ConstantArray, _index: usize) -> VortexResult<Scalar> {
95        Ok(array.scalar.clone())
96    }
97}
98
99impl ValidityVTable<ConstantVTable> for ConstantVTable {
100    fn is_valid(array: &ConstantArray, _index: usize) -> VortexResult<bool> {
101        Ok(!array.scalar().is_null())
102    }
103
104    fn all_valid(array: &ConstantArray) -> VortexResult<bool> {
105        Ok(!array.scalar().is_null())
106    }
107
108    fn all_invalid(array: &ConstantArray) -> VortexResult<bool> {
109        Ok(array.scalar().is_null())
110    }
111
112    fn validity_mask(array: &ConstantArray) -> VortexResult<Mask> {
113        Ok(match array.scalar().is_null() {
114            true => Mask::AllFalse(array.len()),
115            false => Mask::AllTrue(array.len()),
116        })
117    }
118}
119
120impl VisitorVTable<ConstantVTable> for ConstantVTable {
121    fn visit_buffers(array: &ConstantArray, visitor: &mut dyn ArrayBufferVisitor) {
122        let buffer = array
123            .scalar
124            .value()
125            .to_protobytes::<ByteBufferMut>()
126            .freeze();
127        visitor.visit_buffer(&buffer);
128    }
129
130    fn visit_children(_array: &ConstantArray, _visitor: &mut dyn ArrayChildVisitor) {}
131}