vortex_array/arrays/chunked/compute/
elementwise.rs

1use vortex_error::{VortexExpect, VortexResult};
2
3use crate::arrays::{ChunkedArray, ChunkedVTable};
4use crate::compute::{ComputeFn, InvocationArgs, Output};
5use crate::vtable::ComputeVTable;
6use crate::{Array, IntoArray};
7
8impl ComputeVTable<ChunkedVTable> for ChunkedVTable {
9    fn invoke(
10        array: &ChunkedArray,
11        compute_fn: &ComputeFn,
12        args: &InvocationArgs,
13    ) -> VortexResult<Option<Output>> {
14        if compute_fn.is_elementwise() {
15            return invoke_elementwise(array, compute_fn, args);
16        }
17        Ok(None)
18    }
19}
20
21/// Invoke an element-wise compute function over a chunked array.
22fn invoke_elementwise(
23    array: &ChunkedArray,
24    compute_fn: &ComputeFn,
25    args: &InvocationArgs,
26) -> VortexResult<Option<Output>> {
27    assert!(
28        compute_fn.is_elementwise(),
29        "Expected elementwise compute function"
30    );
31    assert!(
32        !args.inputs.is_empty(),
33        "Elementwise compute function requires at least one input"
34    );
35
36    // If not all inputs are arrays, then we pass.
37    if args.inputs.iter().any(|a| a.array().is_none()) {
38        return Ok(None);
39    }
40
41    let mut idx = 0;
42    let mut chunks = Vec::with_capacity(array.nchunks());
43    let mut inputs = Vec::with_capacity(args.inputs.len());
44
45    for chunk in array.non_empty_chunks() {
46        inputs.clear();
47        inputs.push(chunk.clone());
48        for i in 1..args.inputs.len() {
49            let input = args.inputs[i].array().vortex_expect("checked already");
50            let sliced = input.slice(idx, idx + chunk.len())?;
51            inputs.push(sliced);
52        }
53
54        // TODO(ngates): we might want to make invocation args not hold references?
55        let input_refs = inputs.iter().map(|a| a.as_ref().into()).collect::<Vec<_>>();
56
57        // Delegate the compute kernel to the chunk.
58        let result = compute_fn
59            .invoke(&InvocationArgs {
60                inputs: &input_refs,
61                options: args.options,
62            })?
63            .unwrap_array()?;
64
65        chunks.push(result);
66        idx += chunk.len();
67    }
68
69    let return_dtype = compute_fn.return_dtype(args)?;
70    Ok(Some(
71        ChunkedArray::try_new(chunks, return_dtype)?
72            .into_array()
73            .into(),
74    ))
75}
76
77#[cfg(test)]
78mod tests {
79    use vortex_dtype::{DType, Nullability};
80
81    use crate::arrays::{BoolArray, BooleanBuffer, ChunkedArray};
82    use crate::canonical::ToCanonical;
83    use crate::compute::{BooleanOperator, boolean};
84
85    #[test]
86    fn test_bin_bool_chunked() {
87        let arr0 = BoolArray::from_iter(vec![true, false]).to_array();
88        let arr1 = BoolArray::from_iter(vec![false, false, true]).to_array();
89        let chunked1 =
90            ChunkedArray::try_new(vec![arr0, arr1], DType::Bool(Nullability::NonNullable)).unwrap();
91
92        let arr2 = BoolArray::from_iter(vec![Some(false), Some(true)]).to_array();
93        let arr3 = BoolArray::from_iter(vec![Some(false), None, Some(false)]).to_array();
94        let chunked2 =
95            ChunkedArray::try_new(vec![arr2, arr3], DType::Bool(Nullability::Nullable)).unwrap();
96
97        let result = boolean(chunked1.as_ref(), chunked2.as_ref(), BooleanOperator::Or)
98            .unwrap()
99            .to_bool()
100            .unwrap();
101        assert_eq!(
102            result.boolean_buffer(),
103            &BooleanBuffer::from_iter([true, true, false, false, true])
104        );
105    }
106}
vortex_array/arrays/chunked/compute/elementwise.rs

vortex_array/arrays/chunked/compute/
elementwise.rs
