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