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::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(bound: ScalarFn, children: Vec<ArrayRef>, len: usize) -> VortexResult<Self> {
30        let arg_dtypes: Vec<_> = children.iter().map(|c| c.dtype().clone()).collect();
31        let dtype = bound.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(bound.vtable().clone()).into_vtable(),
40            scalar_fn: bound,
41            dtype,
42            len,
43            children,
44            stats: Default::default(),
45        })
46    }
47
48    /// Get the scalar function bound to this array.
49    pub fn scalar_fn(&self) -> &ScalarFn {
50        &self.scalar_fn
51    }
52
53    /// Get the children arrays of this scalar function array.
54    pub fn children(&self) -> &[ArrayRef] {
55        &self.children
56    }
57}