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::arrays::ScalarFnVTable;
11use crate::expr::functions::scalar::ScalarFn;
12use crate::stats::ArrayStats;
13use crate::vtable::ArrayVTable;
14use crate::vtable::ArrayVTableExt;
15
16#[derive(Clone, Debug)]
17pub struct ScalarFnArray {
18    // NOTE(ngates): we should fix vtables so we don't have to hold this
19    pub(super) vtable: ArrayVTable,
20    pub(super) scalar_fn: ScalarFn,
21    pub(super) dtype: DType,
22    pub(super) len: usize,
23    pub(super) children: Vec<ArrayRef>,
24    pub(super) stats: ArrayStats,
25}
26
27impl ScalarFnArray {
28    /// Create a new ScalarFnArray from a scalar function and its children.
29    pub fn try_new(scalar_fn: ScalarFn, children: Vec<ArrayRef>, len: usize) -> 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::new(scalar_fn.vtable().clone()).into_vtable(),
40            scalar_fn,
41            dtype,
42            len,
43            children,
44            stats: Default::default(),
45        })
46    }
47}