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

vortex_array/arrays/extension/vtable/
mod.rs
