vortex_array/arrays/extension/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexExpect;
5use vortex_error::VortexResult;
6use vortex_error::vortex_ensure_eq;
7
8use crate::ArrayRef;
9use crate::EmptyArrayData;
10use crate::array::Array;
11use crate::array::ArrayParts;
12use crate::array::TypedArrayRef;
13use crate::arrays::Extension;
14use crate::dtype::DType;
15use crate::dtype::extension::ExtDType;
16use crate::dtype::extension::ExtDTypeRef;
17use crate::dtype::extension::ExtVTable;
18
19/// The backing storage array for this extension array.
20pub(super) const STORAGE_SLOT: usize = 0;
21pub(super) const NUM_SLOTS: usize = 1;
22pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["storage"];
23
24pub trait ExtensionArrayExt: TypedArrayRef<Extension> {
25    fn ext_dtype(&self) -> &ExtDTypeRef {
26        self.as_ref()
27            .dtype()
28            .as_extension_opt()
29            .vortex_expect("extension array somehow did not have an extension dtype")
30    }
31
32    fn storage_array(&self) -> &ArrayRef {
33        self.as_ref().slots()[STORAGE_SLOT]
34            .as_ref()
35            .vortex_expect("ExtensionArray storage slot")
36    }
37}
38impl<T: TypedArrayRef<Extension>> ExtensionArrayExt for T {}
39
40impl Array<Extension> {
41    /// Constructs a new `ExtensionArray`.
42    ///
43    /// # Panics
44    ///
45    /// Panics if the storage array is not compatible with the extension dtype.
46    pub fn new(ext_dtype: ExtDTypeRef, storage_array: ArrayRef) -> Self {
47        Self::try_new(ext_dtype, storage_array).vortex_expect("Unable to create `ExtensionArray`")
48    }
49
50    /// Tries to construct a new `ExtensionArray`.
51    pub fn try_new(ext_dtype: ExtDTypeRef, storage_array: ArrayRef) -> VortexResult<Self> {
52        vortex_ensure_eq!(
53            ext_dtype.storage_dtype(),
54            storage_array.dtype(),
55            "Tried to create an `ExtensionArray` with an incompatible storage array"
56        );
57
58        let dtype = DType::Extension(ext_dtype);
59        let len = storage_array.len();
60
61        let parts = ArrayParts::new(Extension, dtype, len, EmptyArrayData)
62            .with_slots(vec![Some(storage_array)]);
63
64        Ok(unsafe { Array::from_parts_unchecked(parts) })
65    }
66
67    /// Creates a new [`ExtensionArray`](crate::arrays::ExtensionArray) from a vtable, metadata, and
68    /// a storage array.
69    pub fn try_new_from_vtable<V: ExtVTable>(
70        vtable: V,
71        metadata: V::Metadata,
72        storage_array: ArrayRef,
73    ) -> VortexResult<Self> {
74        let ext_dtype =
75            ExtDType::<V>::try_with_vtable(vtable, metadata, storage_array.dtype().clone())?
76                .erased();
77
78        Self::try_new(ext_dtype, storage_array)
79    }
80}
vortex_array/arrays/extension/array.rs

vortex_array/arrays/extension/
array.rs
