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(
64            Extension,
65            dtype,
66            len,
67            EmptyArrayData,
68            ExtensionSlots {
69                storage: storage_array,
70            }
71            .into_slots(),
72        );
73
74        Ok(unsafe { Array::from_parts_unchecked(parts) })
75    }
76
77    /// Creates a new [`ExtensionArray`](crate::arrays::ExtensionArray) from a vtable, metadata, and
78    /// a storage array.
79    pub fn try_new_from_vtable<V: ExtVTable>(
80        vtable: V,
81        metadata: V::Metadata,
82        storage_array: ArrayRef,
83    ) -> VortexResult<Self> {
84        let ext_dtype =
85            ExtDType::<V>::try_with_vtable(vtable, metadata, storage_array.dtype().clone())?
86                .erased();
87
88        Self::try_new(ext_dtype, storage_array)
89    }
90}
vortex_array/arrays/extension/array.rs

vortex_array/arrays/extension/
array.rs
