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

vortex_array/arrays/extension/
array.rs
