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    pub storage: ArrayRef,
24}
25
26pub trait ExtensionArrayExt: TypedArrayRef<Extension> + ExtensionArraySlotsExt {
27    fn ext_dtype(&self) -> &ExtDTypeRef {
28        self.as_ref()
29            .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(Extension, dtype, len, EmptyArrayData).with_slots(
63            ExtensionSlots {
64                storage: storage_array,
65            }
66            .into_slots(),
67        );
68
69        Ok(unsafe { Array::from_parts_unchecked(parts) })
70    }
71
72    /// Creates a new [`ExtensionArray`](crate::arrays::ExtensionArray) from a vtable, metadata, and
73    /// a storage array.
74    pub fn try_new_from_vtable<V: ExtVTable>(
75        vtable: V,
76        metadata: V::Metadata,
77        storage_array: ArrayRef,
78    ) -> VortexResult<Self> {
79        let ext_dtype =
80            ExtDType::<V>::try_with_vtable(vtable, metadata, storage_array.dtype().clone())?
81                .erased();
82
83        Self::try_new(ext_dtype, storage_array)
84    }
85}
vortex_array/arrays/extension/array.rs

vortex_array/arrays/extension/
array.rs
