vortex_array/arrays/scalar_fn/vtable/
operations.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use crate::ExecutionCtx;
7use crate::IntoArray;
8use crate::array::ArrayView;
9use crate::array::OperationsVTable;
10use crate::arrays::ConstantArray;
11use crate::arrays::scalar_fn::ScalarFnArrayExt;
12use crate::arrays::scalar_fn::vtable::ScalarFnVTable;
13use crate::columnar::Columnar;
14use crate::scalar::Scalar;
15use crate::scalar_fn::VecExecutionArgs;
16
17impl OperationsVTable<ScalarFnVTable> for ScalarFnVTable {
18    fn scalar_at(
19        array: ArrayView<'_, ScalarFnVTable>,
20        index: usize,
21        ctx: &mut ExecutionCtx,
22    ) -> VortexResult<Scalar> {
23        let inputs: Vec<_> = array
24            .children()
25            .iter()
26            .map(|child| Ok(ConstantArray::new(child.execute_scalar(index, ctx)?, 1).into_array()))
27            .collect::<VortexResult<_>>()?;
28
29        let args = VecExecutionArgs::new(inputs, 1);
30        let result = array.scalar_fn().execute(&args, ctx)?;
31
32        let scalar = match result.execute::<Columnar>(ctx)? {
33            Columnar::Canonical(arr) => {
34                tracing::info!(
35                    "Scalar function {} returned non-constant array from execution over all scalar inputs",
36                    array.scalar_fn(),
37                );
38                arr.into_array().execute_scalar(0, ctx)?
39            }
40            Columnar::Constant(constant) => constant.scalar().clone(),
41        };
42
43        debug_assert_eq!(
44            scalar.dtype(),
45            array.dtype(),
46            "Scalar function {} returned dtype {:?} but expected {:?}",
47            array.scalar_fn(),
48            scalar.dtype(),
49            array.dtype()
50        );
51
52        Ok(scalar)
53    }
54}
55
56#[cfg(test)]
57mod tests {
58    use vortex_buffer::buffer;
59    use vortex_error::VortexResult;
60
61    use crate::Canonical;
62    use crate::IntoArray;
63    use crate::LEGACY_SESSION;
64    use crate::VortexSessionExecute;
65    use crate::arrays::BoolArray;
66    use crate::arrays::PrimitiveArray;
67    use crate::arrays::ScalarFnArray;
68    use crate::assert_arrays_eq;
69    use crate::scalar_fn::ScalarFn;
70    use crate::scalar_fn::fns::binary::Binary;
71    use crate::scalar_fn::fns::operators::Operator;
72    use crate::validity::Validity;
73
74    #[test]
75    fn test_scalar_fn_add() -> VortexResult<()> {
76        let lhs = buffer![1i32, 2, 3].into_array();
77        let rhs = buffer![10i32, 20, 30].into_array();
78
79        let scalar_fn = ScalarFn::new(Binary, Operator::Add).erased();
80        let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs], 3)?;
81
82        let result = scalar_fn_array
83            .into_array()
84            .execute::<Canonical>(&mut LEGACY_SESSION.create_execution_ctx())?
85            .into_array();
86        let expected = buffer![11i32, 22, 33].into_array();
87        assert_arrays_eq!(result, expected);
88
89        Ok(())
90    }
91
92    #[test]
93    fn test_scalar_fn_mul() -> VortexResult<()> {
94        let lhs = buffer![2i32, 3, 4].into_array();
95        let rhs = buffer![5i32, 6, 7].into_array();
96
97        let scalar_fn = ScalarFn::new(Binary, Operator::Mul).erased();
98        let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs], 3)?;
99
100        let result = scalar_fn_array
101            .into_array()
102            .execute::<Canonical>(&mut LEGACY_SESSION.create_execution_ctx())?
103            .into_array();
104        let expected = buffer![10i32, 18, 28].into_array();
105        assert_arrays_eq!(result, expected);
106
107        Ok(())
108    }
109
110    #[test]
111    fn test_scalar_fn_with_nullable() -> VortexResult<()> {
112        let lhs = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllValid).into_array();
113        let rhs = PrimitiveArray::new(
114            buffer![10i32, 20, 30],
115            Validity::from_iter([true, false, true]),
116        )
117        .into_array();
118
119        let scalar_fn = ScalarFn::new(Binary, Operator::Add).erased();
120        let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs], 3)?;
121
122        let result = scalar_fn_array
123            .into_array()
124            .execute::<Canonical>(&mut LEGACY_SESSION.create_execution_ctx())?
125            .into_array();
126        let expected = PrimitiveArray::new(
127            buffer![11i32, 0, 33],
128            Validity::from_iter([true, false, true]),
129        )
130        .into_array();
131        assert_arrays_eq!(result, expected);
132
133        Ok(())
134    }
135
136    #[test]
137    fn test_scalar_fn_comparison() -> VortexResult<()> {
138        let lhs = buffer![1i32, 5, 3].into_array();
139        let rhs = buffer![2i32, 5, 1].into_array();
140
141        let scalar_fn = ScalarFn::new(Binary, Operator::Eq).erased();
142        let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs], 3)?;
143
144        let result = scalar_fn_array
145            .into_array()
146            .execute::<Canonical>(&mut LEGACY_SESSION.create_execution_ctx())?
147            .into_array();
148        let expected = BoolArray::from_iter([false, true, false]).into_array();
149        assert_arrays_eq!(result, expected);
150
151        Ok(())
152    }
153}
vortex_array/arrays/scalar_fn/vtable/operations.rs

vortex_array/arrays/scalar_fn/vtable/
operations.rs
