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::ScalarFn;
13use crate::columnar::Columnar;
14use crate::scalar::Scalar;
15use crate::scalar_fn::VecExecutionArgs;
16
17impl OperationsVTable<ScalarFn> for ScalarFn {
18    fn scalar_at(
19        array: ArrayView<'_, ScalarFn>,
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::arrays::scalar_fn::ScalarFnArrayExt;
69    use crate::assert_arrays_eq;
70    use crate::scalar::Scalar;
71    use crate::scalar_fn::TypedScalarFnInstance;
72    use crate::scalar_fn::fns::binary::Binary;
73    use crate::scalar_fn::fns::literal::Literal;
74    use crate::scalar_fn::fns::operators::Operator;
75    use crate::validity::Validity;
76
77    #[test]
78    fn test_scalar_fn_add() -> VortexResult<()> {
79        let lhs = buffer![1i32, 2, 3].into_array();
80        let rhs = buffer![10i32, 20, 30].into_array();
81
82        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Add).erased();
83        let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?;
84
85        assert_eq!(scalar_fn_array.len(), 3);
86
87        let result = scalar_fn_array
88            .into_array()
89            .execute::<Canonical>(&mut LEGACY_SESSION.create_execution_ctx())?
90            .into_array();
91        let expected = buffer![11i32, 22, 33].into_array();
92        assert_arrays_eq!(result, expected);
93
94        Ok(())
95    }
96
97    #[test]
98    fn test_scalar_fn_inferred_len_rejects_mismatched_children() {
99        let lhs = buffer![1i32, 2, 3].into_array();
100        let rhs = buffer![10i32, 20].into_array();
101
102        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Add).erased();
103        let err = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])
104            .expect_err("ScalarFnArray::try_new must reject mismatched child lengths");
105
106        assert!(
107            err.to_string()
108                .contains("ScalarFnArray must have children equal to the array length")
109        );
110    }
111
112    #[test]
113    fn test_scalar_fn_without_children_requires_explicit_len() -> VortexResult<()> {
114        let scalar_fn = TypedScalarFnInstance::new(Literal, Scalar::from(1i32)).erased();
115
116        let Err(err) = ScalarFnArray::try_new(scalar_fn.clone(), vec![]) else {
117            panic!("ScalarFnArray::try_new should reject zero children");
118        };
119        assert!(
120            err.to_string()
121                .contains("ScalarFnArray length cannot be inferred without children")
122        );
123
124        let scalar_fn_array = ScalarFnArray::try_new_with_len(scalar_fn, vec![], 3)?;
125        assert_eq!(scalar_fn_array.len(), 3);
126        assert_eq!(scalar_fn_array.child_count(), 0);
127
128        Ok(())
129    }
130
131    #[test]
132    fn test_scalar_fn_mul() -> VortexResult<()> {
133        let lhs = buffer![2i32, 3, 4].into_array();
134        let rhs = buffer![5i32, 6, 7].into_array();
135
136        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Mul).erased();
137        let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?;
138
139        let result = scalar_fn_array
140            .into_array()
141            .execute::<Canonical>(&mut LEGACY_SESSION.create_execution_ctx())?
142            .into_array();
143        let expected = buffer![10i32, 18, 28].into_array();
144        assert_arrays_eq!(result, expected);
145
146        Ok(())
147    }
148
149    #[test]
150    fn test_scalar_fn_with_nullable() -> VortexResult<()> {
151        let lhs = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllValid).into_array();
152        let rhs = PrimitiveArray::new(
153            buffer![10i32, 20, 30],
154            Validity::from_iter([true, false, true]),
155        )
156        .into_array();
157
158        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Add).erased();
159        let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?;
160
161        let result = scalar_fn_array
162            .into_array()
163            .execute::<Canonical>(&mut LEGACY_SESSION.create_execution_ctx())?
164            .into_array();
165        let expected = PrimitiveArray::new(
166            buffer![11i32, 0, 33],
167            Validity::from_iter([true, false, true]),
168        )
169        .into_array();
170        assert_arrays_eq!(result, expected);
171
172        Ok(())
173    }
174
175    #[test]
176    fn test_scalar_fn_comparison() -> VortexResult<()> {
177        let lhs = buffer![1i32, 5, 3].into_array();
178        let rhs = buffer![2i32, 5, 1].into_array();
179
180        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Eq).erased();
181        let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?;
182
183        let result = scalar_fn_array
184            .into_array()
185            .execute::<Canonical>(&mut LEGACY_SESSION.create_execution_ctx())?
186            .into_array();
187        let expected = BoolArray::from_iter([false, true, false]).into_array();
188        assert_arrays_eq!(result, expected);
189
190        Ok(())
191    }
192}
vortex_array/arrays/scalar_fn/vtable/operations.rs

vortex_array/arrays/scalar_fn/vtable/
operations.rs
