vortex_array/arrays/chunked/vtable/
compute.rs

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

vortex_array/arrays/chunked/vtable/
compute.rs
