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::IntoArray;
76    use crate::VortexSessionExecute;
77    use crate::array::Array;
78    use crate::array::ArrayParts;
79    use crate::array_session;
80    use crate::arrays::BoolArray;
81    use crate::arrays::ConstantArray;
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::dtype::DType;
89    use crate::dtype::Nullability;
90    use crate::dtype::PType;
91    use crate::scalar::Scalar;
92    use crate::scalar_fn::TypedScalarFnInstance;
93    use crate::scalar_fn::fns::binary::Binary;
94    use crate::scalar_fn::fns::list_contains::PreparedSet;
95    use crate::scalar_fn::fns::list_contains::PreparedSetData;
96    use crate::scalar_fn::fns::literal::Literal;
97    use crate::scalar_fn::fns::operators::Operator;
98    use crate::validity::Validity;
99
100    #[test]
101    fn test_scalar_fn_add() -> VortexResult<()> {
102        let mut ctx = array_session().create_execution_ctx();
103        let lhs = buffer![1i32, 2, 3].into_array();
104        let rhs = buffer![10i32, 20, 30].into_array();
105
106        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Add).erased();
107        let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?;
108
109        assert_eq!(scalar_fn_array.len(), 3);
110
111        let expected = buffer![11i32, 22, 33].into_array();
112        assert_arrays_eq!(scalar_fn_array, expected, &mut ctx);
113
114        Ok(())
115    }
116
117    #[test]
118    fn test_scalar_fn_inferred_len_rejects_mismatched_children() {
119        let lhs = buffer![1i32, 2, 3].into_array();
120        let rhs = buffer![10i32, 20].into_array();
121
122        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Add).erased();
123        let err = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])
124            .expect_err("ScalarFnArray::try_new must reject mismatched child lengths");
125
126        assert!(
127            err.to_string()
128                .contains("ScalarFnArray must have children equal to the array length")
129        );
130    }
131
132    #[test]
133    fn rejects_wrong_arity() {
134        let child = buffer![1i32, 2, 3].into_array();
135        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Add).erased();
136
137        let Err(err) = ScalarFnArray::try_new(scalar_fn.clone(), vec![child.clone()]) else {
138            panic!("Binary must reject one child");
139        };
140        assert!(
141            err.to_string()
142                .contains("ScalarFnArray requires 2 children, got 1")
143        );
144
145        let Err(err) = ScalarFnArray::try_new(scalar_fn, vec![child.clone(), child.clone(), child])
146        else {
147            panic!("Binary must reject three children");
148        };
149        assert!(
150            err.to_string()
151                .contains("ScalarFnArray requires 2 children, got 3")
152        );
153    }
154
155    #[test]
156    fn rejects_missing_child_slot() {
157        let lhs = buffer![1i32, 2, 3].into_array();
158        let dtype = lhs.dtype().clone();
159        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Add).erased();
160        let vtable = ScalarFn { id: scalar_fn.id() };
161        let data = ScalarFnData { scalar_fn };
162        let slots = [Some(lhs), None].into_iter().collect::<ArraySlots>();
163
164        let Err(err) =
165            Array::<ScalarFn>::try_from_parts(ArrayParts::new(vtable, dtype, 3, data, slots))
166        else {
167            panic!("ScalarFnArray must reject missing child slots");
168        };
169
170        assert!(
171            err.to_string()
172                .contains("ScalarFnArray requires every child slot to be present")
173        );
174    }
175
176    #[test]
177    fn test_scalar_fn_without_children_requires_explicit_len() -> VortexResult<()> {
178        let list = Scalar::list(
179            DType::Primitive(PType::I32, Nullability::NonNullable),
180            vec![Scalar::from(1i32)],
181            Nullability::NonNullable,
182        );
183        let scalar_fn =
184            TypedScalarFnInstance::new(PreparedSet, PreparedSetData::try_new(list.clone())?)
185                .erased();
186
187        let Err(err) = ScalarFnArray::try_new(scalar_fn.clone(), vec![]) else {
188            panic!("ScalarFnArray::try_new should reject zero children");
189        };
190        assert!(
191            err.to_string()
192                .contains("ScalarFnArray length cannot be inferred without children")
193        );
194
195        let scalar_fn_array = ScalarFnArray::try_new_with_len(scalar_fn, vec![], 3)?;
196        assert_eq!(scalar_fn_array.len(), 3);
197        assert_eq!(scalar_fn_array.child_count(), 0);
198        assert_arrays_eq!(
199            scalar_fn_array,
200            ConstantArray::new(list, 3),
201            &mut array_session().create_execution_ctx()
202        );
203
204        Ok(())
205    }
206
207    #[test]
208    fn rejects_literal() {
209        let scalar_fn = TypedScalarFnInstance::new(Literal, Scalar::from(1i32)).erased();
210
211        let Err(err) = ScalarFnArray::try_new_with_len(scalar_fn.clone(), vec![], 3) else {
212            panic!("ScalarFnArray::try_new_with_len must reject a literal");
213        };
214        assert!(
215            err.to_string()
216                .contains("ScalarFnArray cannot wrap a literal")
217        );
218
219        let vtable = ScalarFn { id: scalar_fn.id() };
220        let data = ScalarFnData { scalar_fn };
221        let parts = ArrayParts::new(vtable, PType::I32.into(), 3, data, ArraySlots::new());
222
223        let Err(err) = Array::<ScalarFn>::try_from_parts(parts) else {
224            panic!("ScalarFnArray must reject a literal from parts");
225        };
226        assert!(
227            err.to_string()
228                .contains("ScalarFnArray cannot wrap a literal")
229        );
230    }
231
232    #[test]
233    fn test_scalar_fn_mul() -> VortexResult<()> {
234        let mut ctx = array_session().create_execution_ctx();
235        let lhs = buffer![2i32, 3, 4].into_array();
236        let rhs = buffer![5i32, 6, 7].into_array();
237
238        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Mul).erased();
239        let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?;
240
241        let expected = buffer![10i32, 18, 28].into_array();
242        assert_arrays_eq!(scalar_fn_array, expected, &mut ctx);
243
244        Ok(())
245    }
246
247    #[test]
248    fn test_scalar_fn_with_nullable() -> VortexResult<()> {
249        let mut ctx = array_session().create_execution_ctx();
250        let lhs = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllValid).into_array();
251        let rhs = PrimitiveArray::new(
252            buffer![10i32, 20, 30],
253            Validity::from_iter([true, false, true]),
254        )
255        .into_array();
256
257        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Add).erased();
258        let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?;
259
260        let expected = PrimitiveArray::new(
261            buffer![11i32, 0, 33],
262            Validity::from_iter([true, false, true]),
263        )
264        .into_array();
265        assert_arrays_eq!(scalar_fn_array, expected, &mut ctx);
266
267        Ok(())
268    }
269
270    #[test]
271    fn test_scalar_fn_comparison() -> VortexResult<()> {
272        let mut ctx = array_session().create_execution_ctx();
273        let lhs = buffer![1i32, 5, 3].into_array();
274        let rhs = buffer![2i32, 5, 1].into_array();
275
276        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Eq).erased();
277        let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?;
278
279        let expected = BoolArray::from_iter([false, true, false]).into_array();
280        assert_arrays_eq!(scalar_fn_array, expected, &mut ctx);
281
282        Ok(())
283    }
284
285    /// Mul(Add(a1, a2), a3)
286    fn scalar_nested_add() -> VortexResult<ArrayRef> {
287        let lhs_n = PrimitiveArray::from_option_iter((0..10i64).map(|i| (i % 4 != 0).then_some(i)));
288        let rhs_r = PrimitiveArray::from_iter(0..10i64);
289
290        let scalar_fn_n = TypedScalarFnInstance::new(Binary, Operator::Add).erased();
291        let args_n = vec![lhs_n.into_array(), rhs_r.into_array()];
292
293        let lhs = ScalarFnArray::try_new(scalar_fn_n, args_n)?.into_array();
294        let rhs = PrimitiveArray::from_iter(0..10i64);
295
296        let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Mul).erased();
297        let args = vec![lhs, rhs.into_array()];
298
299        Ok(ScalarFnArray::try_new(scalar_fn, args)?.into_array())
300    }
301
302    #[test]
303    fn scalar_fn_probe() -> VortexResult<()> {
304        let ctx = &mut array_session().create_execution_ctx();
305        let array = scalar_nested_add()?;
306        let mut probe = array.probe();
307
308        assert!(!probe.execute_is_valid(0, ctx)?);
309        assert!(probe.execute_scalar(0, ctx)?.is_null());
310        assert!(probe.execute_is_valid(9, ctx)?);
311        assert_eq!(probe.execute_scalar(9, ctx)?, Scalar::from(Some(162i64)));
312
313        Ok(())
314    }
315
316    #[test]
317    fn scalar_fn_repeated_probe() -> VortexResult<()> {
318        let ctx = &mut array_session().create_execution_ctx();
319        let array = scalar_nested_add()?;
320        let mut probe = array.repeated_probe();
321
322        assert!(!probe.execute_is_valid(0, ctx)?);
323        assert!(probe.execute_scalar(0, ctx)?.is_null());
324        assert!(probe.execute_is_valid(9, ctx)?);
325        assert_eq!(probe.execute_scalar(9, ctx)?, Scalar::from(Some(162i64)));
326
327        Ok(())
328    }
329
330    #[test]
331    fn scalar_fn_all_valid() -> VortexResult<()> {
332        let ctx = &mut array_session().create_execution_ctx();
333        let array = scalar_nested_add()?;
334        assert!(!array.all_valid(ctx)?);
335        assert!(!array.all_invalid(ctx)?);
336        assert_eq!(array.valid_count(ctx)?, 7);
337        assert_eq!(array.invalid_count(ctx)?, 3);
338        Ok(())
339    }
340}
vortex_array/arrays/scalar_fn/vtable/operations.rs

vortex_array/arrays/scalar_fn/vtable/
operations.rs
