vortex_array/arrays/scalar_fn/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::fmt::Formatter;
6
7use smallvec::SmallVec;
8use vortex_error::VortexExpect;
9use vortex_error::VortexResult;
10use vortex_error::vortex_bail;
11use vortex_error::vortex_ensure;
12
13use crate::ArrayRef;
14use crate::ArraySlots;
15use crate::array::Array;
16use crate::array::ArrayParts;
17use crate::array::TypedArrayRef;
18use crate::arrays::ScalarFn;
19use crate::dtype::DType;
20use crate::scalar_fn::ScalarFnRef;
21
22// ScalarFnArray has a variable number of slots (one per child)
23
24#[derive(Clone, Debug)]
25pub struct ScalarFnData {
26    pub(super) scalar_fn: ScalarFnRef,
27}
28
29impl Display for ScalarFnData {
30    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
31        write!(f, "scalar_fn: {}", self.scalar_fn)
32    }
33}
34
35impl ScalarFnData {
36    /// Get the scalar function bound to this array.
37    #[allow(clippy::inline_always)]
38    #[inline(always)]
39    pub fn scalar_fn(&self) -> &ScalarFnRef {
40        &self.scalar_fn
41    }
42}
43
44pub trait ScalarFnArrayExt: TypedArrayRef<ScalarFn> {
45    fn scalar_fn(&self) -> &ScalarFnRef {
46        &self.scalar_fn
47    }
48
49    fn child_at(&self, idx: usize) -> &ArrayRef {
50        self.as_ref().slots()[idx]
51            .as_ref()
52            .vortex_expect("ScalarFnArray child slot")
53    }
54
55    fn child_count(&self) -> usize {
56        self.as_ref().slots().len()
57    }
58
59    fn nchildren(&self) -> usize {
60        self.child_count()
61    }
62
63    fn get_child(&self, idx: usize) -> &ArrayRef {
64        self.child_at(idx)
65    }
66
67    fn iter_children(&self) -> impl Iterator<Item = &ArrayRef> + '_ {
68        (0..self.child_count()).map(|idx| self.child_at(idx))
69    }
70
71    fn children(&self) -> Vec<ArrayRef> {
72        self.iter_children().cloned().collect()
73    }
74}
75impl<T: TypedArrayRef<ScalarFn>> ScalarFnArrayExt for T {}
76
77impl Array<ScalarFn> {
78    /// Create a new ScalarFnArray from a scalar function and its children.
79    pub fn try_new(
80        scalar_fn: ScalarFnRef,
81        children: impl IntoIterator<Item = ArrayRef>,
82    ) -> VortexResult<Self> {
83        let slots: ArraySlots = children.into_iter().map(Some).collect();
84        let len = Self::infer_len(&slots)?;
85        Self::try_new_from_slots(scalar_fn, slots, len)
86    }
87
88    /// Create a new ScalarFnArray from a scalar function, children, and an explicit length.
89    ///
90    /// This is needed for zero-child scalar functions and deserialization paths where there is no
91    /// child array to infer the length from.
92    pub fn try_new_with_len(
93        scalar_fn: ScalarFnRef,
94        children: impl IntoIterator<Item = ArrayRef>,
95        len: usize,
96    ) -> VortexResult<Self> {
97        let slots: ArraySlots = children.into_iter().map(Some).collect();
98        Self::try_new_from_slots(scalar_fn, slots, len)
99    }
100
101    #[inline]
102    fn try_new_from_slots(
103        scalar_fn: ScalarFnRef,
104        slots: ArraySlots,
105        len: usize,
106    ) -> VortexResult<Self> {
107        Self::validate_arity(&scalar_fn, slots.len())?;
108        Self::validate_children_len(&slots, len)?;
109
110        let arg_dtypes = Self::arg_dtypes(&slots);
111        let dtype = scalar_fn.return_dtype(&arg_dtypes)?;
112
113        Ok(unsafe { Self::new_from_validated_slots(scalar_fn, dtype, slots, len) })
114    }
115
116    /// # Safety
117    /// The caller must ensure arity and child lengths were validated and
118    /// "dtype" is scalar function's return dtype.
119    #[inline]
120    unsafe fn new_from_validated_slots(
121        scalar_fn: ScalarFnRef,
122        dtype: DType,
123        slots: ArraySlots,
124        len: usize,
125    ) -> Self {
126        let vtable = ScalarFn { id: scalar_fn.id() };
127        let data = ScalarFnData { scalar_fn };
128        unsafe { Array::from_parts_unchecked(ArrayParts::new(vtable, dtype, len, data, slots)) }
129    }
130
131    #[inline]
132    fn arg_dtypes(slots: &[Option<ArrayRef>]) -> SmallVec<[DType; 4]> {
133        let mut arg_dtypes = SmallVec::with_capacity(slots.len());
134        for child in slots.iter().flatten() {
135            arg_dtypes.push(child.dtype().clone());
136        }
137        arg_dtypes
138    }
139
140    #[inline]
141    fn infer_len(slots: &[Option<ArrayRef>]) -> VortexResult<usize> {
142        let Some(child) = slots.first().and_then(Option::as_ref) else {
143            vortex_bail!("ScalarFnArray length cannot be inferred without children");
144        };
145        Ok(child.len())
146    }
147
148    #[inline]
149    fn validate_arity(scalar_fn: &ScalarFnRef, child_count: usize) -> VortexResult<()> {
150        let arity = scalar_fn.signature().arity();
151        vortex_ensure!(
152            arity.matches(child_count),
153            "ScalarFnArray requires {arity} children, got {child_count}"
154        );
155        Ok(())
156    }
157
158    #[inline]
159    fn validate_children_len(slots: &[Option<ArrayRef>], len: usize) -> VortexResult<()> {
160        for child in slots.iter().flatten() {
161            vortex_ensure!(
162                child.len() == len,
163                "ScalarFnArray must have children equal to the array length"
164            );
165        }
166        Ok(())
167    }
168}
vortex_array/arrays/scalar_fn/array.rs

vortex_array/arrays/scalar_fn/
array.rs
