vortex_array/arrays/extension/vtable/
mod.rs

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