vortex_array/arrays/extension/vtable/
mod.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_bail;
7use vortex_error::vortex_ensure_eq;
8use vortex_error::vortex_err;
9use vortex_error::vortex_panic;
10use vortex_session::VortexSession;
11use vortex_session::registry::CachedId;
12
13use crate::ArrayRef;
14use crate::EmptyArrayData;
15use crate::ExecutionCtx;
16use crate::ExecutionResult;
17use crate::array::Array;
18use crate::array::ArrayId;
19use crate::array::ArrayParts;
20use crate::array::ArrayView;
21use crate::array::VTable;
22use crate::array::ValidityVTableFromChild;
23use crate::array::with_empty_buffers;
24use crate::arrays::extension::array::ExtensionSlots;
25use crate::arrays::extension::compute::rules::PARENT_RULES;
26use crate::arrays::extension::compute::rules::RULES;
27use crate::buffer::BufferHandle;
28use crate::builders::ArrayBuilder;
29use crate::builders::ExtensionBuilder;
30use crate::dtype::DType;
31use crate::serde::ArrayChildren;
32
33mod kernel;
34mod operations;
35mod validity;
36
37/// An extension array that wraps another array with additional type information.
38///
39/// **⚠️ Unstable API**: This is an experimental feature that may change significantly
40/// in future versions. The extension type system is still evolving.
41///
42/// Unlike Apache Arrow's extension arrays, Vortex extension arrays provide a more flexible
43/// mechanism for adding semantic meaning to existing array types without requiring
44/// changes to the core type system.
45///
46/// ## Design Philosophy
47///
48/// Extension arrays serve as a type-safe wrapper that:
49/// - Preserves the underlying storage format and operations
50/// - Adds semantic type information via `ExtDType`
51/// - Enables custom serialization and deserialization logic
52/// - Allows domain-specific interpretations of generic data
53///
54/// ## Storage and Type Relationship
55///
56/// The extension array maintains a strict contract:
57/// - **Storage array**: Contains the actual data in a standard Vortex encoding
58/// - **Extension type**: Defines how to interpret the storage data semantically
59/// - **Type safety**: The storage array's dtype must match the extension type's storage dtype
60///
61/// ## Use Cases
62///
63/// Extension arrays are ideal for:
64/// - **Custom numeric types**: Units of measurement, currencies
65/// - **Temporal types**: Custom date/time formats, time zones, calendars
66/// - **Domain-specific types**: UUIDs, IP addresses, geographic coordinates
67/// - **Encoded types**: Base64 strings, compressed data, encrypted values
68///
69/// ## Validity and Operations
70///
71/// Extension arrays delegate validity and most operations to their storage array:
72/// - Validity is inherited from the underlying storage
73/// - Slicing preserves the extension type
74/// - Scalar access wraps storage scalars with extension metadata
75#[derive(Clone, Debug)]
76pub struct Extension;
77
78/// A [`Extension`]-encoded Vortex array.
79pub type ExtensionArray = Array<Extension>;
80
81pub(crate) fn initialize(session: &VortexSession) {
82    kernel::initialize(session);
83}
84
85impl VTable for Extension {
86    type TypedArrayData = EmptyArrayData;
87
88    type OperationsVTable = Self;
89    type ValidityVTable = ValidityVTableFromChild;
90
91    fn id(&self) -> ArrayId {
92        static ID: CachedId = CachedId::new("vortex.ext");
93        *ID
94    }
95
96    fn validate(
97        &self,
98        _data: &EmptyArrayData,
99        dtype: &DType,
100        len: usize,
101        slots: &[Option<ArrayRef>],
102    ) -> VortexResult<()> {
103        let storage = slots[ExtensionSlots::STORAGE]
104            .as_ref()
105            .vortex_expect("ExtensionArray storage slot");
106        vortex_ensure_eq!(
107            storage.len(),
108            len,
109            "ExtensionArray length does not match outer length",
110        );
111
112        let ext_dtype = dtype
113            .as_extension_opt()
114            .ok_or_else(|| vortex_err!("not an extension dtype"))?;
115
116        let actual_dtype = DType::Extension(ext_dtype.clone());
117        vortex_ensure_eq!(
118            &actual_dtype,
119            dtype,
120            "ExtensionArray dtype does not match outer dtype",
121        );
122
123        Ok(())
124    }
125
126    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
127        0
128    }
129
130    fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
131        vortex_panic!("ExtensionArray buffer index {idx} out of bounds")
132    }
133
134    fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
135        None
136    }
137
138    fn with_buffers(
139        &self,
140        array: ArrayView<'_, Self>,
141        buffers: &[BufferHandle],
142    ) -> VortexResult<ArrayParts<Self>> {
143        with_empty_buffers(self, array, buffers)
144    }
145
146    fn serialize(
147        _array: ArrayView<'_, Self>,
148        _session: &VortexSession,
149    ) -> VortexResult<Option<Vec<u8>>> {
150        Ok(Some(vec![]))
151    }
152
153    fn deserialize(
154        &self,
155        dtype: &DType,
156        len: usize,
157        metadata: &[u8],
158
159        _buffers: &[BufferHandle],
160        children: &dyn ArrayChildren,
161        _session: &VortexSession,
162    ) -> VortexResult<ArrayParts<Self>> {
163        if !metadata.is_empty() {
164            vortex_bail!(
165                "ExtensionArray expects empty metadata, got {} bytes",
166                metadata.len()
167            );
168        }
169        let DType::Extension(ext_dtype) = dtype else {
170            vortex_bail!("Not an extension DType");
171        };
172        if children.len() != 1 {
173            vortex_bail!("Expected 1 child, got {}", children.len());
174        }
175        let storage = children.get(0, ext_dtype.storage_dtype(), len)?;
176        Ok(ArrayParts::new(
177            self.clone(),
178            dtype.clone(),
179            len,
180            EmptyArrayData,
181            ExtensionSlots { storage }.into_slots(),
182        ))
183    }
184
185    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
186        ExtensionSlots::NAMES[idx].to_string()
187    }
188
189    fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
190        Ok(ExecutionResult::done(array))
191    }
192
193    fn append_to_builder(
194        array: ArrayView<'_, Self>,
195        builder: &mut dyn ArrayBuilder,
196        ctx: &mut ExecutionCtx,
197    ) -> VortexResult<()> {
198        let Some(builder) = builder.as_any_mut().downcast_mut::<ExtensionBuilder>() else {
199            vortex_bail!("append_to_builder for Extension requires an ExtensionBuilder");
200        };
201        builder.append_extension_array(&array.into_owned(), ctx)
202    }
203
204    fn reduce(array: ArrayView<'_, Self>) -> VortexResult<Option<ArrayRef>> {
205        RULES.evaluate(array)
206    }
207
208    fn reduce_parent(
209        array: ArrayView<'_, Self>,
210        parent: &ArrayRef,
211        child_idx: usize,
212    ) -> VortexResult<Option<ArrayRef>> {
213        PARENT_RULES.evaluate(array, parent, child_idx)
214    }
215}
vortex_array/arrays/extension/vtable/mod.rs

vortex_array/arrays/extension/vtable/
mod.rs
