vortex_array/arrays/extension/
array.rs

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

vortex_array/arrays/extension/
array.rs
