vortex_array/arrays/extension/compute/
take.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;
8
9use crate::arrays::{ExtensionArray, ExtensionVTable};
10use crate::compute::{self, TakeKernel, TakeKernelAdapter};
11use crate::{Array, ArrayRef, IntoArray, register_kernel};
12
13impl TakeKernel for ExtensionVTable {
14    fn take(&self, array: &ExtensionArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
15        let taken_storage = compute::take(array.storage(), indices)?;
16        if taken_storage.dtype().nullability() == array.ext_dtype().storage_dtype().nullability() {
17            Ok(ExtensionArray::new(array.ext_dtype().clone(), taken_storage).into_array())
18        } else {
19            // The storage dtype changed (i.e., became nullable due to nullable indices)
20            let ext_dtype = Arc::new(ExtDType::new(
21                array.ext_dtype().id().clone(),
22                Arc::new(taken_storage.dtype().clone()),
23                array.ext_dtype().metadata().cloned(),
24            ));
25            Ok(ExtensionArray::new(ext_dtype, taken_storage).into_array())
26        }
27    }
28}
29
30register_kernel!(TakeKernelAdapter(ExtensionVTable).lift());