vortex_array/arrays/chunked/compute/
mask.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use crate::ArrayRef;
7use crate::ExecutionCtx;
8use crate::IntoArray;
9use crate::array::ArrayView;
10use crate::arrays::Chunked;
11use crate::arrays::ChunkedArray;
12use crate::arrays::chunked::ChunkedArrayExt;
13use crate::scalar_fn::fns::mask::Mask as MaskExpr;
14use crate::scalar_fn::fns::mask::MaskKernel;
15
16impl MaskKernel for Chunked {
17    fn mask(
18        array: ArrayView<'_, Chunked>,
19        mask: &ArrayRef,
20        _ctx: &mut ExecutionCtx,
21    ) -> VortexResult<Option<ArrayRef>> {
22        let chunk_offsets = array.chunk_offset_values();
23        let new_chunks: Vec<ArrayRef> = array
24            .iter_chunks()
25            .enumerate()
26            .map(|(i, chunk)| {
27                let start = chunk_offsets[i];
28                let end = chunk_offsets[i + 1];
29                let chunk_mask = mask.slice(start..end)?;
30                MaskExpr::try_new(chunk.clone(), chunk_mask).map(IntoArray::into_array)
31            })
32            .collect::<VortexResult<_>>()?;
33
34        Ok(Some(
35            ChunkedArray::try_new(new_chunks, array.dtype().as_nullable())?.into_array(),
36        ))
37    }
38}
39
40#[cfg(test)]
41mod test {
42    use rstest::rstest;
43    use vortex_buffer::buffer;
44
45    use crate::IntoArray;
46    use crate::VortexSessionExecute;
47    use crate::array_session;
48    use crate::arrays::ChunkedArray;
49    use crate::arrays::PrimitiveArray;
50    use crate::compute::conformance::mask::test_mask_conformance;
51    use crate::dtype::DType;
52    use crate::dtype::Nullability;
53    use crate::dtype::PType;
54
55    #[rstest]
56    #[case(ChunkedArray::try_new(
57        vec![
58            buffer![0u64, 1].into_array(),
59            buffer![2_u64].into_array(),
60            PrimitiveArray::empty::<u64>(Nullability::NonNullable).into_array(),
61            buffer![3_u64, 4].into_array(),
62        ],
63        DType::Primitive(PType::U64, Nullability::NonNullable),
64    ).unwrap())]
65    #[case(ChunkedArray::try_new(
66        vec![
67            PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(),
68            PrimitiveArray::from_option_iter([Some(4i32), Some(5)]).into_array(),
69        ],
70        DType::Primitive(PType::I32, Nullability::Nullable),
71    ).unwrap())]
72    #[case(ChunkedArray::try_new(
73        vec![
74            buffer![42u8].into_array(),
75        ],
76        DType::Primitive(PType::U8, Nullability::NonNullable),
77    ).unwrap())]
78    #[case(ChunkedArray::try_new(
79        (0..20).map(|i| buffer![i as f32, i as f32 + 0.5].into_array()),
80        DType::Primitive(PType::F32, Nullability::NonNullable),
81    ).unwrap())]
82    fn test_mask_chunked_conformance(#[case] chunked: ChunkedArray) {
83        test_mask_conformance(
84            &chunked.into_array(),
85            &mut array_session().create_execution_ctx(),
86        );
87    }
88}
vortex_array/arrays/chunked/compute/mask.rs

vortex_array/arrays/chunked/compute/
mask.rs
