vortex_array/arrays/extension/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3mod canonical;
4mod kernel;
5mod operations;
6mod validity;
7
8use std::hash::Hasher;
9
10use kernel::PARENT_KERNELS;
11use vortex_error::VortexExpect;
12use vortex_error::VortexResult;
13use vortex_error::vortex_bail;
14use vortex_error::vortex_ensure;
15use vortex_error::vortex_panic;
16use vortex_session::VortexSession;
17use vortex_session::registry::CachedId;
18
19use crate::ArrayEq;
20use crate::ArrayHash;
21use crate::ArrayRef;
22use crate::ExecutionCtx;
23use crate::ExecutionResult;
24use crate::Precision;
25use crate::array::Array;
26use crate::array::ArrayId;
27use crate::array::ArrayView;
28use crate::array::VTable;
29use crate::array::ValidityVTableFromChild;
30use crate::arrays::extension::ExtensionData;
31use crate::arrays::extension::array::SLOT_NAMES;
32use crate::arrays::extension::array::STORAGE_SLOT;
33use crate::arrays::extension::compute::rules::PARENT_RULES;
34use crate::buffer::BufferHandle;
35use crate::dtype::DType;
36use crate::serde::ArrayChildren;
37
38/// A [`Extension`]-encoded Vortex array.
39pub type ExtensionArray = Array<Extension>;
40
41impl ArrayHash for ExtensionData {
42    fn array_hash<H: Hasher>(&self, _state: &mut H, _precision: Precision) {}
43}
44
45impl ArrayEq for ExtensionData {
46    fn array_eq(&self, _other: &Self, _precision: Precision) -> bool {
47        true
48    }
49}
50
51impl VTable for Extension {
52    type ArrayData = ExtensionData;
53
54    type OperationsVTable = Self;
55    type ValidityVTable = ValidityVTableFromChild;
56
57    fn id(&self) -> ArrayId {
58        static ID: CachedId = CachedId::new("vortex.ext");
59        *ID
60    }
61
62    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
63        0
64    }
65
66    fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
67        vortex_panic!("ExtensionArray buffer index {idx} out of bounds")
68    }
69
70    fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
71        None
72    }
73
74    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
75        SLOT_NAMES[idx].to_string()
76    }
77
78    fn serialize(
79        _array: ArrayView<'_, Self>,
80        _session: &VortexSession,
81    ) -> VortexResult<Option<Vec<u8>>> {
82        Ok(Some(vec![]))
83    }
84
85    fn validate(
86        &self,
87        data: &ExtensionData,
88        dtype: &DType,
89        len: usize,
90        slots: &[Option<ArrayRef>],
91    ) -> VortexResult<()> {
92        _ = data;
93        let storage = slots[STORAGE_SLOT]
94            .as_ref()
95            .vortex_expect("ExtensionArray storage slot");
96        vortex_ensure!(
97            storage.len() == len,
98            "ExtensionArray length {} does not match outer length {}",
99            storage.len(),
100            len
101        );
102
103        let actual_dtype = DType::Extension(data.ext_dtype.clone());
104        vortex_ensure!(
105            &actual_dtype == dtype,
106            "ExtensionArray dtype {} does not match outer dtype {}",
107            actual_dtype,
108            dtype
109        );
110
111        Ok(())
112    }
113
114    fn deserialize(
115        &self,
116        dtype: &DType,
117        len: usize,
118        metadata: &[u8],
119
120        _buffers: &[BufferHandle],
121        children: &dyn ArrayChildren,
122        _session: &VortexSession,
123    ) -> VortexResult<crate::array::ArrayParts<Self>> {
124        if !metadata.is_empty() {
125            vortex_bail!(
126                "ExtensionArray expects empty metadata, got {} bytes",
127                metadata.len()
128            );
129        }
130        let DType::Extension(ext_dtype) = dtype else {
131            vortex_bail!("Not an extension DType");
132        };
133        if children.len() != 1 {
134            vortex_bail!("Expected 1 child, got {}", children.len());
135        }
136        let storage = children.get(0, ext_dtype.storage_dtype(), len)?;
137        Ok(crate::array::ArrayParts::new(
138            self.clone(),
139            dtype.clone(),
140            len,
141            ExtensionData::new(ext_dtype.clone(), storage.dtype()),
142        )
143        .with_slots(vec![Some(storage)]))
144    }
145
146    fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
147        Ok(ExecutionResult::done(array))
148    }
149
150    fn reduce_parent(
151        array: ArrayView<'_, Self>,
152        parent: &ArrayRef,
153        child_idx: usize,
154    ) -> VortexResult<Option<ArrayRef>> {
155        PARENT_RULES.evaluate(array, parent, child_idx)
156    }
157
158    fn execute_parent(
159        array: ArrayView<'_, Self>,
160        parent: &ArrayRef,
161        child_idx: usize,
162        ctx: &mut ExecutionCtx,
163    ) -> VortexResult<Option<ArrayRef>> {
164        PARENT_KERNELS.execute(array, parent, child_idx, ctx)
165    }
166}
167
168#[derive(Clone, Debug)]
169pub struct Extension;
vortex_array/arrays/extension/vtable/mod.rs

vortex_array/arrays/extension/vtable/
mod.rs
