Skip to main content

vortex_array/arrays/constant/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Debug;
5use std::hash::Hash;
6
7use vortex_buffer::ByteBufferMut;
8use vortex_error::VortexResult;
9use vortex_error::vortex_ensure;
10use vortex_error::vortex_panic;
11use vortex_session::VortexSession;
12
13use crate::ArrayRef;
14use crate::ExecutionCtx;
15use crate::IntoArray;
16use crate::Precision;
17use crate::arrays::ConstantArray;
18use crate::arrays::constant::compute::rules::PARENT_RULES;
19use crate::arrays::constant::vtable::canonical::constant_canonicalize;
20use crate::buffer::BufferHandle;
21use crate::dtype::DType;
22use crate::scalar::Scalar;
23use crate::scalar::ScalarValue;
24use crate::serde::ArrayChildren;
25use crate::stats::StatsSetRef;
26use crate::vtable;
27use crate::vtable::ArrayId;
28use crate::vtable::VTable;
29pub(crate) mod canonical;
30mod operations;
31mod validity;
32
33vtable!(Constant);
34
35#[derive(Debug)]
36pub struct ConstantVTable;
37
38impl ConstantVTable {
39    pub const ID: ArrayId = ArrayId::new_ref("vortex.constant");
40}
41
42impl VTable for ConstantVTable {
43    type Array = ConstantArray;
44
45    type Metadata = Scalar;
46    type OperationsVTable = Self;
47    type ValidityVTable = Self;
48
49    fn id(_array: &Self::Array) -> ArrayId {
50        Self::ID
51    }
52
53    fn len(array: &ConstantArray) -> usize {
54        array.len
55    }
56
57    fn dtype(array: &ConstantArray) -> &DType {
58        array.scalar.dtype()
59    }
60
61    fn stats(array: &ConstantArray) -> StatsSetRef<'_> {
62        array.stats_set.to_ref(array.as_ref())
63    }
64
65    fn array_hash<H: std::hash::Hasher>(
66        array: &ConstantArray,
67        state: &mut H,
68        _precision: Precision,
69    ) {
70        array.scalar.hash(state);
71        array.len.hash(state);
72    }
73
74    fn array_eq(array: &ConstantArray, other: &ConstantArray, _precision: Precision) -> bool {
75        array.scalar == other.scalar && array.len == other.len
76    }
77
78    fn nbuffers(_array: &ConstantArray) -> usize {
79        1
80    }
81
82    fn buffer(array: &ConstantArray, idx: usize) -> BufferHandle {
83        match idx {
84            0 => BufferHandle::new_host(
85                ScalarValue::to_proto_bytes::<ByteBufferMut>(array.scalar.value()).freeze(),
86            ),
87            _ => vortex_panic!("ConstantArray buffer index {idx} out of bounds"),
88        }
89    }
90
91    fn buffer_name(_array: &ConstantArray, idx: usize) -> Option<String> {
92        match idx {
93            0 => Some("scalar".to_string()),
94            _ => None,
95        }
96    }
97
98    fn nchildren(_array: &ConstantArray) -> usize {
99        0
100    }
101
102    fn child(_array: &ConstantArray, idx: usize) -> ArrayRef {
103        vortex_panic!("ConstantArray child index {idx} out of bounds")
104    }
105
106    fn child_name(_array: &ConstantArray, idx: usize) -> String {
107        vortex_panic!("ConstantArray child_name index {idx} out of bounds")
108    }
109
110    fn metadata(array: &ConstantArray) -> VortexResult<Self::Metadata> {
111        Ok(array.scalar().clone())
112    }
113
114    fn serialize(_metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
115        // HACK: Because the scalar is stored in the buffers, we do not need to serialize the
116        // metadata at all.
117        Ok(Some(vec![]))
118    }
119
120    fn deserialize(
121        _bytes: &[u8],
122        dtype: &DType,
123        _len: usize,
124        buffers: &[BufferHandle],
125        _session: &VortexSession,
126    ) -> VortexResult<Self::Metadata> {
127        vortex_ensure!(
128            buffers.len() == 1,
129            "Expected 1 buffer, got {}",
130            buffers.len()
131        );
132
133        let buffer = buffers[0].clone().try_to_host_sync()?;
134        let bytes: &[u8] = buffer.as_ref();
135
136        let scalar_value = ScalarValue::from_proto_bytes(bytes, dtype)?;
137        let scalar = Scalar::try_new(dtype.clone(), scalar_value)?;
138
139        Ok(scalar)
140    }
141
142    fn build(
143        _dtype: &DType,
144        len: usize,
145        metadata: &Self::Metadata,
146        _buffers: &[BufferHandle],
147        _children: &dyn ArrayChildren,
148    ) -> VortexResult<ConstantArray> {
149        Ok(ConstantArray::new(metadata.clone(), len))
150    }
151
152    fn with_children(_array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
153        vortex_ensure!(
154            children.is_empty(),
155            "ConstantArray has no children, got {}",
156            children.len()
157        );
158        Ok(())
159    }
160
161    fn reduce_parent(
162        array: &Self::Array,
163        parent: &ArrayRef,
164        child_idx: usize,
165    ) -> VortexResult<Option<ArrayRef>> {
166        PARENT_RULES.evaluate(array, parent, child_idx)
167    }
168
169    fn execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
170        Ok(constant_canonicalize(array)?.into_array())
171    }
172}