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

vortex_array/arrays/scalar_fn/vtable/
mod.rs
