vortex_array/arrays/scalar_fn/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3mod operations;
4mod validity;
5
6use std::hash::Hash;
7use std::hash::Hasher;
8use std::marker::PhantomData;
9use std::ops::Deref;
10
11use itertools::Itertools;
12use validity::ValidityArray;
13use vortex_error::VortexResult;
14use vortex_error::vortex_bail;
15use vortex_error::vortex_ensure;
16use vortex_error::vortex_panic;
17use vortex_session::VortexSession;
18
19use crate::ArrayEq;
20use crate::ArrayHash;
21use crate::ArrayRef;
22use crate::EqMode;
23use crate::IntoArray;
24use crate::array::Array;
25use crate::array::ArrayId;
26use crate::array::ArrayParts;
27use crate::array::ArrayView;
28use crate::array::VTable;
29use crate::array::ValidityVTable;
30use crate::array::with_empty_buffers;
31use crate::arrays::scalar_fn::array::ScalarFnArrayExt;
32use crate::arrays::scalar_fn::array::ScalarFnData;
33use crate::arrays::scalar_fn::rules::PARENT_RULES;
34use crate::arrays::scalar_fn::rules::RULES;
35use crate::buffer::BufferHandle;
36use crate::dtype::DType;
37use crate::executor::ExecutionCtx;
38use crate::executor::ExecutionResult;
39use crate::matcher::Matcher;
40use crate::scalar_fn;
41use crate::scalar_fn::ArrayReduceNode;
42use crate::scalar_fn::ReduceNode;
43use crate::scalar_fn::ReduceNodeValidity;
44use crate::scalar_fn::ScalarFnId;
45use crate::scalar_fn::VecExecutionArgs;
46use crate::serde::ArrayChildren;
47use crate::validity::Validity;
48
49/// A [`ScalarFn`]-encoded Vortex array.
50pub type ScalarFnArray = Array<ScalarFn>;
51
52#[derive(Clone, Debug)]
53pub struct ScalarFn {
54    pub(super) id: ScalarFnId,
55}
56
57impl ArrayHash for ScalarFnData {
58    fn array_hash<H: Hasher>(&self, state: &mut H, _accuracy: EqMode) {
59        self.scalar_fn().hash(state);
60    }
61}
62
63impl ArrayEq for ScalarFnData {
64    fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool {
65        self.scalar_fn() == other.scalar_fn()
66    }
67}
68
69impl VTable for ScalarFn {
70    type TypedArrayData = ScalarFnData;
71    type OperationsVTable = Self;
72    type ValidityVTable = Self;
73
74    fn id(&self) -> ArrayId {
75        self.id
76    }
77
78    fn validate(
79        &self,
80        data: &ScalarFnData,
81        dtype: &DType,
82        len: usize,
83        slots: &[Option<ArrayRef>],
84    ) -> VortexResult<()> {
85        let scalar_fn = data.scalar_fn();
86        vortex_ensure!(
87            scalar_fn.id() == self.id,
88            "ScalarFnArray data scalar_fn does not match vtable"
89        );
90
91        let missing_children = slots.iter().filter(|slot| slot.is_none()).count();
92        vortex_ensure!(
93            missing_children == 0,
94            "ScalarFnArray requires every child slot to be present, got {missing_children} missing"
95        );
96
97        let arity = scalar_fn.signature().arity();
98        vortex_ensure!(
99            arity.matches(slots.len()),
100            "ScalarFnArray requires {arity} children, got {}",
101            slots.len()
102        );
103        vortex_ensure!(
104            slots.iter().flatten().all(|c| c.len() == len),
105            "All child arrays must have the same length as the scalar function array"
106        );
107
108        let child_dtypes = slots
109            .iter()
110            .flatten()
111            .map(|c| c.dtype().clone())
112            .collect_vec();
113        vortex_ensure!(
114            scalar_fn.return_dtype(&child_dtypes)? == *dtype,
115            "ScalarFnArray dtype does not match scalar function return dtype"
116        );
117        Ok(())
118    }
119
120    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
121        0
122    }
123
124    fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
125        vortex_panic!("ScalarFnArray buffer index {idx} out of bounds")
126    }
127
128    fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
129        None
130    }
131
132    fn with_buffers(
133        &self,
134        array: ArrayView<'_, Self>,
135        buffers: &[BufferHandle],
136    ) -> VortexResult<ArrayParts<Self>> {
137        with_empty_buffers(self, array, buffers)
138    }
139
140    fn serialize(
141        _array: ArrayView<'_, Self>,
142        _session: &VortexSession,
143    ) -> VortexResult<Option<Vec<u8>>> {
144        // Not supported
145        Ok(None)
146    }
147
148    fn deserialize(
149        &self,
150        _dtype: &DType,
151        _len: usize,
152        _metadata: &[u8],
153        _buffers: &[BufferHandle],
154        _children: &dyn ArrayChildren,
155        _session: &VortexSession,
156    ) -> VortexResult<ArrayParts<Self>> {
157        vortex_bail!("Deserialization of ScalarFnVTable metadata is not supported");
158    }
159
160    fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String {
161        array
162            .scalar_fn()
163            .signature()
164            .child_name(idx)
165            .as_ref()
166            .to_string()
167    }
168
169    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
170        let args = VecExecutionArgs::new(array.children(), array.len());
171        array
172            .scalar_fn()
173            .execute(&args, ctx)
174            .map(ExecutionResult::done)
175    }
176
177    fn reduce(array: ArrayView<'_, Self>) -> VortexResult<Option<ArrayRef>> {
178        RULES.evaluate(array)
179    }
180
181    fn reduce_parent(
182        array: ArrayView<'_, Self>,
183        parent: &ArrayRef,
184        child_idx: usize,
185    ) -> VortexResult<Option<ArrayRef>> {
186        PARENT_RULES.evaluate(array, parent, child_idx)
187    }
188}
189
190/// A matcher that matches any scalar function expression.
191#[derive(Debug)]
192pub struct AnyScalarFn;
193impl Matcher for AnyScalarFn {
194    type Match<'a> = ArrayView<'a, ScalarFn>;
195
196    fn matches(array: &ArrayRef) -> bool {
197        array.is::<ScalarFn>()
198    }
199
200    fn try_match(array: &ArrayRef) -> Option<Self::Match<'_>> {
201        array.as_opt::<ScalarFn>()
202    }
203}
204
205/// A matcher that matches a specific scalar function expression.
206#[derive(Debug, Default)]
207pub struct ExactScalarFn<F: scalar_fn::ScalarFnVTable>(PhantomData<F>);
208
209impl<F: scalar_fn::ScalarFnVTable> Matcher for ExactScalarFn<F> {
210    type Match<'a> = ScalarFnArrayView<'a, F>;
211
212    fn matches(array: &ArrayRef) -> bool {
213        if let Some(scalar_fn_array) = array.as_opt::<ScalarFn>() {
214            scalar_fn_array.data().scalar_fn().is::<F>()
215        } else {
216            false
217        }
218    }
219
220    fn try_match(array: &ArrayRef) -> Option<Self::Match<'_>> {
221        let scalar_fn_array = array.as_opt::<ScalarFn>()?;
222        let scalar_fn_data = scalar_fn_array.data();
223        let scalar_fn = scalar_fn_data.scalar_fn().downcast_ref::<F>()?;
224        Some(ScalarFnArrayView {
225            array,
226            vtable: scalar_fn.vtable(),
227            options: scalar_fn.options(),
228        })
229    }
230}
231
232pub struct ScalarFnArrayView<'a, F: scalar_fn::ScalarFnVTable> {
233    array: &'a ArrayRef,
234    pub vtable: &'a F,
235    pub options: &'a F::Options,
236}
237
238impl<F: scalar_fn::ScalarFnVTable> Deref for ScalarFnArrayView<'_, F> {
239    type Target = ArrayRef;
240
241    fn deref(&self) -> &Self::Target {
242        self.array
243    }
244}
245
246impl ValidityVTable<ScalarFn> for ScalarFn {
247    fn validity(view: ArrayView<'_, ScalarFn>) -> VortexResult<Validity> {
248        let node = ArrayReduceNode::new(view.as_ref());
249        Ok(Validity::Array(match node.validity()? {
250            ReduceNodeValidity::Reduced(reduced) => reduced.into_array(),
251            ReduceNodeValidity::Irreducible => ValidityArray::new(node.into_array()).into_array(),
252        }))
253    }
254}
vortex_array/arrays/scalar_fn/vtable/mod.rs

vortex_array/arrays/scalar_fn/vtable/
mod.rs
