Skip to content

Move array serde from VTable to standalone ArrayPlugin impls #10111

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
mhk197 wants to merge 2 commits into
base: develop
Choose a base branch
from
Draft

Move array serde from VTable to standalone ArrayPlugin impls #10111

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b701797
Move array serde from VTable to standalone ArrayPlugin impls
mhk197 Sep 28, 2026
e50242a
Impl ArrayPlugin directly on DecimalByteParts
mhk197 Sep 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 2 additions & 63 deletions encodings/alp/src/alp/array.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@ use std::fmt::Formatter;
use std::hash::Hash;
use std::hash::Hasher;

use prost::Message;
use vortex_array::Array;
use vortex_array::ArrayEq;
use vortex_array::ArrayHash;
Expand All @@ -33,7 +32,6 @@ use vortex_array::patches::PatchesData;
use vortex_array::patches::PatchesMetadata;
use vortex_array::require_child;
use vortex_array::require_patches;
use vortex_array::serde::ArrayChildren;
use vortex_array::smallvec::smallvec;
use vortex_array::vtable::VTable;
use vortex_array::vtable::ValidityChild;
Expand All @@ -43,14 +41,15 @@ use vortex_error::VortexResult;
use vortex_error::vortex_bail;
use vortex_error::vortex_ensure;
use vortex_error::vortex_panic;
use vortex_session::VortexSession;
use vortex_session::registry::CachedId;

use crate::ALPFloat;
use crate::alp::Exponents;
use crate::alp::decompress::execute_decompress;
use crate::alp::rules::RULES;

mod plugin;

/// A [`ALP`]-encoded Vortex array.
pub type ALPArray = Array<ALP>;

Expand Down Expand Up @@ -111,66 +110,6 @@ impl VTable for ALP {
vortex_array::vtable::with_empty_buffers(self, array, buffers)
}

fn serialize(
array: ArrayView<'_, Self>,
_session: &VortexSession,
) -> VortexResult<Option<Vec<u8>>> {
let exponents = array.exponents();
Ok(Some(
ALPMetadata {
exp_e: exponents.e as u32,
exp_f: exponents.f as u32,
patches: array
.patches()
.map(|p| p.to_metadata(array.len(), array.dtype()))
.transpose()?,
}
.encode_to_vec(),
))
}

fn deserialize(
&self,
dtype: &DType,
len: usize,
metadata: &[u8],
_buffers: &[BufferHandle],
children: &dyn ArrayChildren,
_session: &VortexSession,
) -> VortexResult<ArrayParts<Self>> {
let metadata = ALPMetadata::decode(metadata)?;
let encoded_ptype = match &dtype {
DType::Primitive(PType::F32, n) => DType::Primitive(PType::I32, *n),
DType::Primitive(PType::F64, n) => DType::Primitive(PType::I64, *n),
d => vortex_bail!(MismatchedTypes: "f32 or f64", d),
};
let encoded = children.get(0, &encoded_ptype, len)?;

let patches = metadata
.patches
.map(|p| {
let indices = children.get(1, &p.indices_dtype()?, p.len()?)?;
let values = children.get(2, dtype, p.len()?)?;
let chunk_offsets = p
.chunk_offsets_dtype()?
.map(|dtype| children.get(3, &dtype, usize::try_from(p.chunk_offsets_len())?))
.transpose()?;

Patches::new(len, p.offset()?, indices, values, chunk_offsets)
})
.transpose()?;

let slots = ALPData::make_slots(&encoded, patches.as_ref());
let data = ALPData::new(
Exponents {
e: u8::try_from(metadata.exp_e)?,
f: u8::try_from(metadata.exp_f)?,
},
patches,
);
Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots))
}

fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
ALPSlots::NAMES[idx].to_string()
}
Expand Down
111 changes: 111 additions & 0 deletions encodings/alp/src/alp/array/plugin.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,111 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use prost::Message;
use vortex_array::Array;
use vortex_array::ArrayDeserialization;
use vortex_array::ArrayId;
use vortex_array::ArrayParts;
use vortex_array::ArrayPlugin;
use vortex_array::ArrayRef;
use vortex_array::ArraySerialization;
use vortex_array::IntoArray;
use vortex_array::dtype::DType;
use vortex_array::dtype::PType;
use vortex_array::patches::Patches;
use vortex_array::vtable::VTable;
use vortex_error::VortexResult;
use vortex_error::vortex_bail;
use vortex_error::vortex_ensure;
use vortex_error::vortex_err;
use vortex_session::VortexSession;

use super::ALP;
use super::ALPArrayExt;
use super::ALPData;
use super::ALPMetadata;
use crate::alp::Exponents;

impl ArrayPlugin for ALP {
fn id(&self) -> ArrayId {
VTable::id(self)
}

fn serialize(
&self,
array: &ArrayRef,
_session: &VortexSession,
) -> VortexResult<Option<ArraySerialization>> {
let view = array
.as_opt::<ALP>()
.ok_or_else(|| vortex_err!("ALP plugin cannot serialize {}", array.encoding_id()))?;
let exponents = view.exponents();
let metadata = ALPMetadata {
exp_e: exponents.e as u32,
exp_f: exponents.f as u32,
patches: view
.patches()
.map(|p| p.to_metadata(view.len(), view.dtype()))
.transpose()?,
}
.encode_to_vec();
Ok(Some(ArraySerialization::from_array(
VTable::id(self),
array,
metadata,
)))
}

fn deserialize(
&self,
parts: ArrayDeserialization<'_>,
_session: &VortexSession,
) -> VortexResult<ArrayRef> {
let ArrayDeserialization {
serialized_id,
dtype,
len,
metadata,
buffers: _,
children,
} = parts;
vortex_ensure!(
serialized_id == VTable::id(self),
"ALP plugin does not recognize serialized ID {serialized_id}"
);
let metadata = ALPMetadata::decode(metadata)?;
let encoded_ptype = match &dtype {
DType::Primitive(PType::F32, n) => DType::Primitive(PType::I32, *n),
DType::Primitive(PType::F64, n) => DType::Primitive(PType::I64, *n),
d => vortex_bail!(MismatchedTypes: "f32 or f64", d),
};
let encoded = children.get(0, &encoded_ptype, len)?;

let patches = metadata
.patches
.map(|p| {
let indices = children.get(1, &p.indices_dtype()?, p.len()?)?;
let values = children.get(2, dtype, p.len()?)?;
let chunk_offsets = p
.chunk_offsets_dtype()?
.map(|dtype| children.get(3, &dtype, usize::try_from(p.chunk_offsets_len())?))
.transpose()?;

Patches::new(len, p.offset()?, indices, values, chunk_offsets)
})
.transpose()?;

let slots = ALPData::make_slots(&encoded, patches.as_ref());
let data = ALPData::new(
Exponents {
e: u8::try_from(metadata.exp_e)?,
f: u8::try_from(metadata.exp_f)?,
},
patches,
);
Ok(Array::try_from_parts(
ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots),
)?
.into_array())
}
}
14 changes: 3 additions & 11 deletions encodings/alp/src/alp/plugin.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
//!
//! This enables zero-cost backward compatibility with previously written datasets.

use vortex_array::Array;
use vortex_array::ArrayDeserialization;
use vortex_array::ArrayId;
use vortex_array::ArrayPlugin;
Expand Down Expand Up @@ -57,16 +56,9 @@ impl ArrayPlugin for ALPPatchedPlugin {
"ALP plugin does not recognize serialized ID {}",
parts.serialized_id,
);
let alp_array = Array::<ALP>::try_from_parts(ArrayVTable::deserialize(
&ALP,
parts.dtype,
parts.len,
parts.metadata,
parts.buffers,
parts.children,
session,
)?)
.map_err(|_| vortex_err!("ALP plugin should only deserialize vortex.alp"))?;
let alp_array = ArrayPlugin::deserialize(&ALP, parts, session)?
.try_downcast::<ALP>()
.map_err(|_| vortex_err!("ALP plugin should only deserialize vortex.alp"))?;

// Check if there are interior patches to externalize.
let Some(patches) = alp_array.patches() else {
Expand Down
102 changes: 2 additions & 100 deletions encodings/alp/src/alp_rd/array.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,6 @@ use std::fmt::Formatter;
use std::hash::Hash;
use std::hash::Hasher;

use itertools::Itertools;
use prost::Message;
use vortex_array::Array;
use vortex_array::ArrayEq;
use vortex_array::ArrayHash;
Expand All @@ -35,7 +33,6 @@ use vortex_array::patches::PatchesData;
use vortex_array::patches::PatchesMetadata;
use vortex_array::require_child;
use vortex_array::require_patches;
use vortex_array::serde::ArrayChildren;
use vortex_array::smallvec::smallvec;
use vortex_array::validity::Validity;
use vortex_array::vtable::VTable;
Expand All @@ -46,14 +43,14 @@ use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use vortex_error::vortex_bail;
use vortex_error::vortex_ensure;
use vortex_error::vortex_err;
use vortex_error::vortex_panic;
use vortex_session::VortexSession;
use vortex_session::registry::CachedId;

use crate::alp_rd::rules::RULES;
use crate::alp_rd_decode;

mod plugin;

/// A [`ALPRD`]-encoded Vortex array.
pub type ALPRDArray = Array<ALPRD>;

Expand Down Expand Up @@ -136,101 +133,6 @@ impl VTable for ALPRD {
vortex_array::vtable::with_empty_buffers(self, array, buffers)
}

fn serialize(
array: ArrayView<'_, Self>,
_session: &VortexSession,
) -> VortexResult<Option<Vec<u8>>> {
let dict = array
.left_parts_dictionary()
.iter()
.map(|&i| i as u32)
.collect::<Vec<_>>();

Ok(Some(
ALPRDMetadata {
right_bit_width: array.right_bit_width() as u32,
dict_len: array.left_parts_dictionary().len() as u32,
dict,
left_parts_ptype: array.left_parts().dtype().as_ptype() as i32,
patches: array
.left_parts_patches()
.map(|p| p.to_metadata(array.len(), p.dtype()))
.transpose()?,
}
.encode_to_vec(),
))
}

#[allow(clippy::disallowed_methods)]
fn deserialize(
&self,
dtype: &DType,
len: usize,
metadata: &[u8],
_buffers: &[BufferHandle],
children: &dyn ArrayChildren,
_session: &VortexSession,
) -> VortexResult<ArrayParts<Self>> {
let metadata = ALPRDMetadata::decode(metadata)?;
if children.len() < 2 {
vortex_bail!(
"Expected at least 2 children for ALPRD encoding, found {}",
children.len()
);
}

let left_parts_dtype = DType::Primitive(metadata.left_parts_ptype(), dtype.nullability());
let left_parts = children.get(0, &left_parts_dtype, len)?;
let left_parts_dictionary: Buffer<u16> = metadata.dict.as_slice()
[0..metadata.dict_len as usize]
.iter()
.map(|&i| {
u16::try_from(i)
.map_err(|_| vortex_err!("left_parts_dictionary code {i} does not fit in u16"))
})
.try_collect()?;

let right_parts_dtype = match &dtype {
DType::Primitive(PType::F32, _) => {
DType::Primitive(PType::U32, Nullability::NonNullable)
}
DType::Primitive(PType::F64, _) => {
DType::Primitive(PType::U64, Nullability::NonNullable)
}
_ => vortex_bail!("Expected f32 or f64 dtype, got {:?}", dtype),
};
let right_parts = children.get(1, &right_parts_dtype, len)?;

let left_parts_patches = metadata
.patches
.map(|p| {
let indices = children.get(2, &p.indices_dtype()?, p.len()?)?;
let values = children.get(3, &left_parts_dtype.as_nonnullable(), p.len()?)?;

Patches::new(
len,
p.offset()?,
indices,
values,
// TODO(0ax1): handle chunk offsets
None,
)
})
.transpose()?;
let slots = ALPRDData::make_slots(&left_parts, &right_parts, left_parts_patches.as_ref());
let data = ALPRDData::new(
left_parts_dictionary,
u8::try_from(metadata.right_bit_width).map_err(|_| {
vortex_err!(
"right_bit_width {} out of u8 range",
metadata.right_bit_width
)
})?,
left_parts_patches,
);
Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots))
}

fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
ALPRDSlots::NAMES[idx].to_string()
}
Expand Down
Loading
Loading