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::ArrayRef;
11use crate::IntoArray;
12use crate::arrays::ExtensionArray;
13use crate::arrays::ExtensionVTable;
14use crate::compute::MaskKernel;
15use crate::compute::MaskKernelAdapter;
16use crate::compute::{self};
17use crate::register_kernel;
18
19impl MaskKernel for ExtensionVTable {
20    fn mask(&self, array: &ExtensionArray, mask_array: &Mask) -> VortexResult<ArrayRef> {
21        let masked_storage = compute::mask(array.storage(), mask_array)?;
22        if masked_storage.dtype().nullability() == array.ext_dtype().storage_dtype().nullability() {
23            Ok(ExtensionArray::new(array.ext_dtype().clone(), masked_storage).into_array())
24        } else {
25            // The storage dtype changed (i.e., became nullable due to masking)
26            let ext_dtype = Arc::new(ExtDType::new(
27                array.ext_dtype().id().clone(),
28                Arc::new(masked_storage.dtype().clone()),
29                array.ext_dtype().metadata().cloned(),
30            ));
31            Ok(ExtensionArray::new(ext_dtype, masked_storage).into_array())
32        }
33    }
34}
35
36register_kernel!(MaskKernelAdapter(ExtensionVTable).lift());