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::array_slots;
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#[array_slots(Extension)]
21pub struct ExtensionSlots {
22    /// The backing storage array for this extension array.
23    #[slot(0)]
24    pub storage: ArrayRef,
25}
26
27pub trait ExtensionArrayExt: TypedArrayRef<Extension> + ExtensionArraySlotsExt {
28    fn ext_dtype(&self) -> &ExtDTypeRef {
29        self.as_ref()
30            .dtype()
31            .as_extension_opt()
32            .vortex_expect("extension array somehow did not have an extension dtype")
33    }
34
35    /// Returns the backing storage array.
36    fn storage_array(&self) -> &ArrayRef {
37        self.storage()
38    }
39}
40impl<T: TypedArrayRef<Extension>> ExtensionArrayExt for T {}
41
42impl Array<Extension> {
43    /// Constructs a new `ExtensionArray`.
44    ///
45    /// # Panics
46    ///
47    /// Panics if the storage array is not compatible with the extension dtype.
48    pub fn new(ext_dtype: ExtDTypeRef, storage_array: ArrayRef) -> Self {
49        Self::try_new(ext_dtype, storage_array).vortex_expect("Unable to create `ExtensionArray`")
50    }
51
52    /// Tries to construct a new `ExtensionArray`.
53    pub fn try_new(ext_dtype: ExtDTypeRef, storage_array: ArrayRef) -> VortexResult<Self> {
54        vortex_ensure_eq!(
55            ext_dtype.storage_dtype(),
56            storage_array.dtype(),
57            "Tried to create an `ExtensionArray` with an incompatible storage array"
58        );
59
60        let dtype = DType::Extension(ext_dtype);
61        let len = storage_array.len();
62
63        let parts = ArrayParts::new(Extension, dtype, len, EmptyArrayData).with_slots(
64            ExtensionSlots {
65                storage: storage_array,
66            }
67            .into_slots(),
68        );
69
70        Ok(unsafe { Array::from_parts_unchecked(parts) })
71    }
72
73    /// Creates a new [`ExtensionArray`](crate::arrays::ExtensionArray) from a vtable, metadata, and
74    /// a storage array.
75    pub fn try_new_from_vtable<V: ExtVTable>(
76        vtable: V,
77        metadata: V::Metadata,
78        storage_array: ArrayRef,
79    ) -> VortexResult<Self> {
80        let ext_dtype =
81            ExtDType::<V>::try_with_vtable(vtable, metadata, storage_array.dtype().clone())?
82                .erased();
83
84        Self::try_new(ext_dtype, storage_array)
85    }
86}
vortex_array/arrays/extension/array.rs

vortex_array/arrays/extension/
array.rs
