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

vortex_array/arrays/chunked/compute/
mask.rs
