vortex_array/arrays/constant/
mod.rs

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