vortex_array/arrays/extension/compute/
mask.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use vortex_dtype::ExtDType;
7use vortex_error::VortexResult;
8use vortex_mask::Mask;
9
10use crate::arrays::{ExtensionArray, ExtensionVTable};
11use crate::compute::{self, MaskKernel, MaskKernelAdapter};
12use crate::{ArrayRef, IntoArray, register_kernel};
13
14impl MaskKernel for ExtensionVTable {
15    fn mask(&self, array: &ExtensionArray, mask_array: &Mask) -> VortexResult<ArrayRef> {
16        let masked_storage = compute::mask(array.storage(), mask_array)?;
17        if masked_storage.dtype().nullability() == array.ext_dtype().storage_dtype().nullability() {
18            Ok(ExtensionArray::new(array.ext_dtype().clone(), masked_storage).into_array())
19        } else {
20            // The storage dtype changed (i.e., became nullable due to masking)
21            let ext_dtype = Arc::new(ExtDType::new(
22                array.ext_dtype().id().clone(),
23                Arc::new(masked_storage.dtype().clone()),
24                array.ext_dtype().metadata().cloned(),
25            ));
26            Ok(ExtensionArray::new(ext_dtype, masked_storage).into_array())
27        }
28    }
29}
30
31register_kernel!(MaskKernelAdapter(ExtensionVTable).lift());