Skip to main content

vortex_array/arrays/scalar_fn/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5use vortex_error::vortex_ensure;
6
7use crate::ArrayRef;
8use crate::DynArray;
9use crate::arrays::ScalarFnVTable;
10use crate::dtype::DType;
11use crate::scalar_fn::ScalarFnRef;
12use crate::stats::ArrayStats;
13
14#[derive(Clone, Debug)]
15pub struct ScalarFnArray {
16    pub(super) vtable: ScalarFnVTable,
17    pub(super) dtype: DType,
18    pub(super) len: usize,
19    pub(super) children: Vec<ArrayRef>,
20    pub(super) stats: ArrayStats,
21}
22
23impl ScalarFnArray {
24    /// Create a new ScalarFnArray from a scalar function and its children.
25    pub fn try_new(
26        scalar_fn: ScalarFnRef,
27        children: Vec<ArrayRef>,
28        len: usize,
29    ) -> VortexResult<Self> {
30        let arg_dtypes: Vec<_> = children.iter().map(|c| c.dtype().clone()).collect();
31        let dtype = scalar_fn.return_dtype(&arg_dtypes)?;
32
33        vortex_ensure!(
34            children.iter().all(|c| c.len() == len),
35            "ScalarFnArray must have children equal to the array length"
36        );
37
38        Ok(Self {
39            vtable: ScalarFnVTable { scalar_fn },
40            dtype,
41            len,
42            children,
43            stats: Default::default(),
44        })
45    }
46
47    /// Get the scalar function bound to this array.
48    #[allow(clippy::same_name_method)]
49    #[inline(always)]
50    pub fn scalar_fn(&self) -> &ScalarFnRef {
51        &self.vtable.scalar_fn
52    }
53
54    /// Get the children arrays of this scalar function array.
55    pub fn children(&self) -> &[ArrayRef] {
56        &self.children
57    }
58}