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 vortex_error::VortexExpect;
8use vortex_error::VortexResult;
9use vortex_error::vortex_bail;
10use vortex_error::vortex_ensure;
11
12use crate::ArrayRef;
13use crate::ArraySlots;
14use crate::array::Array;
15use crate::array::ArrayParts;
16use crate::array::TypedArrayRef;
17use crate::arrays::ScalarFn;
18use crate::scalar_fn::ScalarFnRef;
19
20// ScalarFnArray has a variable number of slots (one per child)
21
22#[derive(Clone, Debug)]
23pub struct ScalarFnData {
24    pub(super) scalar_fn: ScalarFnRef,
25}
26
27impl Display for ScalarFnData {
28    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
29        write!(f, "scalar_fn: {}", self.scalar_fn)
30    }
31}
32
33impl ScalarFnData {
34    /// Get the scalar function bound to this array.
35    #[allow(clippy::inline_always)]
36    #[inline(always)]
37    pub fn scalar_fn(&self) -> &ScalarFnRef {
38        &self.scalar_fn
39    }
40}
41
42pub trait ScalarFnArrayExt: TypedArrayRef<ScalarFn> {
43    fn scalar_fn(&self) -> &ScalarFnRef {
44        &self.scalar_fn
45    }
46
47    fn child_at(&self, idx: usize) -> &ArrayRef {
48        self.as_ref().slots()[idx]
49            .as_ref()
50            .vortex_expect("ScalarFnArray child slot")
51    }
52
53    fn child_count(&self) -> usize {
54        self.as_ref().slots().len()
55    }
56
57    fn nchildren(&self) -> usize {
58        self.child_count()
59    }
60
61    fn get_child(&self, idx: usize) -> &ArrayRef {
62        self.child_at(idx)
63    }
64
65    fn iter_children(&self) -> impl Iterator<Item = &ArrayRef> + '_ {
66        (0..self.child_count()).map(|idx| self.child_at(idx))
67    }
68
69    fn children(&self) -> Vec<ArrayRef> {
70        self.iter_children().cloned().collect()
71    }
72}
73impl<T: TypedArrayRef<ScalarFn>> ScalarFnArrayExt for T {}
74
75impl Array<ScalarFn> {
76    /// Create a new ScalarFnArray from a scalar function and its children.
77    pub fn try_new(scalar_fn: ScalarFnRef, children: Vec<ArrayRef>) -> VortexResult<Self> {
78        let len = Self::infer_len(&children)?;
79        Self::try_new_with_len(scalar_fn, children, len)
80    }
81
82    /// Create a new ScalarFnArray from a scalar function, children, and an explicit length.
83    ///
84    /// This is needed for zero-child scalar functions and deserialization paths where there is no
85    /// child array to infer the length from.
86    pub fn try_new_with_len(
87        scalar_fn: ScalarFnRef,
88        children: Vec<ArrayRef>,
89        len: usize,
90    ) -> VortexResult<Self> {
91        Self::validate_arity(&scalar_fn, children.len())?;
92        Self::validate_children_len(&children, len)?;
93
94        let arg_dtypes: Vec<_> = children.iter().map(|c| c.dtype().clone()).collect();
95        let dtype = scalar_fn.return_dtype(&arg_dtypes)?;
96        let data = ScalarFnData {
97            scalar_fn: scalar_fn.clone(),
98        };
99        let vtable = ScalarFn { id: scalar_fn.id() };
100
101        Ok(unsafe {
102            Array::from_parts_unchecked(
103                ArrayParts::new(vtable, dtype, len, data)
104                    .with_slots(children.into_iter().map(Some).collect::<ArraySlots>()),
105            )
106        })
107    }
108
109    fn infer_len(children: &[ArrayRef]) -> VortexResult<usize> {
110        let Some(child) = children.first() else {
111            vortex_bail!("ScalarFnArray length cannot be inferred without children");
112        };
113        Ok(child.len())
114    }
115
116    fn validate_arity(scalar_fn: &ScalarFnRef, child_count: usize) -> VortexResult<()> {
117        let arity = scalar_fn.signature().arity();
118        vortex_ensure!(
119            arity.matches(child_count),
120            "ScalarFnArray requires {arity} children, got {child_count}"
121        );
122        Ok(())
123    }
124
125    fn validate_children_len(children: &[ArrayRef], len: usize) -> VortexResult<()> {
126        vortex_ensure!(
127            children.iter().all(|c| c.len() == len),
128            "ScalarFnArray must have children equal to the array length"
129        );
130        Ok(())
131    }
132}
vortex_array/arrays/scalar_fn/array.rs

vortex_array/arrays/scalar_fn/
array.rs
