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::ProbeState;
9use crate::array::ArrayView;
10use crate::array::OperationsVTable;
11use crate::arrays::ConstantArray;
12use crate::arrays::scalar_fn::ScalarFnArrayExt;
13use crate::arrays::scalar_fn::vtable::ScalarFn;
14use crate::columnar::Columnar;
15use crate::scalar::Scalar;
16use crate::scalar_fn::VecExecutionArgs;
17
18impl OperationsVTable<ScalarFn> for ScalarFn {
19    type ProbeState = ();
20
21    /// The evaluated scalar carries the row's nullness, so no separate validity check is needed.
22    fn probe_scalar(
23        state: &mut ProbeState<'_, ScalarFn>,
24        index: usize,
25        ctx: &mut ExecutionCtx,
26    ) -> VortexResult<Scalar> {
27        Self::scalar_at(state.array(), index, ctx)
28    }
29
30    fn scalar_at(
31        array: ArrayView<'_, ScalarFn>,
32        index: usize,
33        ctx: &mut ExecutionCtx,
34    ) -> VortexResult<Scalar> {
35        let inputs: Vec<_> = array
36            .children()
37            .iter()
38            .map(|child| Ok(ConstantArray::new(child.execute_scalar(index, ctx)?, 1).into_array()))
39            .collect::<VortexResult<_>>()?;
40
41        let args = VecExecutionArgs::new(inputs, 1);
42        let result = array.scalar_fn().execute(&args, ctx)?;
43
44        let scalar = match result.execute::<Columnar>(ctx)? {
45            Columnar::Canonical(arr) => {
46                tracing::info!(
47                    "Scalar function {} returned non-constant array from execution over all scalar inputs",
48                    array.scalar_fn(),
49                );
50                arr.into_array().execute_scalar(0, ctx)?
51            }
52            Columnar::Constant(constant) => constant.scalar().clone(),
53        };
54
55        debug_assert_eq!(
56            scalar.dtype(),
57            array.dtype(),
58            "Scalar function {} returned dtype {:?} but expected {:?}",
59            array.scalar_fn(),
60            scalar.dtype(),
61            array.dtype()
62        );
63
64        Ok(scalar)
65    }
66}
67
68#[cfg(test)]
69mod tests {
70    use vortex_buffer::buffer;
71    use vortex_error::VortexResult;
72
73    use crate::ArrayRef;
74    use crate::ArraySlots;
75    use crate::Canonical;
76    use crate::IntoArray;
77    use crate::VortexSessionExecute;
78    use crate::array::Array;
79    use crate::array::ArrayParts;
80    use crate::array_session;
81    use crate::arrays::BoolArray;
82    use crate::arrays::PrimitiveArray;
83    use crate::arrays::ScalarFnArray;
84    use crate::arrays::scalar_fn::ScalarFnArrayExt;
85    use crate::arrays::scalar_fn::array::ScalarFnData;
86    use crate::arrays::scalar_fn::vtable::ScalarFn;
87    use crate::assert_arrays_eq;
88    use crate::scalar::Scalar;
89    use crate::scalar_fn::TypedScalarFnInstance;
90    use crate::scalar_fn::fns::binary::Binary;
91    use crate::scalar_fn::fns::literal::Literal;
92    use crate::scalar_fn::fns::operators::Operator;
93    use crate::validity::Validity;
94
95    #[test]
96    fn test_scalar_fn_add() -> VortexResult<()> {
97        let mut ctx = array_session().create_execution_ctx();
98        let lhs = buffer![1i32, 2, 3].into_array();
99        let rhs = buffer![10i32, 20, 30].into_array();
100
101        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Add).erased();
102        let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?;
103
104        assert_eq!(scalar_fn_array.len(), 3);
105
106        let result = scalar_fn_array
107            .into_array()
108            .execute::<Canonical>(&mut array_session().create_execution_ctx())?
109            .into_array();
110        let expected = buffer![11i32, 22, 33].into_array();
111        assert_arrays_eq!(result, expected, &mut ctx);
112
113        Ok(())
114    }
115
116    #[test]
117    fn test_scalar_fn_inferred_len_rejects_mismatched_children() {
118        let lhs = buffer![1i32, 2, 3].into_array();
119        let rhs = buffer![10i32, 20].into_array();
120
121        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Add).erased();
122        let err = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])
123            .expect_err("ScalarFnArray::try_new must reject mismatched child lengths");
124
125        assert!(
126            err.to_string()
127                .contains("ScalarFnArray must have children equal to the array length")
128        );
129    }
130
131    #[test]
132    fn rejects_wrong_arity() {
133        let child = buffer![1i32, 2, 3].into_array();
134        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Add).erased();
135
136        let Err(err) = ScalarFnArray::try_new(scalar_fn.clone(), vec![child.clone()]) else {
137            panic!("Binary must reject one child");
138        };
139        assert!(
140            err.to_string()
141                .contains("ScalarFnArray requires 2 children, got 1")
142        );
143
144        let Err(err) = ScalarFnArray::try_new(scalar_fn, vec![child.clone(), child.clone(), child])
145        else {
146            panic!("Binary must reject three children");
147        };
148        assert!(
149            err.to_string()
150                .contains("ScalarFnArray requires 2 children, got 3")
151        );
152    }
153
154    #[test]
155    fn rejects_missing_child_slot() {
156        let lhs = buffer![1i32, 2, 3].into_array();
157        let dtype = lhs.dtype().clone();
158        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Add).erased();
159        let vtable = ScalarFn { id: scalar_fn.id() };
160        let data = ScalarFnData { scalar_fn };
161        let slots = [Some(lhs), None].into_iter().collect::<ArraySlots>();
162
163        let Err(err) =
164            Array::<ScalarFn>::try_from_parts(ArrayParts::new(vtable, dtype, 3, data, slots))
165        else {
166            panic!("ScalarFnArray must reject missing child slots");
167        };
168
169        assert!(
170            err.to_string()
171                .contains("ScalarFnArray requires every child slot to be present")
172        );
173    }
174
175    #[test]
176    fn test_scalar_fn_without_children_requires_explicit_len() -> VortexResult<()> {
177        let scalar_fn = TypedScalarFnInstance::new(Literal, Scalar::from(1i32)).erased();
178
179        let Err(err) = ScalarFnArray::try_new(scalar_fn.clone(), vec![]) else {
180            panic!("ScalarFnArray::try_new should reject zero children");
181        };
182        assert!(
183            err.to_string()
184                .contains("ScalarFnArray length cannot be inferred without children")
185        );
186
187        let scalar_fn_array = ScalarFnArray::try_new_with_len(scalar_fn, vec![], 3)?;
188        assert_eq!(scalar_fn_array.len(), 3);
189        assert_eq!(scalar_fn_array.child_count(), 0);
190
191        Ok(())
192    }
193
194    #[test]
195    fn test_scalar_fn_mul() -> VortexResult<()> {
196        let mut ctx = array_session().create_execution_ctx();
197        let lhs = buffer![2i32, 3, 4].into_array();
198        let rhs = buffer![5i32, 6, 7].into_array();
199
200        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Mul).erased();
201        let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?;
202
203        let result = scalar_fn_array
204            .into_array()
205            .execute::<Canonical>(&mut array_session().create_execution_ctx())?
206            .into_array();
207        let expected = buffer![10i32, 18, 28].into_array();
208        assert_arrays_eq!(result, expected, &mut ctx);
209
210        Ok(())
211    }
212
213    #[test]
214    fn test_scalar_fn_with_nullable() -> VortexResult<()> {
215        let mut ctx = array_session().create_execution_ctx();
216        let lhs = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllValid).into_array();
217        let rhs = PrimitiveArray::new(
218            buffer![10i32, 20, 30],
219            Validity::from_iter([true, false, true]),
220        )
221        .into_array();
222
223        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Add).erased();
224        let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?;
225
226        let result = scalar_fn_array
227            .into_array()
228            .execute::<Canonical>(&mut array_session().create_execution_ctx())?
229            .into_array();
230        let expected = PrimitiveArray::new(
231            buffer![11i32, 0, 33],
232            Validity::from_iter([true, false, true]),
233        )
234        .into_array();
235        assert_arrays_eq!(result, expected, &mut ctx);
236
237        Ok(())
238    }
239
240    #[test]
241    fn test_scalar_fn_comparison() -> VortexResult<()> {
242        let mut ctx = array_session().create_execution_ctx();
243        let lhs = buffer![1i32, 5, 3].into_array();
244        let rhs = buffer![2i32, 5, 1].into_array();
245
246        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Eq).erased();
247        let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?;
248
249        let result = scalar_fn_array
250            .into_array()
251            .execute::<Canonical>(&mut array_session().create_execution_ctx())?
252            .into_array();
253        let expected = BoolArray::from_iter([false, true, false]).into_array();
254        assert_arrays_eq!(result, expected, &mut ctx);
255
256        Ok(())
257    }
258
259    /// Mul(Add(a1, a2), a3)
260    fn scalar_nested_add() -> VortexResult<ArrayRef> {
261        let lhs_n = PrimitiveArray::from_option_iter((0..10i64).map(|i| (i % 4 != 0).then_some(i)));
262        let rhs_r = PrimitiveArray::from_iter(0..10i64);
263
264        let scalar_fn_n = TypedScalarFnInstance::new(Binary, Operator::Add).erased();
265        let args_n = vec![lhs_n.into_array(), rhs_r.into_array()];
266
267        let lhs = ScalarFnArray::try_new(scalar_fn_n, args_n)?.into_array();
268        let rhs = PrimitiveArray::from_iter(0..10i64);
269
270        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Mul).erased();
271        let args = vec![lhs, rhs.into_array()];
272
273        Ok(ScalarFnArray::try_new(scalar_fn, args)?.into_array())
274    }
275
276    #[test]
277    fn scalar_fn_probe() -> VortexResult<()> {
278        let ctx = &mut array_session().create_execution_ctx();
279        let array = scalar_nested_add()?;
280        let mut probe = array.probe();
281
282        assert!(!probe.execute_is_valid(0, ctx)?);
283        assert!(probe.execute_scalar(0, ctx)?.is_null());
284        assert!(probe.execute_is_valid(9, ctx)?);
285        assert_eq!(probe.execute_scalar(9, ctx)?, Scalar::from(Some(162i64)));
286
287        Ok(())
288    }
289
290    #[test]
291    fn scalar_fn_repeated_probe() -> VortexResult<()> {
292        let ctx = &mut array_session().create_execution_ctx();
293        let array = scalar_nested_add()?;
294        let mut probe = array.repeated_probe();
295
296        assert!(!probe.execute_is_valid(0, ctx)?);
297        assert!(probe.execute_scalar(0, ctx)?.is_null());
298        assert!(probe.execute_is_valid(9, ctx)?);
299        assert_eq!(probe.execute_scalar(9, ctx)?, Scalar::from(Some(162i64)));
300
301        Ok(())
302    }
303
304    #[test]
305    fn scalar_fn_all_valid() -> VortexResult<()> {
306        let ctx = &mut array_session().create_execution_ctx();
307        let array = scalar_nested_add()?;
308        assert!(!array.all_valid(ctx)?);
309        assert!(!array.all_invalid(ctx)?);
310        assert_eq!(array.valid_count(ctx)?, 7);
311        assert_eq!(array.invalid_count(ctx)?, 3);
312        Ok(())
313    }
314}
vortex_array/arrays/scalar_fn/vtable/operations.rs

vortex_array/arrays/scalar_fn/vtable/
operations.rs
