vortex_array/array/chunked/compute/
binary_numeric.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
use vortex_error::VortexResult;
use vortex_scalar::BinaryNumericOperator;

use crate::array::{ChunkedArray, ChunkedEncoding};
use crate::compute::{binary_numeric, slice, BinaryNumericFn};
use crate::{ArrayDType as _, ArrayData, IntoArrayData};

impl BinaryNumericFn<ChunkedArray> for ChunkedEncoding {
    fn binary_numeric(
        &self,
        array: &ChunkedArray,
        rhs: &ArrayData,
        op: BinaryNumericOperator,
    ) -> VortexResult<Option<ArrayData>> {
        let mut start = 0;

        let mut new_chunks = Vec::with_capacity(array.nchunks());
        for chunk in array.chunks() {
            let end = start + chunk.len();
            new_chunks.push(binary_numeric(&chunk, &slice(rhs, start, end)?, op)?);
            start = end;
        }

        ChunkedArray::try_new(new_chunks, array.dtype().clone())
            .map(IntoArrayData::into_array)
            .map(Some)
    }
}