mod array;
mod canonical;
mod kernel;
mod operations;
mod validity;
mod visitor;
use kernel::PARENT_KERNELS;
use vortex_dtype::DType;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use vortex_error::vortex_bail;
use vortex_error::vortex_ensure;
use vortex_session::VortexSession;
use crate::ArrayRef;
use crate::EmptyMetadata;
use crate::ExecutionCtx;
use crate::arrays::extension::ExtensionArray;
use crate::arrays::extension::compute::rules::PARENT_RULES;
use crate::buffer::BufferHandle;
use crate::serde::ArrayChildren;
use crate::vtable;
use crate::vtable::ArrayId;
use crate::vtable::VTable;
use crate::vtable::ValidityVTableFromChild;
vtable!(Extension);
impl VTable for ExtensionVTable {
type Array = ExtensionArray;
type Metadata = EmptyMetadata;
type ArrayVTable = Self;
type OperationsVTable = Self;
type ValidityVTable = ValidityVTableFromChild;
type VisitorVTable = Self;
fn id(_array: &Self::Array) -> ArrayId {
Self::ID
}
fn metadata(_array: &ExtensionArray) -> VortexResult<Self::Metadata> {
Ok(EmptyMetadata)
}
fn serialize(_metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
Ok(Some(vec![]))
}
fn deserialize(
_bytes: &[u8],
_dtype: &DType,
_len: usize,
_buffers: &[BufferHandle],
_session: &VortexSession,
) -> VortexResult<Self::Metadata> {
Ok(EmptyMetadata)
}
fn build(
dtype: &DType,
len: usize,
_metadata: &Self::Metadata,
_buffers: &[BufferHandle],
children: &dyn ArrayChildren,
) -> VortexResult<ExtensionArray> {
let DType::Extension(ext_dtype) = dtype else {
vortex_bail!("Not an extension DType");
};
if children.len() != 1 {
vortex_bail!("Expected 1 child, got {}", children.len());
}
let storage = children.get(0, ext_dtype.storage_dtype(), len)?;
Ok(ExtensionArray::new(ext_dtype.clone(), storage))
}
fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
vortex_ensure!(
children.len() == 1,
"ExtensionArray expects exactly 1 child (storage), got {}",
children.len()
);
array.storage = children
.into_iter()
.next()
.vortex_expect("children length already validated");
Ok(())
}
fn execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
Ok(array.to_array())
}
fn reduce_parent(
array: &Self::Array,
parent: &ArrayRef,
child_idx: usize,
) -> VortexResult<Option<ArrayRef>> {
PARENT_RULES.evaluate(array, parent, child_idx)
}
fn execute_parent(
array: &Self::Array,
parent: &ArrayRef,
child_idx: usize,
ctx: &mut ExecutionCtx,
) -> VortexResult<Option<ArrayRef>> {
PARENT_KERNELS.execute(array, parent, child_idx, ctx)
}
}
#[derive(Debug)]
pub struct ExtensionVTable;
impl ExtensionVTable {
pub const ID: ArrayId = ArrayId::new_ref("vortex.ext");
}