vortex_array/arrays/extension/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod array;
5mod canonical;
6mod operations;
7mod rules;
8mod validity;
9mod visitor;
10
11use vortex_dtype::DType;
12use vortex_error::VortexExpect;
13use vortex_error::VortexResult;
14use vortex_error::vortex_bail;
15use vortex_error::vortex_ensure;
16use vortex_vector::Vector;
17
18use crate::ArrayRef;
19use crate::EmptyMetadata;
20use crate::VectorExecutor;
21use crate::arrays::extension::ExtensionArray;
22use crate::arrays::extension::vtable::rules::PARENT_RULES;
23use crate::buffer::BufferHandle;
24use crate::executor::ExecutionCtx;
25use crate::serde::ArrayChildren;
26use crate::vtable;
27use crate::vtable::ArrayId;
28use crate::vtable::ArrayVTable;
29use crate::vtable::ArrayVTableExt;
30use crate::vtable::NotSupported;
31use crate::vtable::VTable;
32use crate::vtable::ValidityVTableFromChild;
33
34vtable!(Extension);
35
36impl VTable for ExtensionVTable {
37    type Array = ExtensionArray;
38
39    type Metadata = EmptyMetadata;
40
41    type ArrayVTable = Self;
42    type CanonicalVTable = Self;
43    type OperationsVTable = Self;
44    type ValidityVTable = ValidityVTableFromChild;
45    type VisitorVTable = Self;
46    type ComputeVTable = NotSupported;
47    type EncodeVTable = NotSupported;
48
49    fn id(&self) -> ArrayId {
50        ArrayId::new_ref("vortex.ext")
51    }
52
53    fn encoding(_array: &Self::Array) -> ArrayVTable {
54        ExtensionVTable.as_vtable()
55    }
56
57    fn metadata(_array: &ExtensionArray) -> VortexResult<Self::Metadata> {
58        Ok(EmptyMetadata)
59    }
60
61    fn serialize(_metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
62        Ok(Some(vec![]))
63    }
64
65    fn deserialize(_buffer: &[u8]) -> VortexResult<Self::Metadata> {
66        Ok(EmptyMetadata)
67    }
68
69    fn build(
70        &self,
71        dtype: &DType,
72        len: usize,
73        _metadata: &Self::Metadata,
74        _buffers: &[BufferHandle],
75        children: &dyn ArrayChildren,
76    ) -> VortexResult<ExtensionArray> {
77        let DType::Extension(ext_dtype) = dtype else {
78            vortex_bail!("Not an extension DType");
79        };
80        if children.len() != 1 {
81            vortex_bail!("Expected 1 child, got {}", children.len());
82        }
83        let storage = children.get(0, ext_dtype.storage_dtype(), len)?;
84        Ok(ExtensionArray::new(ext_dtype.clone(), storage))
85    }
86
87    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
88        vortex_ensure!(
89            children.len() == 1,
90            "ExtensionArray expects exactly 1 child (storage), got {}",
91            children.len()
92        );
93        array.storage = children
94            .into_iter()
95            .next()
96            .vortex_expect("children length already validated");
97        Ok(())
98    }
99
100    fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<Vector> {
101        array.storage().execute(ctx)
102    }
103
104    fn reduce_parent(
105        array: &Self::Array,
106        parent: &ArrayRef,
107        child_idx: usize,
108    ) -> VortexResult<Option<ArrayRef>> {
109        PARENT_RULES.evaluate(array, parent, child_idx)
110    }
111}
112
113#[derive(Debug)]
114pub struct ExtensionVTable;