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