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 {len}",
110            storage.len(),
111        );
112
113        let ext_dtype = dtype
114            .as_extension_opt()
115            .ok_or_else(|| vortex_err!("not an extension dtype"))?;
116
117        let actual_dtype = DType::Extension(ext_dtype.clone());
118        vortex_ensure_eq!(
119            &actual_dtype,
120            dtype,
121            "ExtensionArray dtype {actual_dtype} does not match outer dtype {dtype}",
122        );
123
124        Ok(())
125    }
126
127    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
128        0
129    }
130
131    fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
132        vortex_panic!("ExtensionArray buffer index {idx} out of bounds")
133    }
134
135    fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
136        None
137    }
138
139    fn with_buffers(
140        &self,
141        array: ArrayView<'_, Self>,
142        buffers: &[BufferHandle],
143    ) -> VortexResult<ArrayParts<Self>> {
144        with_empty_buffers(self, array, buffers)
145    }
146
147    fn serialize(
148        _array: ArrayView<'_, Self>,
149        _session: &VortexSession,
150    ) -> VortexResult<Option<Vec<u8>>> {
151        Ok(Some(vec![]))
152    }
153
154    fn deserialize(
155        &self,
156        dtype: &DType,
157        len: usize,
158        metadata: &[u8],
159
160        _buffers: &[BufferHandle],
161        children: &dyn ArrayChildren,
162        _session: &VortexSession,
163    ) -> VortexResult<ArrayParts<Self>> {
164        if !metadata.is_empty() {
165            vortex_bail!(
166                "ExtensionArray expects empty metadata, got {} bytes",
167                metadata.len()
168            );
169        }
170        let DType::Extension(ext_dtype) = dtype else {
171            vortex_bail!("Not an extension DType");
172        };
173        if children.len() != 1 {
174            vortex_bail!("Expected 1 child, got {}", children.len());
175        }
176        let storage = children.get(0, ext_dtype.storage_dtype(), len)?;
177        Ok(
178            ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData)
179                .with_slots(ExtensionSlots { storage }.into_slots()),
180        )
181    }
182
183    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
184        ExtensionSlots::NAMES[idx].to_string()
185    }
186
187    fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
188        Ok(ExecutionResult::done(array))
189    }
190
191    fn append_to_builder(
192        array: ArrayView<'_, Self>,
193        builder: &mut dyn ArrayBuilder,
194        ctx: &mut ExecutionCtx,
195    ) -> VortexResult<()> {
196        let Some(builder) = builder.as_any_mut().downcast_mut::<ExtensionBuilder>() else {
197            vortex_bail!("append_to_builder for Extension requires an ExtensionBuilder");
198        };
199        builder.append_extension_array(&array.into_owned(), ctx)
200    }
201
202    fn reduce(array: ArrayView<'_, Self>) -> VortexResult<Option<ArrayRef>> {
203        RULES.evaluate(array)
204    }
205
206    fn reduce_parent(
207        array: ArrayView<'_, Self>,
208        parent: &ArrayRef,
209        child_idx: usize,
210    ) -> VortexResult<Option<ArrayRef>> {
211        PARENT_RULES.evaluate(array, parent, child_idx)
212    }
213}
vortex_array/arrays/extension/vtable/mod.rs

vortex_array/arrays/extension/vtable/
mod.rs
