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;
5use std::fmt::Display;
6use std::fmt::Formatter;
7use std::hash::Hash;
8use std::hash::Hasher;
9use std::marker::PhantomData;
10use std::ops::Deref;
11
12use itertools::Itertools;
13use vortex_error::VortexResult;
14use vortex_error::vortex_bail;
15use vortex_error::vortex_ensure;
16use vortex_error::vortex_panic;
17use vortex_session::VortexSession;
18use vortex_session::registry::CachedId;
19
20use crate::ArrayEq;
21use crate::ArrayHash;
22use crate::ArrayRef;
23use crate::EqMode;
24use crate::array::Array;
25use crate::array::ArrayId;
26use crate::array::ArrayParts;
27use crate::array::ArrayView;
28use crate::array::VTable;
29use crate::array::with_empty_buffers;
30use crate::arrays::scalar_fn::array::ScalarFnArrayExt;
31use crate::arrays::scalar_fn::array::ScalarFnData;
32use crate::arrays::scalar_fn::rules::PARENT_RULES;
33use crate::arrays::scalar_fn::rules::RULES;
34use crate::buffer::BufferHandle;
35use crate::dtype::DType;
36use crate::executor::ExecutionCtx;
37use crate::executor::ExecutionResult;
38use crate::expr::Expression;
39use crate::expr::display::ExprDisplay;
40use crate::matcher::Matcher;
41use crate::scalar_fn;
42use crate::scalar_fn::Arity;
43use crate::scalar_fn::ChildName;
44use crate::scalar_fn::ExecutionArgs;
45use crate::scalar_fn::ScalarFnId;
46use crate::scalar_fn::ScalarFnVTableExt;
47use crate::scalar_fn::VecExecutionArgs;
48use crate::serde::ArrayChildren;
49
50/// A [`ScalarFn`]-encoded Vortex array.
51pub type ScalarFnArray = Array<ScalarFn>;
52
53#[derive(Clone, Debug)]
54pub struct ScalarFn {
55    pub(super) id: ScalarFnId,
56}
57
58impl ArrayHash for ScalarFnData {
59    fn array_hash<H: Hasher>(&self, state: &mut H, _accuracy: EqMode) {
60        self.scalar_fn().hash(state);
61    }
62}
63
64impl ArrayEq for ScalarFnData {
65    fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool {
66        self.scalar_fn() == other.scalar_fn()
67    }
68}
69
70impl VTable for ScalarFn {
71    type TypedArrayData = ScalarFnData;
72    type OperationsVTable = Self;
73    type ValidityVTable = Self;
74
75    fn id(&self) -> ArrayId {
76        self.id
77    }
78
79    fn validate(
80        &self,
81        data: &ScalarFnData,
82        dtype: &DType,
83        len: usize,
84        slots: &[Option<ArrayRef>],
85    ) -> VortexResult<()> {
86        let scalar_fn = data.scalar_fn();
87        vortex_ensure!(
88            scalar_fn.id() == self.id,
89            "ScalarFnArray data scalar_fn does not match vtable"
90        );
91
92        let missing_children = slots.iter().filter(|slot| slot.is_none()).count();
93        vortex_ensure!(
94            missing_children == 0,
95            "ScalarFnArray requires every child slot to be present, got {missing_children} missing"
96        );
97
98        let arity = scalar_fn.signature().arity();
99        vortex_ensure!(
100            arity.matches(slots.len()),
101            "ScalarFnArray requires {arity} children, got {}",
102            slots.len()
103        );
104        vortex_ensure!(
105            slots.iter().flatten().all(|c| c.len() == len),
106            "All child arrays must have the same length as the scalar function array"
107        );
108
109        let child_dtypes = slots
110            .iter()
111            .flatten()
112            .map(|c| c.dtype().clone())
113            .collect_vec();
114        vortex_ensure!(
115            scalar_fn.return_dtype(&child_dtypes)? == *dtype,
116            "ScalarFnArray dtype does not match scalar function return dtype"
117        );
118        Ok(())
119    }
120
121    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
122        0
123    }
124
125    fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
126        vortex_panic!("ScalarFnArray buffer index {idx} out of bounds")
127    }
128
129    fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
130        None
131    }
132
133    fn with_buffers(
134        &self,
135        array: ArrayView<'_, Self>,
136        buffers: &[BufferHandle],
137    ) -> VortexResult<ArrayParts<Self>> {
138        with_empty_buffers(self, array, buffers)
139    }
140
141    fn serialize(
142        _array: ArrayView<'_, Self>,
143        _session: &VortexSession,
144    ) -> VortexResult<Option<Vec<u8>>> {
145        // Not supported
146        Ok(None)
147    }
148
149    fn deserialize(
150        &self,
151        _dtype: &DType,
152        _len: usize,
153        _metadata: &[u8],
154        _buffers: &[BufferHandle],
155        _children: &dyn ArrayChildren,
156        _session: &VortexSession,
157    ) -> VortexResult<ArrayParts<Self>> {
158        vortex_bail!("Deserialization of ScalarFnVTable metadata is not supported");
159    }
160
161    fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String {
162        array
163            .scalar_fn()
164            .signature()
165            .child_name(idx)
166            .as_ref()
167            .to_string()
168    }
169
170    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
171        let args = VecExecutionArgs::new(array.children(), array.len());
172        array
173            .scalar_fn()
174            .execute(&args, ctx)
175            .map(ExecutionResult::done)
176    }
177
178    fn reduce(array: ArrayView<'_, Self>) -> VortexResult<Option<ArrayRef>> {
179        RULES.evaluate(array)
180    }
181
182    fn reduce_parent(
183        array: ArrayView<'_, Self>,
184        parent: &ArrayRef,
185        child_idx: usize,
186    ) -> VortexResult<Option<ArrayRef>> {
187        PARENT_RULES.evaluate(array, parent, child_idx)
188    }
189}
190
191/// A matcher that matches any scalar function expression.
192#[derive(Debug)]
193pub struct AnyScalarFn;
194impl Matcher for AnyScalarFn {
195    type Match<'a> = ArrayView<'a, ScalarFn>;
196
197    fn matches(array: &ArrayRef) -> bool {
198        array.is::<ScalarFn>()
199    }
200
201    fn try_match(array: &ArrayRef) -> Option<Self::Match<'_>> {
202        array.as_opt::<ScalarFn>()
203    }
204}
205
206/// A matcher that matches a specific scalar function expression.
207#[derive(Debug, Default)]
208pub struct ExactScalarFn<F: scalar_fn::ScalarFnVTable>(PhantomData<F>);
209
210impl<F: scalar_fn::ScalarFnVTable> Matcher for ExactScalarFn<F> {
211    type Match<'a> = ScalarFnArrayView<'a, F>;
212
213    fn matches(array: &ArrayRef) -> bool {
214        if let Some(scalar_fn_array) = array.as_opt::<ScalarFn>() {
215            scalar_fn_array.data().scalar_fn().is::<F>()
216        } else {
217            false
218        }
219    }
220
221    fn try_match(array: &ArrayRef) -> Option<Self::Match<'_>> {
222        let scalar_fn_array = array.as_opt::<ScalarFn>()?;
223        let scalar_fn_data = scalar_fn_array.data();
224        let scalar_fn = scalar_fn_data.scalar_fn().downcast_ref::<F>()?;
225        Some(ScalarFnArrayView {
226            array,
227            vtable: scalar_fn.vtable(),
228            options: scalar_fn.options(),
229        })
230    }
231}
232
233pub struct ScalarFnArrayView<'a, F: scalar_fn::ScalarFnVTable> {
234    array: &'a ArrayRef,
235    pub vtable: &'a F,
236    pub options: &'a F::Options,
237}
238
239impl<F: scalar_fn::ScalarFnVTable> Deref for ScalarFnArrayView<'_, F> {
240    type Target = ArrayRef;
241
242    fn deref(&self) -> &Self::Target {
243        self.array
244    }
245}
246
247// Used only in this method to allow constrained using of Expression evaluate.
248#[derive(Clone)]
249struct ArrayExpr;
250
251#[derive(Clone, Debug)]
252struct FakeEq<T>(T);
253
254impl<T> PartialEq<Self> for FakeEq<T> {
255    fn eq(&self, _other: &Self) -> bool {
256        false
257    }
258}
259
260impl<T> Eq for FakeEq<T> {}
261
262impl<T> Hash for FakeEq<T> {
263    fn hash<H: Hasher>(&self, _state: &mut H) {}
264}
265
266impl Display for FakeEq<ArrayRef> {
267    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
268        write!(f, "{}", self.0.encoding_id())
269    }
270}
271
272impl scalar_fn::ScalarFnVTable for ArrayExpr {
273    type Options = FakeEq<ArrayRef>;
274
275    fn id(&self) -> ScalarFnId {
276        static ID: CachedId = CachedId::new("vortex.array");
277        *ID
278    }
279
280    fn arity(&self, _options: &Self::Options) -> Arity {
281        Arity::Exact(0)
282    }
283
284    fn child_name(&self, _options: &Self::Options, _child_idx: usize) -> ChildName {
285        todo!()
286    }
287
288    fn fmt_sql(
289        &self,
290        options: &Self::Options,
291        _expr: &dyn ExprDisplay,
292        f: &mut Formatter<'_>,
293    ) -> std::fmt::Result {
294        write!(f, "{}", options.0.encoding_id())
295    }
296
297    fn return_dtype(&self, options: &Self::Options, _arg_dtypes: &[DType]) -> VortexResult<DType> {
298        Ok(options.0.dtype().clone())
299    }
300
301    fn execute(
302        &self,
303        options: &Self::Options,
304        _args: &dyn ExecutionArgs,
305        ctx: &mut ExecutionCtx,
306    ) -> VortexResult<ArrayRef> {
307        crate::Executable::execute(options.0.clone(), ctx)
308    }
309
310    fn validity(
311        &self,
312        options: &Self::Options,
313        _expression: &Expression,
314    ) -> VortexResult<Option<Expression>> {
315        let validity_array = options.0.validity()?.to_array(options.0.len());
316        Ok(Some(ArrayExpr.new_expr(FakeEq(validity_array), [])))
317    }
318
319    fn is_strict(&self, _options: &Self::Options) -> bool {
320        true
321    }
322}
vortex_array/arrays/scalar_fn/vtable/mod.rs

vortex_array/arrays/scalar_fn/vtable/
mod.rs
