diff --git a/encodings/alp/src/alp/array.rs b/encodings/alp/src/alp/array.rs index 3e2010a1ad5..bb210d76cb3 100644 --- a/encodings/alp/src/alp/array.rs +++ b/encodings/alp/src/alp/array.rs @@ -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; @@ -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; @@ -43,7 +41,6 @@ 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; @@ -51,6 +48,8 @@ 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; @@ -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>> { - 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> { - 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() } diff --git a/encodings/alp/src/alp/array/plugin.rs b/encodings/alp/src/alp/array/plugin.rs new file mode 100644 index 00000000000..9a2486c7247 --- /dev/null +++ b/encodings/alp/src/alp/array/plugin.rs @@ -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> { + let view = array + .as_opt::() + .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 { + 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()) + } +} diff --git a/encodings/alp/src/alp/plugin.rs b/encodings/alp/src/alp/plugin.rs index 5756dc5a991..35750732003 100644 --- a/encodings/alp/src/alp/plugin.rs +++ b/encodings/alp/src/alp/plugin.rs @@ -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; @@ -57,16 +56,9 @@ impl ArrayPlugin for ALPPatchedPlugin { "ALP plugin does not recognize serialized ID {}", parts.serialized_id, ); - let alp_array = Array::::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::() + .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 { diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array.rs index b6ee50d7b1f..1fc30d779e8 100644 --- a/encodings/alp/src/alp_rd/array.rs +++ b/encodings/alp/src/alp_rd/array.rs @@ -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; @@ -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; @@ -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; @@ -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>> { - let dict = array - .left_parts_dictionary() - .iter() - .map(|&i| i as u32) - .collect::>(); - - 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> { - 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 = 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() } diff --git a/encodings/alp/src/alp_rd/array/plugin.rs b/encodings/alp/src/alp_rd/array/plugin.rs new file mode 100644 index 00000000000..998b87b8907 --- /dev/null +++ b/encodings/alp/src/alp_rd/array/plugin.rs @@ -0,0 +1,149 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use itertools::Itertools; +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::Nullability; +use vortex_array::dtype::PType; +use vortex_array::patches::Patches; +use vortex_array::vtable::VTable; +use vortex_buffer::Buffer; +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::ALPRD; +use super::ALPRDArrayExt; +use super::ALPRDArraySlotsExt; +use super::ALPRDData; +use super::ALPRDMetadata; + +impl ArrayPlugin for ALPRD { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array + .as_opt::() + .ok_or_else(|| vortex_err!("ALPRD plugin cannot serialize {}", array.encoding_id()))?; + let dict = view + .left_parts_dictionary() + .iter() + .map(|&i| i as u32) + .collect::>(); + + let metadata = ALPRDMetadata { + right_bit_width: view.right_bit_width() as u32, + dict_len: view.left_parts_dictionary().len() as u32, + dict, + left_parts_ptype: view.left_parts().dtype().as_ptype() as i32, + patches: view + .left_parts_patches() + .map(|p| p.to_metadata(view.len(), p.dtype())) + .transpose()?, + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata, + ))) + } + + #[allow(clippy::disallowed_methods)] + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers: _, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "ALPRD plugin does not recognize serialized ID {serialized_id}" + ); + 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 = 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(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/encodings/bytebool/src/array.rs b/encodings/bytebool/src/array.rs index 02bb4fcb2b8..9dc44e2db91 100644 --- a/encodings/bytebool/src/array.rs +++ b/encodings/bytebool/src/array.rs @@ -24,7 +24,6 @@ use vortex_array::arrays::BoolArray; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::scalar::Scalar; -use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::OperationsVTable; use vortex_array::vtable::VTable; @@ -34,12 +33,12 @@ use vortex_array::vtable::validity_to_child; use vortex_buffer::BitBufferMut; use vortex_buffer::ByteBuffer; 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; +mod plugin; + /// A [`ByteBool`]-encoded Vortex array. pub type ByteBoolArray = Array; @@ -113,47 +112,6 @@ impl VTable for ByteBool { ) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - if !metadata.is_empty() { - vortex_bail!( - "ByteBoolArray expects empty metadata, got {} bytes", - metadata.len() - ); - } - let validity = if children.is_empty() { - Validity::from(dtype.nullability()) - } else if children.len() == 1 { - let validity = children.get(0, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("Expected 0 or 1 child, got {}", children.len()); - }; - - if buffers.len() != 1 { - vortex_bail!("Expected 1 buffer, got {}", buffers.len()); - } - let buffer = buffers[0].clone(); - - let data = ByteBoolData::new(buffer); - let slots = ByteBoolData::make_slots(&validity, len); - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { ByteBoolSlots::NAMES[idx].to_string() } @@ -338,6 +296,7 @@ mod tests { use vortex_array::serde::SerializedArray; use vortex_array::session::ArraySessionExt; use vortex_buffer::ByteBufferMut; + use vortex_session::VortexSession; use vortex_session::registry::ReadContext; use super::*; diff --git a/encodings/bytebool/src/array/plugin.rs b/encodings/bytebool/src/array/plugin.rs new file mode 100644 index 00000000000..8bae86b88f5 --- /dev/null +++ b/encodings/bytebool/src/array/plugin.rs @@ -0,0 +1,88 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +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::validity::Validity; +use vortex_array::vtable::VTable; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_session::VortexSession; + +use super::ByteBool; +use super::ByteBoolData; + +impl ArrayPlugin for ByteBool { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + array.is::(), + "ByteBool plugin cannot serialize {}", + array.encoding_id() + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "ByteBool plugin does not recognize serialized ID {serialized_id}" + ); + if !metadata.is_empty() { + vortex_bail!( + "ByteBoolArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + let validity = if children.is_empty() { + Validity::from(dtype.nullability()) + } else if children.len() == 1 { + let validity = children.get(0, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + vortex_bail!("Expected 0 or 1 child, got {}", children.len()); + }; + + if buffers.len() != 1 { + vortex_bail!("Expected 1 buffer, got {}", buffers.len()); + } + let buffer = buffers[0].clone(); + + let data = ByteBoolData::new(buffer); + let slots = ByteBoolData::make_slots(&validity, len); + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/encodings/datetime-parts/src/array.rs b/encodings/datetime-parts/src/array.rs index 30594c78367..ff26fc5c6db 100644 --- a/encodings/datetime-parts/src/array.rs +++ b/encodings/datetime-parts/src/array.rs @@ -6,7 +6,6 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hasher; -use prost::Message; use vortex_array::AnyCanonical; use vortex_array::Array; use vortex_array::ArrayEq; @@ -24,10 +23,8 @@ use vortex_array::arrays::Primitive; use vortex_array::arrays::TemporalArray; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::require_child; -use vortex_array::serde::ArrayChildren; use vortex_array::smallvec::smallvec; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityChild; @@ -37,7 +34,6 @@ 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::TemporalParts; @@ -45,6 +41,8 @@ use crate::canonical::decode_to_temporal; use crate::compute::rules::PARENT_RULES; use crate::split_temporal; +mod plugin; + /// A [`DateTimeParts`]-encoded Vortex array. pub type DateTimePartsArray = Array; @@ -130,58 +128,6 @@ impl VTable for DateTimeParts { vortex_array::vtable::with_empty_buffers(self, array, buffers) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - DateTimePartsMetadata { - days_ptype: PType::try_from(array.days().dtype())? as i32, - seconds_ptype: PType::try_from(array.seconds().dtype())? as i32, - subseconds_ptype: PType::try_from(array.subseconds().dtype())? as i32, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - let metadata = DateTimePartsMetadata::decode(metadata)?; - if children.len() != 3 { - vortex_bail!( - "Expected 3 children for datetime-parts encoding, found {}", - children.len() - ) - } - - let days = children.get( - 0, - &DType::Primitive(metadata.get_days_ptype()?, dtype.nullability()), - len, - )?; - let seconds = children.get( - 1, - &DType::Primitive(metadata.get_seconds_ptype()?, Nullability::NonNullable), - len, - )?; - let subseconds = children.get( - 2, - &DType::Primitive(metadata.get_subseconds_ptype()?, Nullability::NonNullable), - len, - )?; - - let slots = smallvec![Some(days), Some(seconds), Some(subseconds)]; - let data = DateTimePartsData {}; - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { DateTimePartsSlots::NAMES[idx].to_string() } diff --git a/encodings/datetime-parts/src/array/plugin.rs b/encodings/datetime-parts/src/array/plugin.rs new file mode 100644 index 00000000000..2a8131f5593 --- /dev/null +++ b/encodings/datetime-parts/src/array/plugin.rs @@ -0,0 +1,106 @@ +// 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::Nullability; +use vortex_array::dtype::PType; +use vortex_array::smallvec::smallvec; +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::DateTimeParts; +use super::DateTimePartsArraySlotsExt; +use super::DateTimePartsData; +use super::DateTimePartsMetadata; + +impl ArrayPlugin for DateTimeParts { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_opt::().ok_or_else(|| { + vortex_err!( + "DateTimeParts plugin cannot serialize {}", + array.encoding_id() + ) + })?; + let metadata = DateTimePartsMetadata { + days_ptype: PType::try_from(view.days().dtype())? as i32, + seconds_ptype: PType::try_from(view.seconds().dtype())? as i32, + subseconds_ptype: PType::try_from(view.subseconds().dtype())? as i32, + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers: _, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "DateTimeParts plugin does not recognize serialized ID {serialized_id}" + ); + let metadata = DateTimePartsMetadata::decode(metadata)?; + if children.len() != 3 { + vortex_bail!( + "Expected 3 children for datetime-parts encoding, found {}", + children.len() + ) + } + + let days = children.get( + 0, + &DType::Primitive(metadata.get_days_ptype()?, dtype.nullability()), + len, + )?; + let seconds = children.get( + 1, + &DType::Primitive(metadata.get_seconds_ptype()?, Nullability::NonNullable), + len, + )?; + let subseconds = children.get( + 2, + &DType::Primitive(metadata.get_subseconds_ptype()?, Nullability::NonNullable), + len, + )?; + + let slots = smallvec![Some(days), Some(seconds), Some(subseconds)]; + let data = DateTimePartsData {}; + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs index c1826ae46d4..5104ebe95f0 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs @@ -25,7 +25,6 @@ use vortex_array::dtype::DecimalDType; use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; -use vortex_array::serde::ArrayChildren; use vortex_array::vtable::OperationsVTable; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityChild; @@ -35,7 +34,6 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use super::MAX_LOWER_PARTS; use super::assemble::assemble_decimal; @@ -252,25 +250,6 @@ impl VTable for DecimalByteParts { vortex_array::vtable::with_empty_buffers(self, array, buffers) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - vortex_bail!("DecimalByteParts serialization requires DecimalBytePartsPlugin") - } - - fn deserialize( - &self, - _dtype: &DType, - _len: usize, - _metadata: &[u8], - _buffers: &[BufferHandle], - _children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_bail!("DecimalByteParts deserialization requires DecimalBytePartsPlugin") - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { DecimalBytePartsSlots::slot_name(idx) } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index 65bb8222f7a..a0ef8a6674f 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -31,7 +31,6 @@ mod split; mod testing; pub use array::*; -pub use plugin::DecimalBytePartsPlugin; pub use plugin::DecimalBytePartsV2Metadata; pub use plugin::decimal_byte_parts_v1_id; pub use plugin::decimal_byte_parts_v2_id; diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/mod.rs index 3896cd51c96..8e4a3840938 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/mod.rs @@ -44,15 +44,9 @@ pub fn decimal_byte_parts_v2_id() -> ArrayId { /// Each version owns its metadata schema and serde functions. The plugin writes v1 whenever an /// array has no lower parts, so such arrays stay readable by older readers, and v2 otherwise. /// The v2 format itself accepts any lower part count up to the maximum. -/// -/// Register this plugin, or call [`crate::initialize`], to enable both formats. Direct registration -/// of [`DecimalByteParts`] does not support serde. -#[derive(Clone, Debug)] -pub struct DecimalBytePartsPlugin; - -impl ArrayPlugin for DecimalBytePartsPlugin { +impl ArrayPlugin for DecimalByteParts { fn id(&self) -> ArrayId { - VTable::id(&DecimalByteParts) + VTable::id(self) } fn serialized_ids(&self) -> Vec { diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/tests.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/tests.rs index 92c9d1edbb7..de4ba8bf0d2 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/tests.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/tests.rs @@ -244,40 +244,6 @@ fn serialization_requires_v2_permission() -> VortexResult<()> { Ok(()) } -#[test] -fn bare_vtable_refuses_serde() -> VortexResult<()> { - let session = array_session(); - session.arrays().register(DecimalByteParts); - let msp = msp(); - let array = DecimalByteParts::try_new(msp.clone(), DecimalDType::new(19, 2))?.into_array(); - let result = session.array_serialize(&array); - assert!( - result.as_ref().is_err_and(|err| err - .to_string() - .contains("DecimalByteParts serialization requires DecimalBytePartsPlugin")), - "expected unsupported VTable serialization, got {result:?}" - ); - - let id = VTable::id(&DecimalByteParts); - let plugin = session - .arrays() - .registry() - .get(&id) - .vortex_expect("registered"); - let children = vec![msp]; - let result = plugin.deserialize( - ArrayDeserialization::new(id, array.dtype(), array.len(), &[8, 7], &[], &children), - &session, - ); - assert!( - result.as_ref().is_err_and(|err| err - .to_string() - .contains("DecimalByteParts deserialization requires DecimalBytePartsPlugin")), - "expected unsupported VTable deserialization, got {result:?}" - ); - Ok(()) -} - fn msp() -> ArrayRef { buffer![1i64, 2, 3].into_array() } @@ -317,7 +283,8 @@ fn deserialize_with( children: Vec, ) -> VortexResult { let dtype = DType::Decimal(DecimalDType::new(38, 2), Nullability::NonNullable); - DecimalBytePartsPlugin.deserialize( + ArrayPlugin::deserialize( + &DecimalByteParts, ArrayDeserialization::new(serialized_id, &dtype, 3, metadata, &[], &children), &array_session(), ) diff --git a/encodings/decimal-byte-parts/src/lib.rs b/encodings/decimal-byte-parts/src/lib.rs index 2557555eac8..ca65e43140e 100644 --- a/encodings/decimal-byte-parts/src/lib.rs +++ b/encodings/decimal-byte-parts/src/lib.rs @@ -24,7 +24,7 @@ use vortex_session::VortexSession; pub fn initialize(session: &VortexSession) { // One plugin owns both serialized formats: registering it reads either ID and writes the // one that fits the array. Which of them a writer may emit is decided by its editions. - session.arrays().register(DecimalBytePartsPlugin); + session.arrays().register(DecimalByteParts); compute::kernel::initialize(session); session.aggregate_fns().register_aggregate_kernel( diff --git a/encodings/fastlanes/src/bitpacking/plugin.rs b/encodings/fastlanes/src/bitpacking/plugin.rs index 3ff07db7e5c..efb6567d611 100644 --- a/encodings/fastlanes/src/bitpacking/plugin.rs +++ b/encodings/fastlanes/src/bitpacking/plugin.rs @@ -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; @@ -56,16 +55,11 @@ impl ArrayPlugin for BitPackedPatchedPlugin { "BitPacked plugin does not recognize serialized ID {}", parts.serialized_id, ); - let bitpacked = Array::::try_from_parts(ArrayVTable::deserialize( - &BitPacked, - parts.dtype, - parts.len, - parts.metadata, - parts.buffers, - parts.children, - session, - )?) - .map_err(|_| vortex_err!("BitPacked plugin should only deserialize fastlanes.bitpacked"))?; + let bitpacked = ArrayPlugin::deserialize(&BitPacked, parts, session)? + .try_downcast::() + .map_err(|_| { + vortex_err!("BitPacked plugin should only deserialize fastlanes.bitpacked") + })?; // Create a new BitPackedArray without the interior patches installed. let Some(patches) = bitpacked.patches() else { diff --git a/encodings/fastlanes/src/bitpacking/vtable/mod.rs b/encodings/fastlanes/src/bitpacking/vtable/mod.rs index 68fbf1b41d3..61525c6b13d 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/mod.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/mod.rs @@ -4,7 +4,6 @@ use std::hash::Hash; use std::hash::Hasher; -use prost::Message; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -27,16 +26,13 @@ use vortex_array::patches::PatchesData; use vortex_array::patches::PatchesMetadata; use vortex_array::require_patches; use vortex_array::require_validity; -use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::child_to_validity; use vortex_array::vtable::validity_to_child; 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; @@ -52,6 +48,7 @@ use crate::bitpacking::array::PATCH_SLOTS; use crate::bitpacking::vtable::rules::RULES; mod kernels; mod operations; +mod plugin; mod rules; mod validity; @@ -160,101 +157,6 @@ impl VTable for BitPacked { ) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - BitPackedMetadata { - bit_width: array.bit_width() as u32, - offset: array.offset() 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> { - let metadata = BitPackedMetadata::decode(metadata)?; - if buffers.len() != 1 { - vortex_bail!("Expected 1 buffer, got {}", buffers.len()); - } - let packed = buffers[0].clone(); - - let load_validity = |child_idx: usize| { - if children.len() == child_idx { - Ok(Validity::from(dtype.nullability())) - } else if children.len() == child_idx + 1 { - let validity = children.get(child_idx, &Validity::DTYPE, len)?; - Ok(Validity::Array(validity)) - } else { - vortex_bail!( - "Expected {} or {} children, got {}", - child_idx, - child_idx + 1, - children.len() - ); - } - }; - - let validity_idx = match &metadata.patches { - None => 0, - Some(patches_meta) if patches_meta.chunk_offsets_dtype()?.is_some() => 3, - Some(_) => 2, - }; - - let validity = load_validity(validity_idx)?; - - let patches = metadata - .patches - .map(|p| { - let indices = children.get(0, &p.indices_dtype()?, p.len()?)?; - let values = children.get(1, dtype, p.len()?)?; - let chunk_offsets = p - .chunk_offsets_dtype()? - .map(|dtype| children.get(2, &dtype, p.chunk_offsets_len() as usize)) - .transpose()?; - - Patches::new(len, p.offset()?, indices, values, chunk_offsets) - }) - .transpose()?; - - let slots = { - let mut s = ArraySlots::with_capacity(4); - PatchesData::push_slots(&mut s, patches.as_ref()); - s.push(validity_to_child(&validity, len)); - s - }; - let data = BitPackedData::try_new( - packed, - patches, - u8::try_from(metadata.bit_width).map_err(|_| { - vortex_err!( - "BitPackedMetadata bit_width {} does not fit in u8", - metadata.bit_width - ) - })?, - u16::try_from(metadata.offset).map_err(|_| { - vortex_err!( - "BitPackedMetadata offset {} does not fit in u16", - metadata.offset - ) - })?, - )?; - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, diff --git a/encodings/fastlanes/src/bitpacking/vtable/plugin.rs b/encodings/fastlanes/src/bitpacking/vtable/plugin.rs new file mode 100644 index 00000000000..f47512bb9e9 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking/vtable/plugin.rs @@ -0,0 +1,147 @@ +// 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::ArraySlots; +use vortex_array::IntoArray; +use vortex_array::patches::Patches; +use vortex_array::patches::PatchesData; +use vortex_array::validity::Validity; +use vortex_array::vtable::VTable; +use vortex_array::vtable::validity_to_child; +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::BitPacked; +use super::BitPackedMetadata; +use crate::BitPackedArrayExt; +use crate::BitPackedData; + +impl ArrayPlugin for BitPacked { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_opt::().ok_or_else(|| { + vortex_err!("BitPacked plugin cannot serialize {}", array.encoding_id()) + })?; + let metadata = BitPackedMetadata { + bit_width: view.bit_width() as u32, + offset: view.offset() 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 { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "BitPacked plugin does not recognize serialized ID {serialized_id}" + ); + let metadata = BitPackedMetadata::decode(metadata)?; + if buffers.len() != 1 { + vortex_bail!("Expected 1 buffer, got {}", buffers.len()); + } + let packed = buffers[0].clone(); + + let load_validity = |child_idx: usize| { + if children.len() == child_idx { + Ok(Validity::from(dtype.nullability())) + } else if children.len() == child_idx + 1 { + let validity = children.get(child_idx, &Validity::DTYPE, len)?; + Ok(Validity::Array(validity)) + } else { + vortex_bail!( + "Expected {} or {} children, got {}", + child_idx, + child_idx + 1, + children.len() + ); + } + }; + + let validity_idx = match &metadata.patches { + None => 0, + Some(patches_meta) if patches_meta.chunk_offsets_dtype()?.is_some() => 3, + Some(_) => 2, + }; + + let validity = load_validity(validity_idx)?; + + let patches = metadata + .patches + .map(|p| { + let indices = children.get(0, &p.indices_dtype()?, p.len()?)?; + let values = children.get(1, dtype, p.len()?)?; + let chunk_offsets = p + .chunk_offsets_dtype()? + .map(|dtype| children.get(2, &dtype, p.chunk_offsets_len() as usize)) + .transpose()?; + + Patches::new(len, p.offset()?, indices, values, chunk_offsets) + }) + .transpose()?; + + let slots = { + let mut s = ArraySlots::with_capacity(4); + PatchesData::push_slots(&mut s, patches.as_ref()); + s.push(validity_to_child(&validity, len)); + s + }; + let data = BitPackedData::try_new( + packed, + patches, + u8::try_from(metadata.bit_width).map_err(|_| { + vortex_err!( + "BitPackedMetadata bit_width {} does not fit in u8", + metadata.bit_width + ) + })?, + u16::try_from(metadata.offset).map_err(|_| { + vortex_err!( + "BitPackedMetadata offset {} does not fit in u16", + metadata.offset + ) + })?, + )?; + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/encodings/fastlanes/src/delta/vtable/mod.rs b/encodings/fastlanes/src/delta/vtable/mod.rs index 1f2e3823689..535f0632b7a 100644 --- a/encodings/fastlanes/src/delta/vtable/mod.rs +++ b/encodings/fastlanes/src/delta/vtable/mod.rs @@ -4,7 +4,6 @@ use std::hash::Hash; use std::hash::Hasher; -use prost::Message; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -19,19 +18,13 @@ use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; -use vortex_array::dtype::PType; -use vortex_array::serde::ArrayChildren; use vortex_array::vtable::VTable; use vortex_error::VortexResult; 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::DeltaData; -use crate::delta::array::DeltaArrayExt; -use crate::delta::array::DeltaArraySlotsExt; use crate::delta::array::DeltaSlots; use crate::delta::array::DeltaSlotsView; use crate::delta::array::delta_decompress::delta_decompress; @@ -39,6 +32,7 @@ use crate::delta::array::lane_count; use crate::delta_compress; mod operations; +mod plugin; mod rules; mod slice; mod validity; @@ -127,57 +121,6 @@ impl VTable for Delta { DeltaSlots::NAMES[idx].to_string() } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - DeltaMetadata { - deltas_len: array.deltas().len() as u64, - offset: array.offset() as u32, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_ensure!( - buffers.is_empty(), - "DeltaArray expects 0 buffers, got {}", - buffers.len() - ); - vortex_ensure!( - children.len() == 2, - "DeltaArray expects 2 children, got {}", - children.len() - ); - let metadata = DeltaMetadata::decode(metadata)?; - let ptype = PType::try_from(dtype)?; - let lanes = lane_count(ptype); - - // Compute the length of the bases array - let deltas_len = usize::try_from(metadata.deltas_len) - .map_err(|_| vortex_err!("deltas_len {} overflowed usize", metadata.deltas_len))?; - let num_chunks = deltas_len / 1024; - let remainder_base_size = if deltas_len % 1024 > 0 { 1 } else { 0 }; - let bases_len = num_chunks * lanes + remainder_base_size; - - let bases = children.get(0, dtype, bases_len)?; - let deltas = children.get(1, dtype, deltas_len)?; - - let data = DeltaData::try_new(metadata.offset as usize)?; - let slots = DeltaSlots { bases, deltas }.into_slots(); - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done( delta_decompress(&array, ctx)?.into_array(), diff --git a/encodings/fastlanes/src/delta/vtable/plugin.rs b/encodings/fastlanes/src/delta/vtable/plugin.rs new file mode 100644 index 00000000000..0bf5e4798a4 --- /dev/null +++ b/encodings/fastlanes/src/delta/vtable/plugin.rs @@ -0,0 +1,101 @@ +// 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::PType; +use vortex_array::vtable::VTable; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use super::Delta; +use super::DeltaMetadata; +use crate::DeltaData; +use crate::delta::array::DeltaArrayExt; +use crate::delta::array::DeltaArraySlotsExt; +use crate::delta::array::DeltaSlots; +use crate::delta::array::lane_count; + +impl ArrayPlugin for Delta { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array + .as_opt::() + .ok_or_else(|| vortex_err!("Delta plugin cannot serialize {}", array.encoding_id()))?; + let metadata = DeltaMetadata { + deltas_len: view.deltas().len() as u64, + offset: view.offset() as u32, + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Delta plugin does not recognize serialized ID {serialized_id}" + ); + vortex_ensure!( + buffers.is_empty(), + "DeltaArray expects 0 buffers, got {}", + buffers.len() + ); + vortex_ensure!( + children.len() == 2, + "DeltaArray expects 2 children, got {}", + children.len() + ); + let metadata = DeltaMetadata::decode(metadata)?; + let ptype = PType::try_from(dtype)?; + let lanes = lane_count(ptype); + + // Compute the length of the bases array + let deltas_len = usize::try_from(metadata.deltas_len) + .map_err(|_| vortex_err!("deltas_len {} overflowed usize", metadata.deltas_len))?; + let num_chunks = deltas_len / 1024; + let remainder_base_size = if deltas_len % 1024 > 0 { 1 } else { 0 }; + let bases_len = num_chunks * lanes + remainder_base_size; + + let bases = children.get(0, dtype, bases_len)?; + let deltas = children.get(1, dtype, deltas_len)?; + + let data = DeltaData::try_new(metadata.offset as usize)?; + let slots = DeltaSlots { bases, deltas }.into_slots(); + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/encodings/fastlanes/src/for/vtable/mod.rs b/encodings/fastlanes/src/for/vtable/mod.rs index 452abcd4562..49f6d349e69 100644 --- a/encodings/fastlanes/src/for/vtable/mod.rs +++ b/encodings/fastlanes/src/for/vtable/mod.rs @@ -20,20 +20,16 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::scalar::Scalar; -use vortex_array::scalar::ScalarValue; -use vortex_array::serde::ArrayChildren; use vortex_array::smallvec::smallvec; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTableFromChild; 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::FoRData; -use crate::r#for::array::FoRArrayExt; use crate::r#for::array::FoRSlots; use crate::r#for::array::FoRSlotsView; use crate::r#for::array::for_decompress::decompress; @@ -41,6 +37,7 @@ use crate::r#for::vtable::rules::PARENT_RULES; mod kernels; mod operations; +mod plugin; mod rules; mod slice; mod validity; @@ -110,46 +107,6 @@ impl VTable for FoR { FoRSlots::NAMES[idx].to_string() } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - // Note that we **only** serialize the optional scalar value (not including the dtype). - Ok(Some(ScalarValue::to_proto_bytes( - array.reference_scalar().value(), - ))) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - vortex_ensure!( - buffers.is_empty(), - "FoRArray expects 0 buffers, got {}", - buffers.len() - ); - if children.len() != 1 { - vortex_bail!( - "Expected 1 child for FoR encoding, found {}", - children.len() - ) - } - - let scalar_value = ScalarValue::from_proto_bytes(metadata, dtype, session)?; - let reference = Scalar::try_new(dtype.clone(), scalar_value)?; - let encoded = children.get(0, dtype, len)?; - let slots = smallvec![Some(encoded)]; - - let data = FoRData::try_new(reference)?; - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, diff --git a/encodings/fastlanes/src/for/vtable/plugin.rs b/encodings/fastlanes/src/for/vtable/plugin.rs new file mode 100644 index 00000000000..b550d40d79a --- /dev/null +++ b/encodings/fastlanes/src/for/vtable/plugin.rs @@ -0,0 +1,87 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +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::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_array::smallvec::smallvec; +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::FoR; +use crate::FoRData; +use crate::r#for::array::FoRArrayExt; + +impl ArrayPlugin for FoR { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array + .as_opt::() + .ok_or_else(|| vortex_err!("FoR plugin cannot serialize {}", array.encoding_id()))?; + // Note that we **only** serialize the optional scalar value (not including the dtype). + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + ScalarValue::to_proto_bytes(view.reference_scalar().value()), + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "FoR plugin does not recognize serialized ID {serialized_id}" + ); + vortex_ensure!( + buffers.is_empty(), + "FoRArray expects 0 buffers, got {}", + buffers.len() + ); + if children.len() != 1 { + vortex_bail!( + "Expected 1 child for FoR encoding, found {}", + children.len() + ) + } + + let scalar_value = ScalarValue::from_proto_bytes(metadata, dtype, session)?; + let reference = Scalar::try_new(dtype.clone(), scalar_value)?; + let encoded = children.get(0, dtype, len)?; + let slots = smallvec![Some(encoded)]; + + let data = FoRData::try_new(reference)?; + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/encodings/fastlanes/src/lib.rs b/encodings/fastlanes/src/lib.rs index 43d83c6fc7f..be1b9e34994 100644 --- a/encodings/fastlanes/src/lib.rs +++ b/encodings/fastlanes/src/lib.rs @@ -89,7 +89,6 @@ pub fn initialize(session: &VortexSession) { session.arrays().register(Delta); session.arrays().register(FoR); session.arrays().register(RLE); - session.arrays().register(TransposedBool); bitpacking::initialize(session); r#for::initialize(session); rle::initialize(session); diff --git a/encodings/fastlanes/src/rle/vtable/mod.rs b/encodings/fastlanes/src/rle/vtable/mod.rs index a3f6f01fa5a..0ecfb07be33 100644 --- a/encodings/fastlanes/src/rle/vtable/mod.rs +++ b/encodings/fastlanes/src/rle/vtable/mod.rs @@ -4,7 +4,6 @@ use std::hash::Hash; use std::hash::Hasher; -use prost::Message; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -21,23 +20,20 @@ use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::serde::ArrayChildren; use vortex_array::vtable::VTable; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::RLEData; -use crate::rle::array::RLEArrayExt; -use crate::rle::array::RLEArraySlotsExt; use crate::rle::array::RLESlots; use crate::rle::array::RLESlotsView; use crate::rle::array::rle_decompress::rle_decompress; use crate::rle::vtable::rules::RULES; mod operations; +mod plugin; mod rules; mod validity; @@ -133,70 +129,6 @@ impl VTable for RLE { RLESlots::NAMES[idx].to_string() } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - RLEMetadata { - values_len: array.values().len() as u64, - indices_len: array.indices().len() as u64, - indices_ptype: PType::try_from(array.indices().dtype())? as i32, - values_idx_offsets_len: array.values_idx_offsets().len() as u64, - values_idx_offsets_ptype: PType::try_from(array.values_idx_offsets().dtype())? - as i32, - offset: array.offset() as u64, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_ensure!( - buffers.is_empty(), - "RLEArray expects 0 buffers, got {}", - buffers.len() - ); - let metadata = RLEMetadata::decode(metadata)?; - let values = children.get( - 0, - &DType::Primitive(dtype.as_ptype(), Nullability::NonNullable), - usize::try_from(metadata.values_len)?, - )?; - - let indices = children.get( - 1, - &DType::Primitive(metadata.indices_ptype(), dtype.nullability()), - usize::try_from(metadata.indices_len)?, - )?; - - let values_idx_offsets = children.get( - 2, - &DType::Primitive( - metadata.values_idx_offsets_ptype(), - Nullability::NonNullable, - ), - usize::try_from(metadata.values_idx_offsets_len)?, - )?; - - let slots = RLESlots { - values, - indices, - values_idx_offsets, - } - .into_slots(); - let data = RLEData::try_new(metadata.offset as usize)?; - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done( rle_decompress(&array, ctx)?.into_array(), diff --git a/encodings/fastlanes/src/rle/vtable/plugin.rs b/encodings/fastlanes/src/rle/vtable/plugin.rs new file mode 100644 index 00000000000..aeea142141b --- /dev/null +++ b/encodings/fastlanes/src/rle/vtable/plugin.rs @@ -0,0 +1,114 @@ +// 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::Nullability; +use vortex_array::dtype::PType; +use vortex_array::vtable::VTable; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use super::RLE; +use super::RLEMetadata; +use crate::RLEData; +use crate::rle::array::RLEArrayExt; +use crate::rle::array::RLEArraySlotsExt; +use crate::rle::array::RLESlots; + +impl ArrayPlugin for RLE { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array + .as_opt::() + .ok_or_else(|| vortex_err!("RLE plugin cannot serialize {}", array.encoding_id()))?; + let metadata = RLEMetadata { + values_len: view.values().len() as u64, + indices_len: view.indices().len() as u64, + indices_ptype: PType::try_from(view.indices().dtype())? as i32, + values_idx_offsets_len: view.values_idx_offsets().len() as u64, + values_idx_offsets_ptype: PType::try_from(view.values_idx_offsets().dtype())? as i32, + offset: view.offset() as u64, + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "RLE plugin does not recognize serialized ID {serialized_id}" + ); + vortex_ensure!( + buffers.is_empty(), + "RLEArray expects 0 buffers, got {}", + buffers.len() + ); + let metadata = RLEMetadata::decode(metadata)?; + let values = children.get( + 0, + &DType::Primitive(dtype.as_ptype(), Nullability::NonNullable), + usize::try_from(metadata.values_len)?, + )?; + + let indices = children.get( + 1, + &DType::Primitive(metadata.indices_ptype(), dtype.nullability()), + usize::try_from(metadata.indices_len)?, + )?; + + let values_idx_offsets = children.get( + 2, + &DType::Primitive( + metadata.values_idx_offsets_ptype(), + Nullability::NonNullable, + ), + usize::try_from(metadata.values_idx_offsets_len)?, + )?; + + let slots = RLESlots { + values, + indices, + values_idx_offsets, + } + .into_slots(); + let data = RLEData::try_new(metadata.offset as usize)?; + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/encodings/fastlanes/src/transposed_bool.rs b/encodings/fastlanes/src/transposed_bool.rs index 92da3b3b79c..3472f3286f5 100644 --- a/encodings/fastlanes/src/transposed_bool.rs +++ b/encodings/fastlanes/src/transposed_bool.rs @@ -27,7 +27,6 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar::Scalar; -use vortex_array::serde::ArrayChildren; use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; use vortex_array::vtable::OperationsVTable; @@ -35,10 +34,8 @@ use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; use vortex_error::VortexExpect; 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::FL_CHUNK_SIZE; @@ -206,25 +203,6 @@ impl VTable for TransposedBool { vortex_array::vtable::with_empty_buffers(self, array, buffers) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - vortex_bail!("Cannot serialise TransposedBoolArray"); - } - - fn deserialize( - &self, - _dtype: &DType, - _len: usize, - _metadata: &[u8], - _buffers: &[BufferHandle], - _children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_bail!("Cannot deserialise TransposedBoolArray"); - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { match idx { TRANSPOSED_SLOT => "transposed".to_string(), diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index 86b8fd0c34d..b9e03a13663 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -13,7 +13,6 @@ use fsst::Compressor; use fsst::Decompressor; use fsst::Symbol; use num_traits::AsPrimitive; -use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -57,7 +56,6 @@ 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::canonical::FSST_DECODE_SLACK; @@ -66,6 +64,8 @@ use crate::canonical::canonicalize_fsst; use crate::canonical::fsst_decode_bytes; use crate::rules::RULES; +mod plugin; + /// A [`FSST`]-encoded Vortex array. pub type FSSTArray = Array; @@ -188,128 +188,6 @@ impl VTable for FSST { ) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - let codes_offsets = array.codes_offsets(); - Ok(Some( - FSSTMetadata { - uncompressed_lengths_ptype: array.uncompressed_lengths().dtype().as_ptype().into(), - codes_offsets_ptype: codes_offsets.dtype().as_ptype().into(), - } - .encode_to_vec(), - )) - } - - /// Deserializes an FSST array from its serialized components. - /// - /// Supports two serialization formats: - /// - /// ## Legacy format (2 buffers, 2 children) - /// - /// The original FSST layout stored the compressed codes as a full `VarBinArray` child. - /// - **Buffers**: `[symbols, symbol_lengths]` - /// - **Children**: `[codes (VarBinArray), uncompressed_lengths (Primitive)]` - /// - /// The codes VarBinArray child is decomposed: its bytes become the `codes_bytes` buffer, - /// and its offsets/validity are extracted into slots. - /// See `FSST::deserialize_legacy`. - /// - /// ## Current format (3 buffers, 2-3 children) - /// - /// The current layout stores the compressed bytes as a raw buffer alongside the symbol - /// table, with offsets and validity as separate children. - /// - **Buffers**: `[symbols, symbol_lengths, compressed_codes_bytes]` - /// - **Children**: `[uncompressed_lengths, codes_offsets, (optional) codes_validity]` - /// - /// The `codes_bytes` buffer is stored directly in `FSSTData`. A `VarBinArray` for the - /// codes can be reconstructed on demand via [`FSSTArrayExt::codes()`] using the bytes - /// from `FSSTData` combined with offsets and validity from the array's slots. - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - let metadata = FSSTMetadata::decode(metadata)?; - let symbols = Buffer::::from_byte_buffer(buffers[0].clone().try_to_host_sync()?); - let symbol_lengths = Buffer::::from_byte_buffer(buffers[1].clone().try_to_host_sync()?); - - let mut ctx = session.create_execution_ctx(); - if buffers.len() == 2 { - return Self::deserialize_legacy( - self, - dtype, - len, - &metadata, - &symbols, - &symbol_lengths, - children, - &mut ctx, - ); - } - - if buffers.len() == 3 { - let uncompressed_lengths = children.get( - 0, - &DType::Primitive( - metadata.get_uncompressed_lengths_ptype()?, - Nullability::NonNullable, - ), - len, - )?; - - let codes_bytes = buffers[2].clone(); - let codes_offsets = children.get( - 1, - &DType::Primitive( - PType::try_from(metadata.codes_offsets_ptype)?, - Nullability::NonNullable, - ), - // VarBin offsets are len + 1 - len + 1, - )?; - - let codes_validity = if children.len() == 2 { - Validity::from(dtype.nullability()) - } else if children.len() == 3 { - let validity = children.get(2, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("Expected 2 or 3 children, got {}", children.len()); - }; - - FSSTData::validate_parts( - symbols.as_slice(), - symbol_lengths.as_slice(), - &codes_bytes, - &codes_offsets, - dtype.nullability(), - &uncompressed_lengths, - dtype, - len, - &mut ctx, - )?; - let slots = FSSTSlots { - uncompressed_lengths, - codes_offsets, - codes_validity: validity_to_child(&codes_validity, len), - } - .into_slots(); - let data = FSSTData::try_new(symbols, symbol_lengths, codes_bytes, len)?; - return Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)); - } - - vortex_bail!( - "InvalidArgument: Expected 2 or 3 buffers, got {}", - buffers.len() - ); - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { FSSTSlots::NAMES[idx].to_string() } diff --git a/encodings/fsst/src/array/plugin.rs b/encodings/fsst/src/array/plugin.rs new file mode 100644 index 00000000000..70ed6deef75 --- /dev/null +++ b/encodings/fsst/src/array/plugin.rs @@ -0,0 +1,181 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use fsst::Symbol; +use prost::Message as _; +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::VortexSessionExecute; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::validity::Validity; +use vortex_array::vtable::VTable; +use vortex_array::vtable::validity_to_child; +use vortex_buffer::Buffer; +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::FSST; +use super::FSSTArraySlotsExt; +use super::FSSTData; +use super::FSSTMetadata; +use super::FSSTSlots; + +impl ArrayPlugin for FSST { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array + .as_opt::() + .ok_or_else(|| vortex_err!("FSST plugin cannot serialize {}", array.encoding_id()))?; + let codes_offsets = view.codes_offsets(); + let metadata = FSSTMetadata { + uncompressed_lengths_ptype: view.uncompressed_lengths().dtype().as_ptype().into(), + codes_offsets_ptype: codes_offsets.dtype().as_ptype().into(), + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata, + ))) + } + + /// Deserializes an FSST array from its serialized components. + /// + /// Supports two serialization formats: + /// + /// ## Legacy format (2 buffers, 2 children) + /// + /// The original FSST layout stored the compressed codes as a full `VarBinArray` child. + /// - **Buffers**: `[symbols, symbol_lengths]` + /// - **Children**: `[codes (VarBinArray), uncompressed_lengths (Primitive)]` + /// + /// The codes VarBinArray child is decomposed: its bytes become the `codes_bytes` buffer, + /// and its offsets/validity are extracted into slots. + /// See `FSST::deserialize_legacy`. + /// + /// ## Current format (3 buffers, 2-3 children) + /// + /// The current layout stores the compressed bytes as a raw buffer alongside the symbol + /// table, with offsets and validity as separate children. + /// - **Buffers**: `[symbols, symbol_lengths, compressed_codes_bytes]` + /// - **Children**: `[uncompressed_lengths, codes_offsets, (optional) codes_validity]` + /// + /// The `codes_bytes` buffer is stored directly in `FSSTData`. A `VarBinArray` for the + /// codes can be reconstructed on demand via [`FSSTArrayExt::codes()`] using the bytes + /// from `FSSTData` combined with offsets and validity from the array's slots. + /// + /// [`FSSTArrayExt::codes()`]: super::FSSTArrayExt::codes + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "FSST plugin does not recognize serialized ID {serialized_id}" + ); + let metadata = FSSTMetadata::decode(metadata)?; + let symbols = Buffer::::from_byte_buffer(buffers[0].clone().try_to_host_sync()?); + let symbol_lengths = Buffer::::from_byte_buffer(buffers[1].clone().try_to_host_sync()?); + + let mut ctx = session.create_execution_ctx(); + if buffers.len() == 2 { + return Ok(Array::try_from_parts(Self::deserialize_legacy( + self, + dtype, + len, + &metadata, + &symbols, + &symbol_lengths, + children, + &mut ctx, + )?)? + .into_array()); + } + + if buffers.len() == 3 { + let uncompressed_lengths = children.get( + 0, + &DType::Primitive( + metadata.get_uncompressed_lengths_ptype()?, + Nullability::NonNullable, + ), + len, + )?; + + let codes_bytes = buffers[2].clone(); + let codes_offsets = children.get( + 1, + &DType::Primitive( + PType::try_from(metadata.codes_offsets_ptype)?, + Nullability::NonNullable, + ), + // VarBin offsets are len + 1 + len + 1, + )?; + + let codes_validity = if children.len() == 2 { + Validity::from(dtype.nullability()) + } else if children.len() == 3 { + let validity = children.get(2, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + vortex_bail!("Expected 2 or 3 children, got {}", children.len()); + }; + + FSSTData::validate_parts( + symbols.as_slice(), + symbol_lengths.as_slice(), + &codes_bytes, + &codes_offsets, + dtype.nullability(), + &uncompressed_lengths, + dtype, + len, + &mut ctx, + )?; + let slots = FSSTSlots { + uncompressed_lengths, + codes_offsets, + codes_validity: validity_to_child(&codes_validity, len), + } + .into_slots(); + let data = FSSTData::try_new(symbols, symbol_lengths, codes_bytes, len)?; + return Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()); + } + + vortex_bail!( + "InvalidArgument: Expected 2 or 3 buffers, got {}", + buffers.len() + ); + } +} diff --git a/encodings/onpair/src/array.rs b/encodings/onpair/src/array.rs index 912fc27c278..51e38b7284f 100644 --- a/encodings/onpair/src/array.rs +++ b/encodings/onpair/src/array.rs @@ -14,7 +14,6 @@ use onpair::CompactDictionary; use onpair::CompactDictionaryView; use onpair::Dictionary; use onpair::DictionaryStorage; -use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -31,12 +30,10 @@ use vortex_array::builders::ArrayBuilder; use vortex_array::builders::VarBinBuilder; use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; use vortex_array::dtype::OffsetBuilderPType; use vortex_array::dtype::PType; use vortex_array::match_each_integer_ptype; use vortex_array::match_each_varbin_builder; -use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; @@ -49,7 +46,6 @@ 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::canonical::OnPairDecodePlan; @@ -58,6 +54,8 @@ use crate::canonical::onpair_decode_bytes; use crate::decode::collect_widened; use crate::rules::RULES; +mod plugin; + /// An [`OnPair`]-encoded Vortex array. pub type OnPairArray = Array; @@ -507,98 +505,6 @@ impl VTable for OnPair { ) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - let dict_size = u32::try_from(array.dict_offsets().len().saturating_sub(1)) - .map_err(|_| vortex_err!("OnPair dict_size exceeds u32"))?; - let codes_len = array.codes().len() as u64; - Ok(Some( - OnPairMetadata { - uncompressed_lengths_ptype: array.uncompressed_lengths().dtype().as_ptype().into(), - dict_size, - codes_len, - dict_offsets_ptype: array.dict_offsets().dtype().as_ptype().into(), - codes_ptype: array.codes().dtype().as_ptype().into(), - codes_offsets_ptype: array.codes_offsets().dtype().as_ptype().into(), - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - if buffers.len() != 1 { - vortex_bail!(InvalidArgument: "Expected 1 buffer, got {}", buffers.len()); - } - let metadata = OnPairMetadata::decode(metadata)?; - let uncompressed_ptype = metadata.get_uncompressed_lengths_ptype()?; - - // Slot children do not persist their own lengths, so metadata records - // the dictionary and code-stream sizes needed to deserialize them. - let dict_offsets_len = metadata.dict_size as usize + 1; - let codes_len = usize::try_from(metadata.codes_len) - .map_err(|_| vortex_err!("codes_len {} overflows usize", metadata.codes_len))?; - // The cascading compressor may have narrowed any of these integer - // children to a tighter ptype; the recorded ptype tells the framework - // exactly which dtype to materialise as. - let dict_offsets_ptype = PType::try_from(metadata.dict_offsets_ptype).map_err(|_| { - vortex_err!("invalid dict_offsets_ptype {}", metadata.dict_offsets_ptype) - })?; - let codes_ptype = PType::try_from(metadata.codes_ptype) - .map_err(|_| vortex_err!("invalid codes_ptype {}", metadata.codes_ptype))?; - let codes_offsets_ptype = PType::try_from(metadata.codes_offsets_ptype).map_err(|_| { - vortex_err!( - "invalid codes_offsets_ptype {}", - metadata.codes_offsets_ptype - ) - })?; - let dict_offsets = children.get( - 0, - &DType::Primitive(dict_offsets_ptype, Nullability::NonNullable), - dict_offsets_len, - )?; - let codes = children.get( - 1, - &DType::Primitive(codes_ptype, Nullability::NonNullable), - codes_len, - )?; - let codes_offsets = children.get( - 2, - &DType::Primitive(codes_offsets_ptype, Nullability::NonNullable), - len + 1, - )?; - let uncompressed_lengths = children.get( - 3, - &DType::Primitive(uncompressed_ptype, Nullability::NonNullable), - len, - )?; - let validity = match children.len() { - 4 => Validity::from(dtype.nullability()), - 5 => Validity::Array(children.get(4, &Validity::DTYPE, len)?), - other => vortex_bail!(InvalidArgument: "Expected 4 or 5 children, got {other}"), - }; - - let data = OnPairData::new(buffers[0].clone()); - let slots = OnPairSlots { - dict_offsets, - codes, - codes_offsets, - uncompressed_lengths, - validity: validity_to_child(&validity, len), - } - .into_slots(); - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { OnPairSlots::NAMES[idx].to_string() } diff --git a/encodings/onpair/src/array/plugin.rs b/encodings/onpair/src/array/plugin.rs new file mode 100644 index 00000000000..981ece3c045 --- /dev/null +++ b/encodings/onpair/src/array/plugin.rs @@ -0,0 +1,145 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use prost::Message as _; +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::Nullability; +use vortex_array::dtype::PType; +use vortex_array::validity::Validity; +use vortex_array::vtable::VTable; +use vortex_array::vtable::validity_to_child; +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::OnPair; +use super::OnPairArraySlotsExt; +use super::OnPairData; +use super::OnPairMetadata; +use super::OnPairSlots; + +impl ArrayPlugin for OnPair { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array + .as_opt::() + .ok_or_else(|| vortex_err!("OnPair plugin cannot serialize {}", array.encoding_id()))?; + let dict_size = u32::try_from(view.dict_offsets().len().saturating_sub(1)) + .map_err(|_| vortex_err!("OnPair dict_size exceeds u32"))?; + let codes_len = view.codes().len() as u64; + let metadata = OnPairMetadata { + uncompressed_lengths_ptype: view.uncompressed_lengths().dtype().as_ptype().into(), + dict_size, + codes_len, + dict_offsets_ptype: view.dict_offsets().dtype().as_ptype().into(), + codes_ptype: view.codes().dtype().as_ptype().into(), + codes_offsets_ptype: view.codes_offsets().dtype().as_ptype().into(), + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "OnPair plugin does not recognize serialized ID {serialized_id}" + ); + if buffers.len() != 1 { + vortex_bail!(InvalidArgument: "Expected 1 buffer, got {}", buffers.len()); + } + let metadata = OnPairMetadata::decode(metadata)?; + let uncompressed_ptype = metadata.get_uncompressed_lengths_ptype()?; + + // Slot children do not persist their own lengths, so metadata records + // the dictionary and code-stream sizes needed to deserialize them. + let dict_offsets_len = metadata.dict_size as usize + 1; + let codes_len = usize::try_from(metadata.codes_len) + .map_err(|_| vortex_err!("codes_len {} overflows usize", metadata.codes_len))?; + // The cascading compressor may have narrowed any of these integer + // children to a tighter ptype; the recorded ptype tells the framework + // exactly which dtype to materialise as. + let dict_offsets_ptype = PType::try_from(metadata.dict_offsets_ptype).map_err(|_| { + vortex_err!("invalid dict_offsets_ptype {}", metadata.dict_offsets_ptype) + })?; + let codes_ptype = PType::try_from(metadata.codes_ptype) + .map_err(|_| vortex_err!("invalid codes_ptype {}", metadata.codes_ptype))?; + let codes_offsets_ptype = PType::try_from(metadata.codes_offsets_ptype).map_err(|_| { + vortex_err!( + "invalid codes_offsets_ptype {}", + metadata.codes_offsets_ptype + ) + })?; + let dict_offsets = children.get( + 0, + &DType::Primitive(dict_offsets_ptype, Nullability::NonNullable), + dict_offsets_len, + )?; + let codes = children.get( + 1, + &DType::Primitive(codes_ptype, Nullability::NonNullable), + codes_len, + )?; + let codes_offsets = children.get( + 2, + &DType::Primitive(codes_offsets_ptype, Nullability::NonNullable), + len + 1, + )?; + let uncompressed_lengths = children.get( + 3, + &DType::Primitive(uncompressed_ptype, Nullability::NonNullable), + len, + )?; + let validity = match children.len() { + 4 => Validity::from(dtype.nullability()), + 5 => Validity::Array(children.get(4, &Validity::DTYPE, len)?), + other => vortex_bail!(InvalidArgument: "Expected 4 or 5 children, got {other}"), + }; + + let data = OnPairData::new(buffers[0].clone()); + let slots = OnPairSlots { + dict_offsets, + codes, + codes_offsets, + uncompressed_lengths, + validity: validity_to_child(&validity, len), + } + .into_slots(); + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/encodings/parquet-variant/src/vtable.rs b/encodings/parquet-variant/src/vtable.rs index 894a732e4c1..d9432414114 100644 --- a/encodings/parquet-variant/src/vtable.rs +++ b/encodings/parquet-variant/src/vtable.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use prost::Message; use vortex_array::Array; use vortex_array::ArrayId; use vortex_array::ArrayParts; @@ -16,17 +15,13 @@ use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::proto::dtype as pb; -use vortex_array::serde::ArrayChildren; -use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::child_to_validity; -use vortex_array::vtable::validity_to_child; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::array::ParquetVariantArraySlotsExt; @@ -34,6 +29,8 @@ use crate::array::ParquetVariantSlots; use crate::array::core_storage_without_typed_value; use crate::array::logical_shredded_from_parquet_typed_value; +mod plugin; + /// VTable for Arrow's canonical `arrow.parquet.variant` extension storage. /// /// `ParquetVariantArray` preserves semi-structured data stored as Parquet Variant values in a @@ -181,98 +178,6 @@ impl VTable for ParquetVariant { ParquetVariantSlots::NAMES[idx].to_string() } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - let typed_value_dtype = array - .typed_value() - .map(|tv| tv.dtype().try_into()) - .transpose()?; - Ok(Some( - ParquetVariantMetadataProto { - has_value: array.value().is_some(), - typed_value_dtype, - value_nullable: array.value().is_some_and(|v| v.dtype().is_nullable()), - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - vortex_ensure!( - buffers.is_empty(), - "ParquetVariantArray expects 0 buffers, got {}", - buffers.len() - ); - - let proto = ParquetVariantMetadataProto::decode(metadata)?; - let typed_value_dtype = match proto.typed_value_dtype.as_ref() { - Some(dtype) => Some(DType::from_proto(dtype, session)?), - None => None, - }; - - vortex_ensure!(matches!(dtype, DType::Variant(_)), "Expected Variant DType"); - let has_typed_value = typed_value_dtype.is_some(); - vortex_ensure!( - proto.has_value || has_typed_value, - "At least one of value or typed_value must be present" - ); - - let expected_children = 1 + proto.has_value as usize + has_typed_value as usize; - vortex_ensure!( - children.len() == expected_children || children.len() == expected_children + 1, - "Expected {} or {} children, got {}", - expected_children, - expected_children + 1, - children.len() - ); - - let (validity, mut child_idx) = if children.len() == expected_children { - (Validity::from(dtype.nullability()), 0) - } else { - (Validity::Array(children.get(0, &Validity::DTYPE, len)?), 1) - }; - let variant_metadata = - children.get(child_idx, &DType::Binary(Nullability::NonNullable), len)?; - child_idx += 1; - - let value = if proto.has_value { - let v = children.get(child_idx, &DType::Binary(proto.value_nullable.into()), len)?; - child_idx += 1; - Some(v) - } else { - None - }; - - let typed_value = if has_typed_value { - // typed_value can be any type — primitive, list, struct, etc. - let dtype = typed_value_dtype - .ok_or_else(|| vortex_err!("typed_value_dtype missing for typed_value child"))?; - let tv = children.get(child_idx, &dtype, len)?; - Some(tv) - } else { - None - }; - - let slots = ParquetVariantSlots { - validity: validity_to_child(&validity, len), - metadata: variant_metadata, - value, - typed_value, - } - .into_slots(); - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots(slots)) - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { let shredded = array .typed_value() diff --git a/encodings/parquet-variant/src/vtable/plugin.rs b/encodings/parquet-variant/src/vtable/plugin.rs new file mode 100644 index 00000000000..e0ac28040f3 --- /dev/null +++ b/encodings/parquet-variant/src/vtable/plugin.rs @@ -0,0 +1,146 @@ +// 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::EmptyArrayData; +use vortex_array::IntoArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::validity::Validity; +use vortex_array::vtable::VTable; +use vortex_array::vtable::validity_to_child; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use super::ParquetVariant; +use super::ParquetVariantMetadataProto; +use crate::array::ParquetVariantArraySlotsExt; +use crate::array::ParquetVariantSlots; + +impl ArrayPlugin for ParquetVariant { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_opt::().ok_or_else(|| { + vortex_err!( + "ParquetVariant plugin cannot serialize {}", + array.encoding_id() + ) + })?; + let typed_value_dtype = view + .typed_value() + .map(|tv| tv.dtype().try_into()) + .transpose()?; + let metadata = ParquetVariantMetadataProto { + has_value: view.value().is_some(), + typed_value_dtype, + value_nullable: view.value().is_some_and(|v| v.dtype().is_nullable()), + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "ParquetVariant plugin does not recognize serialized ID {serialized_id}" + ); + vortex_ensure!( + buffers.is_empty(), + "ParquetVariantArray expects 0 buffers, got {}", + buffers.len() + ); + + let proto = ParquetVariantMetadataProto::decode(metadata)?; + let typed_value_dtype = match proto.typed_value_dtype.as_ref() { + Some(dtype) => Some(DType::from_proto(dtype, session)?), + None => None, + }; + + vortex_ensure!(matches!(dtype, DType::Variant(_)), "Expected Variant DType"); + let has_typed_value = typed_value_dtype.is_some(); + vortex_ensure!( + proto.has_value || has_typed_value, + "At least one of value or typed_value must be present" + ); + + let expected_children = 1 + proto.has_value as usize + has_typed_value as usize; + vortex_ensure!( + children.len() == expected_children || children.len() == expected_children + 1, + "Expected {} or {} children, got {}", + expected_children, + expected_children + 1, + children.len() + ); + + let (validity, mut child_idx) = if children.len() == expected_children { + (Validity::from(dtype.nullability()), 0) + } else { + (Validity::Array(children.get(0, &Validity::DTYPE, len)?), 1) + }; + let variant_metadata = + children.get(child_idx, &DType::Binary(Nullability::NonNullable), len)?; + child_idx += 1; + + let value = if proto.has_value { + let v = children.get(child_idx, &DType::Binary(proto.value_nullable.into()), len)?; + child_idx += 1; + Some(v) + } else { + None + }; + + let typed_value = if has_typed_value { + // typed_value can be any type — primitive, list, struct, etc. + let dtype = typed_value_dtype + .ok_or_else(|| vortex_err!("typed_value_dtype missing for typed_value child"))?; + let tv = children.get(child_idx, &dtype, len)?; + Some(tv) + } else { + None + }; + + let slots = ParquetVariantSlots { + validity: validity_to_child(&validity, len), + metadata: variant_metadata, + value, + typed_value, + } + .into_slots(); + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots(slots), + )? + .into_array()) + } +} diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index 2d9d563f473..ec80b56e543 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -17,7 +17,6 @@ use pco::match_number_enum; use pco::wrapped::ChunkDecompressor; use pco::wrapped::FileCompressor; use pco::wrapped::FileDecompressor; -use prost::Message; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -39,7 +38,6 @@ use vortex_array::dtype::DType; use vortex_array::dtype::PType; use vortex_array::dtype::half; use vortex_array::scalar::Scalar; -use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::OperationsVTable; use vortex_array::vtable::VTable; @@ -50,16 +48,16 @@ use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_error::VortexError; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::PcoChunkInfo; use crate::PcoMetadata; use crate::PcoPageInfo; +mod plugin; + // Overall approach here: // Chunk the array into Pco chunks (currently using the default recommended size // for good compression), and into finer-grained Pco pages. As we go, write each @@ -190,60 +188,6 @@ impl VTable for Pco { ) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(array.metadata.clone().encode_to_vec())) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - let metadata = PcoMetadata::decode(metadata)?; - let validity = if children.is_empty() { - Validity::from(dtype.nullability()) - } else if children.len() == 1 { - let validity = children.get(0, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("PcoArray expected 0 or 1 child, got {}", children.len()); - }; - - vortex_ensure!(buffers.len() >= metadata.chunks.len()); - let chunk_metas = buffers[..metadata.chunks.len()] - .iter() - .map(|b| b.clone().try_to_host_sync()) - .collect::>>()?; - let pages = buffers[metadata.chunks.len()..] - .iter() - .map(|b| b.clone().try_to_host_sync()) - .collect::>>()?; - - let expected_n_pages = metadata - .chunks - .iter() - .map(|info| info.pages.len()) - .sum::(); - vortex_ensure!(pages.len() == expected_n_pages); - - let slots = PcoSlots { - validity: validity_to_child(&validity, len), - } - .into_slots(); - // SAFETY: `Array::try_from_parts`, which consumes these parts, validates the data before - // publishing the array. - let data = - unsafe { PcoData::new_unchecked(chunk_metas, pages, dtype.as_ptype(), metadata, len) }; - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { PcoSlots::NAMES[idx].to_string() } diff --git a/encodings/pco/src/array/plugin.rs b/encodings/pco/src/array/plugin.rs new file mode 100644 index 00000000000..2d787f4da3c --- /dev/null +++ b/encodings/pco/src/array/plugin.rs @@ -0,0 +1,104 @@ +// 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::validity::Validity; +use vortex_array::vtable::VTable; +use vortex_array::vtable::validity_to_child; +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::Pco; +use super::PcoData; +use super::PcoSlots; +use crate::PcoMetadata; + +impl ArrayPlugin for Pco { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array + .as_opt::() + .ok_or_else(|| vortex_err!("Pco plugin cannot serialize {}", array.encoding_id()))?; + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + view.metadata.clone().encode_to_vec(), + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Pco plugin does not recognize serialized ID {serialized_id}" + ); + let metadata = PcoMetadata::decode(metadata)?; + let validity = if children.is_empty() { + Validity::from(dtype.nullability()) + } else if children.len() == 1 { + let validity = children.get(0, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + vortex_bail!("PcoArray expected 0 or 1 child, got {}", children.len()); + }; + + vortex_ensure!(buffers.len() >= metadata.chunks.len()); + let chunk_metas = buffers[..metadata.chunks.len()] + .iter() + .map(|b| b.clone().try_to_host_sync()) + .collect::>>()?; + let pages = buffers[metadata.chunks.len()..] + .iter() + .map(|b| b.clone().try_to_host_sync()) + .collect::>>()?; + + let expected_n_pages = metadata + .chunks + .iter() + .map(|info| info.pages.len()) + .sum::(); + vortex_ensure!(pages.len() == expected_n_pages); + + let slots = PcoSlots { + validity: validity_to_child(&validity, len), + } + .into_slots(); + // SAFETY: `Array::try_from_parts`, which consumes these parts, validates the data before + // publishing the array. + let data = + unsafe { PcoData::new_unchecked(chunk_metas, pages, dtype.as_ptype(), metadata, len) }; + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index 3a0b10c5c4a..b4ca8aaf016 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -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; @@ -30,10 +29,8 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::legacy_session; -use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; @@ -42,7 +39,6 @@ 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::compress::runend_decode_decimal; @@ -54,6 +50,8 @@ use crate::ops::find_physical_index; use crate::ops::find_slice_end_index; use crate::rules::RULES; +mod plugin; + /// A [`RunEnd`]-encoded Vortex array. pub type RunEndArray = Array; @@ -133,42 +131,6 @@ impl VTable for RunEnd { vortex_array::vtable::with_empty_buffers(self, array, buffers) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - RunEndMetadata { - ends_ptype: PType::try_from(array.ends().dtype()) - .vortex_expect("Must be a valid PType") as i32, - num_runs: array.ends().len() as u64, - offset: array.offset() as u64, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - let metadata = RunEndMetadata::decode(metadata)?; - let ends_dtype = DType::Primitive(metadata.ends_ptype(), Nullability::NonNullable); - let runs = usize::try_from(metadata.num_runs).vortex_expect("Must be a valid usize"); - let ends = children.get(0, &ends_dtype, runs)?; - - let values = children.get(1, dtype, runs)?; - let offset = usize::try_from(metadata.offset).vortex_expect("Offset must be a valid usize"); - let slots = RunEndSlots { ends, values }.into_slots(); - let data = RunEndData::new(offset); - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { RunEndSlots::NAMES[idx].to_string() } diff --git a/encodings/runend/src/array/plugin.rs b/encodings/runend/src/array/plugin.rs new file mode 100644 index 00000000000..a28fa8d5747 --- /dev/null +++ b/encodings/runend/src/array/plugin.rs @@ -0,0 +1,88 @@ +// 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::Nullability; +use vortex_array::dtype::PType; +use vortex_array::vtable::VTable; +use vortex_error::VortexExpect as _; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use super::RunEnd; +use super::RunEndArrayExt; +use super::RunEndArraySlotsExt; +use super::RunEndData; +use super::RunEndMetadata; +use super::RunEndSlots; + +impl ArrayPlugin for RunEnd { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array + .as_opt::() + .ok_or_else(|| vortex_err!("RunEnd plugin cannot serialize {}", array.encoding_id()))?; + let metadata = RunEndMetadata { + ends_ptype: PType::try_from(view.ends().dtype()).vortex_expect("Must be a valid PType") + as i32, + num_runs: view.ends().len() as u64, + offset: view.offset() as u64, + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers: _, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "RunEnd plugin does not recognize serialized ID {serialized_id}" + ); + let metadata = RunEndMetadata::decode(metadata)?; + let ends_dtype = DType::Primitive(metadata.ends_ptype(), Nullability::NonNullable); + let runs = usize::try_from(metadata.num_runs).vortex_expect("Must be a valid usize"); + let ends = children.get(0, &ends_dtype, runs)?; + + let values = children.get(1, dtype, runs)?; + let offset = usize::try_from(metadata.offset).vortex_expect("Offset must be a valid usize"); + let slots = RunEndSlots { ends, values }.into_slots(); + let data = RunEndData::new(offset); + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/encodings/sequence/src/array.rs b/encodings/sequence/src/array.rs index f3f44e64f6e..38d6e101699 100644 --- a/encodings/sequence/src/array.rs +++ b/encodings/sequence/src/array.rs @@ -7,7 +7,6 @@ use std::hash::Hash; use std::hash::Hasher; use num_traits::AsPrimitive; -use prost::Message; use smallvec::smallvec; use vortex_array::Array; use vortex_array::ArrayEq; @@ -23,7 +22,6 @@ use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; -use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::PType; use vortex_array::expr::stats::Precision as StatPrecision; use vortex_array::expr::stats::Stat; @@ -33,7 +31,6 @@ use vortex_array::proto::scalar::ScalarValue as ProtoScalarValue; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; -use vortex_array::serde::ArrayChildren; use vortex_array::stats::StatsSet; use vortex_array::validity::Validity; use vortex_array::vtable::OperationsVTable; @@ -45,7 +42,6 @@ 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::compress::sequence_decompress; @@ -53,6 +49,8 @@ use crate::eval; use crate::eval::SequenceValue; use crate::rules::RULES; +mod plugin; + /// A [`Sequence`]-encoded Vortex array. pub type SequenceArray = Array; @@ -333,80 +331,6 @@ impl VTable for Sequence { vortex_array::vtable::with_empty_buffers(self, array, buffers) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - let metadata = SequenceMetadata { - base: Some((&array.base()).into()), - multiplier: Some((&array.multiplier()).into()), - }; - - Ok(Some(metadata.encode_to_vec())) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - vortex_ensure!( - buffers.is_empty(), - "SequenceArray expects 0 buffers, got {}", - buffers.len() - ); - vortex_ensure!( - children.is_empty(), - "SequenceArray expects 0 children, got {}", - children.len() - ); - let DType::Primitive(output_ptype, _) = dtype else { - vortex_bail!( - "only primitive dtypes are supported in SequenceArray currently, got {dtype}" - ); - }; - let metadata = SequenceMetadata::decode(metadata)?; - - let base_metadata = metadata - .base - .as_ref() - .ok_or_else(|| vortex_err!("base required"))?; - - let multiplier_metadata = metadata - .multiplier - .as_ref() - .ok_or_else(|| vortex_err!("multiplier required"))?; - - // We go via Scalar to validate that the value is valid for the ptype. - let base = Scalar::from_proto_value( - base_metadata, - &DType::Primitive(*output_ptype, NonNullable), - session, - )? - .as_primitive() - .pvalue() - .vortex_expect("sequence array base should be a non-nullable primitive"); - - // The serialized step preserves signedness independently of the output ptype. - let multiplier_ptype = SequenceData::multiplier_ptype_from_proto(multiplier_metadata)?; - let multiplier = Scalar::from_proto_value( - multiplier_metadata, - &DType::Primitive(multiplier_ptype, NonNullable), - session, - )? - .as_primitive() - .pvalue() - .vortex_expect("sequence array multiplier should be a non-nullable primitive"); - - let data = - SequenceData::try_new(base, multiplier, *output_ptype, dtype.nullability(), len)?; - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { SLOT_NAMES[idx].to_string() } diff --git a/encodings/sequence/src/array/plugin.rs b/encodings/sequence/src/array/plugin.rs new file mode 100644 index 00000000000..1d0b5e89c84 --- /dev/null +++ b/encodings/sequence/src/array/plugin.rs @@ -0,0 +1,125 @@ +// 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::Nullability::NonNullable; +use vortex_array::scalar::Scalar; +use vortex_array::vtable::VTable; +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_session::VortexSession; + +use super::Sequence; +use super::SequenceData; +use super::SequenceMetadata; + +impl ArrayPlugin for Sequence { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_opt::().ok_or_else(|| { + vortex_err!("Sequence plugin cannot serialize {}", array.encoding_id()) + })?; + let metadata = SequenceMetadata { + base: Some((&view.base()).into()), + multiplier: Some((&view.multiplier()).into()), + }; + + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata.encode_to_vec(), + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Sequence plugin does not recognize serialized ID {serialized_id}" + ); + vortex_ensure!( + buffers.is_empty(), + "SequenceArray expects 0 buffers, got {}", + buffers.len() + ); + vortex_ensure!( + children.is_empty(), + "SequenceArray expects 0 children, got {}", + children.len() + ); + let DType::Primitive(output_ptype, _) = dtype else { + vortex_bail!( + "only primitive dtypes are supported in SequenceArray currently, got {dtype}" + ); + }; + let metadata = SequenceMetadata::decode(metadata)?; + + let base_metadata = metadata + .base + .as_ref() + .ok_or_else(|| vortex_err!("base required"))?; + + let multiplier_metadata = metadata + .multiplier + .as_ref() + .ok_or_else(|| vortex_err!("multiplier required"))?; + + // We go via Scalar to validate that the value is valid for the ptype. + let base = Scalar::from_proto_value( + base_metadata, + &DType::Primitive(*output_ptype, NonNullable), + session, + )? + .as_primitive() + .pvalue() + .vortex_expect("sequence array base should be a non-nullable primitive"); + + // The serialized step preserves signedness independently of the output ptype. + let multiplier_ptype = SequenceData::multiplier_ptype_from_proto(multiplier_metadata)?; + let multiplier = Scalar::from_proto_value( + multiplier_metadata, + &DType::Primitive(multiplier_ptype, NonNullable), + session, + )? + .as_primitive() + .pvalue() + .vortex_expect("sequence array multiplier should be a non-nullable primitive"); + + let data = + SequenceData::try_new(base, multiplier, *output_ptype, dtype.nullability(), len)?; + Ok( + Array::try_from_parts(ArrayParts::new(self.clone(), dtype.clone(), len, data))? + .into_array(), + ) + } +} diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index 79edc647c37..88e36383d05 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -7,7 +7,6 @@ use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; -use prost::Message as _; use vortex_array::AnyCanonical; use vortex_array::Array; use vortex_array::ArrayEq; @@ -43,7 +42,6 @@ use vortex_array::require_opt_child; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_array::scalar_fn::fns::operators::Operator; -use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; @@ -52,7 +50,6 @@ use vortex_error::VortexExpect as _; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; -use vortex_error::vortex_ensure_eq; use vortex_error::vortex_panic; use vortex_mask::AllOr; use vortex_mask::Mask; @@ -68,6 +65,7 @@ mod canonical; mod compute; mod kernel; mod ops; +mod plugin; mod rules; mod slice; @@ -244,64 +242,6 @@ impl VTable for Sparse { vortex_array::vtable::unsupported_buffer_replacement(array, buffers) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - let patches = array.patches().to_metadata(array.len(), array.dtype())?; - let metadata = SparseMetadata { patches }; - - // Note that we DO NOT serialize the fill value since that is stored in the buffers. - Ok(Some(metadata.encode_to_vec())) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - let metadata = SparseMetadata::decode(metadata)?; - - // Once we have the patches metadata, we need to get the fill value from the buffers. - - if buffers.len() != 1 { - vortex_bail!("Expected 1 buffer, got {}", buffers.len()); - } - let scalar_bytes: &[u8] = &buffers[0].clone().try_to_host_sync()?; - - let scalar_value = ScalarValue::from_proto_bytes(scalar_bytes, dtype, session)?; - let fill_value = Scalar::try_new(dtype.clone(), scalar_value)?; - - vortex_ensure_eq!( - children.len(), - 2, - "SparseArray expects 2 children for sparse encoding, found {}", - children.len() - ); - - let patch_indices = children.get( - 0, - &metadata.patches.indices_dtype()?, - metadata.patches.len()?, - )?; - let patch_values = children.get(1, dtype, metadata.patches.len()?)?; - - let patches = Patches::new( - len, - metadata.patches.offset()?, - patch_indices, - patch_values, - None, - )?; - let slots = SparseData::make_slots(&patches); - let data = SparseData::from_patches(&patches, fill_value)?; - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { SparseSlots::NAMES[idx].to_string() } diff --git a/encodings/sparse/src/plugin.rs b/encodings/sparse/src/plugin.rs new file mode 100644 index 00000000000..d65c538247e --- /dev/null +++ b/encodings/sparse/src/plugin.rs @@ -0,0 +1,110 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use prost::Message as _; +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::patches::Patches; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_array::vtable::VTable; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use crate::Sparse; +use crate::SparseData; +use crate::SparseExt; +use crate::SparseMetadata; + +impl ArrayPlugin for Sparse { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array + .as_opt::() + .ok_or_else(|| vortex_err!("Sparse plugin cannot serialize {}", array.encoding_id()))?; + let patches = view.patches().to_metadata(view.len(), view.dtype())?; + let metadata = SparseMetadata { patches }; + + // Note that we DO NOT serialize the fill value since that is stored in the buffers. + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata.encode_to_vec(), + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Sparse plugin does not recognize serialized ID {serialized_id}" + ); + let metadata = SparseMetadata::decode(metadata)?; + + // Once we have the patches metadata, we need to get the fill value from the buffers. + + if buffers.len() != 1 { + vortex_bail!("Expected 1 buffer, got {}", buffers.len()); + } + let scalar_bytes: &[u8] = &buffers[0].clone().try_to_host_sync()?; + + let scalar_value = ScalarValue::from_proto_bytes(scalar_bytes, dtype, session)?; + let fill_value = Scalar::try_new(dtype.clone(), scalar_value)?; + + vortex_ensure_eq!( + children.len(), + 2, + "SparseArray expects 2 children for sparse encoding, found {}", + children.len() + ); + + let patch_indices = children.get( + 0, + &metadata.patches.indices_dtype()?, + metadata.patches.len()?, + )?; + let patch_values = children.get(1, dtype, metadata.patches.len()?)?; + + let patches = Patches::new( + len, + metadata.patches.offset()?, + patch_indices, + patch_values, + None, + )?; + let slots = SparseData::make_slots(&patches); + let data = SparseData::from_patches(&patches, fill_value)?; + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/encodings/zigzag/src/array.rs b/encodings/zigzag/src/array.rs index 09076fcbbdf..6f6ebe3235d 100644 --- a/encodings/zigzag/src/array.rs +++ b/encodings/zigzag/src/array.rs @@ -23,7 +23,6 @@ use vortex_array::dtype::DType; use vortex_array::dtype::PType; use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::scalar::Scalar; -use vortex_array::serde::ArrayChildren; use vortex_array::smallvec::smallvec; use vortex_array::vtable::OperationsVTable; use vortex_array::vtable::VTable; @@ -34,7 +33,6 @@ 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 zigzag::ZigZag as ExternalZigZag; @@ -42,6 +40,8 @@ use crate::compute::ZigZagEncoded; use crate::rules::RULES; use crate::zigzag_decode; +mod plugin; + /// A [`ZigZag`]-encoded Vortex array. pub type ZigZagArray = Array; @@ -97,41 +97,6 @@ impl VTable for ZigZag { vortex_array::vtable::with_empty_buffers(self, array, buffers) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - if !metadata.is_empty() { - vortex_bail!( - "ZigZagArray expects empty metadata, got {} bytes", - metadata.len() - ); - } - if children.len() != 1 { - vortex_bail!("Expected 1 child, got {}", children.len()); - } - - let ptype = PType::try_from(dtype)?; - let encoded_type = DType::Primitive(ptype.to_unsigned(), dtype.nullability()); - - let encoded = children.get(0, &encoded_type, len)?; - let slots = smallvec![Some(encoded.clone())]; - let data = ZigZagData::try_new(encoded.dtype())?; - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { ZigZagSlots::NAMES[idx].to_string() } diff --git a/encodings/zigzag/src/array/plugin.rs b/encodings/zigzag/src/array/plugin.rs new file mode 100644 index 00000000000..6bf651688eb --- /dev/null +++ b/encodings/zigzag/src/array/plugin.rs @@ -0,0 +1,84 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +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::smallvec::smallvec; +use vortex_array::vtable::VTable; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_session::VortexSession; + +use super::ZigZag; +use super::ZigZagData; + +impl ArrayPlugin for ZigZag { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + array.is::(), + "ZigZag plugin cannot serialize {}", + array.encoding_id() + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers: _, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "ZigZag plugin does not recognize serialized ID {serialized_id}" + ); + if !metadata.is_empty() { + vortex_bail!( + "ZigZagArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + if children.len() != 1 { + vortex_bail!("Expected 1 child, got {}", children.len()); + } + + let ptype = PType::try_from(dtype)?; + let encoded_type = DType::Primitive(ptype.to_unsigned(), dtype.nullability()); + + let encoded = children.get(0, &encoded_type, len)?; + let slots = smallvec![Some(encoded.clone())]; + let data = ZigZagData::try_new(encoded.dtype())?; + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index ce9d5d38b28..00f7a2fe634 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -12,7 +12,6 @@ use std::sync::Arc; use itertools::Itertools as _; use num_traits::AsPrimitive; -use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -39,7 +38,6 @@ use vortex_array::dtype::DType; use vortex_array::dtype::OffsetBuilderPType; use vortex_array::match_each_varbin_builder; use vortex_array::scalar::Scalar; -use vortex_array::serde::ArrayChildren; use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; use vortex_array::vtable::OperationsVTable; @@ -60,7 +58,6 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_mask::AllOr; use vortex_mask::Mask; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use zstd::zstd_safe::WriteBuf; @@ -68,6 +65,8 @@ use crate::ZstdFrameMetadata; use crate::ZstdMetadata; use crate::validate_frame_content_size; +mod plugin; + // Zstd doesn't support training dictionaries on very few samples. const MIN_SAMPLES_FOR_DICTIONARY: usize = 8; type ViewLen = u32; @@ -212,57 +211,6 @@ impl VTable for Zstd { ) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(array.metadata.clone().encode_to_vec())) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - let metadata = ZstdMetadata::decode(metadata)?; - let validity = if children.is_empty() { - Validity::from(dtype.nullability()) - } else if children.len() == 1 { - let validity = children.get(0, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("ZstdArray expected 0 or 1 child, got {}", children.len()); - }; - - let (dictionary_buffer, compressed_buffers) = if metadata.dictionary_size == 0 { - // no dictionary - ( - None, - buffers - .iter() - .map(|b| b.clone().try_to_host_sync()) - .collect::>>()?, - ) - } else { - // with dictionary - ( - Some(buffers[0].clone().try_to_host_sync()?), - buffers[1..] - .iter() - .map(|b| b.clone().try_to_host_sync()) - .collect::>>()?, - ) - }; - - let slots = smallvec![validity_to_child(&validity, len)]; - let data = ZstdData::new(dictionary_buffer, compressed_buffers, metadata, len); - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { ZstdSlots::NAMES[idx].to_string() } diff --git a/encodings/zstd/src/array/plugin.rs b/encodings/zstd/src/array/plugin.rs new file mode 100644 index 00000000000..8d366968190 --- /dev/null +++ b/encodings/zstd/src/array/plugin.rs @@ -0,0 +1,101 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use prost::Message as _; +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::smallvec::smallvec; +use vortex_array::validity::Validity; +use vortex_array::vtable::VTable; +use vortex_array::vtable::validity_to_child; +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::Zstd; +use super::ZstdData; +use crate::ZstdMetadata; + +impl ArrayPlugin for Zstd { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array + .as_opt::() + .ok_or_else(|| vortex_err!("Zstd plugin cannot serialize {}", array.encoding_id()))?; + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + view.metadata.clone().encode_to_vec(), + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Zstd plugin does not recognize serialized ID {serialized_id}" + ); + let metadata = ZstdMetadata::decode(metadata)?; + let validity = if children.is_empty() { + Validity::from(dtype.nullability()) + } else if children.len() == 1 { + let validity = children.get(0, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + vortex_bail!("ZstdArray expected 0 or 1 child, got {}", children.len()); + }; + + let (dictionary_buffer, compressed_buffers) = if metadata.dictionary_size == 0 { + // no dictionary + ( + None, + buffers + .iter() + .map(|b| b.clone().try_to_host_sync()) + .collect::>>()?, + ) + } else { + // with dictionary + ( + Some(buffers[0].clone().try_to_host_sync()?), + buffers[1..] + .iter() + .map(|b| b.clone().try_to_host_sync()) + .collect::>>()?, + ) + }; + + let slots = smallvec![validity_to_child(&validity, len)]; + let data = ZstdData::new(dictionary_buffer, compressed_buffers, metadata, len); + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/encodings/zstd/src/zstd_buffers.rs b/encodings/zstd/src/zstd_buffers.rs index 51d338777ca..579126b60e6 100644 --- a/encodings/zstd/src/zstd_buffers.rs +++ b/encodings/zstd/src/zstd_buffers.rs @@ -8,7 +8,6 @@ use std::hash::Hash; use std::hash::Hasher; use std::sync::Arc; -use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayDeserialization; use vortex_array::ArrayEq; @@ -24,7 +23,6 @@ use vortex_array::ExecutionResult; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::scalar::Scalar; -use vortex_array::serde::ArrayChildren; use vortex_array::session::ArraySessionExt; use vortex_array::validity::Validity; use vortex_array::vtable::OperationsVTable; @@ -39,9 +37,10 @@ use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::ZstdBuffersMetadata; use crate::validate_frame_content_size; +mod plugin; + /// A [`ZstdBuffers`]-encoded Vortex array. pub type ZstdBuffersArray = Array; @@ -445,69 +444,6 @@ impl VTable for ZstdBuffers { format!("child_{idx}") } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - let children: Vec<&ArrayRef> = array.slots().iter().flatten().collect(); - let child_dtypes = children - .iter() - .map(|child| child.dtype().try_into()) - .collect::>>()?; - let child_lens = children.iter().map(|child| child.len() as u64).collect(); - - Ok(Some( - ZstdBuffersMetadata { - inner_encoding_id: array.inner_encoding_id.to_string(), - inner_metadata: array.inner_metadata.clone(), - uncompressed_sizes: array.uncompressed_sizes.clone(), - buffer_alignments: array.buffer_alignments.clone(), - child_dtypes, - child_lens, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - let metadata = ZstdBuffersMetadata::decode(metadata)?; - let compressed_buffers: Vec = buffers.to_vec(); - - // Children belong to inner encodings, and serialization doesn't - // preserve their dtypes and values. Check dtypes are recovered from - // metadata. - vortex_ensure_eq!(metadata.child_dtypes.len(), children.len()); - vortex_ensure_eq!(metadata.child_lens.len(), children.len()); - - let slots: ArraySlots = (0..children.len()) - .map(|i| { - let child_dtype = DType::from_proto(&metadata.child_dtypes[i], session)?; - let child_len = usize::try_from(metadata.child_lens[i])?; - children.get(i, &child_dtype, child_len).map(Some) - }) - .collect::>>()? - .into(); - - let data = ZstdBuffersData { - inner_encoding_id: array_id_from_string(&metadata.inner_encoding_id), - inner_metadata: metadata.inner_metadata.clone(), - compressed_buffers, - uncompressed_sizes: metadata.uncompressed_sizes.clone(), - buffer_alignments: metadata.buffer_alignments.clone(), - }; - - data.validate()?; - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - // with_slots handles child replacement via the slots mechanism fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/encodings/zstd/src/zstd_buffers/plugin.rs b/encodings/zstd/src/zstd_buffers/plugin.rs new file mode 100644 index 00000000000..3f142bb54c6 --- /dev/null +++ b/encodings/zstd/src/zstd_buffers/plugin.rs @@ -0,0 +1,116 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use prost::Message as _; +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::ArraySlots; +use vortex_array::IntoArray; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::vtable::VTable; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use super::ZstdBuffers; +use super::ZstdBuffersData; +use super::array_id_from_string; +use crate::ZstdBuffersMetadata; + +impl ArrayPlugin for ZstdBuffers { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_opt::().ok_or_else(|| { + vortex_err!( + "ZstdBuffers plugin cannot serialize {}", + array.encoding_id() + ) + })?; + let children: Vec<&ArrayRef> = view.slots().iter().flatten().collect(); + let child_dtypes = children + .iter() + .map(|child| child.dtype().try_into()) + .collect::>>()?; + let child_lens = children.iter().map(|child| child.len() as u64).collect(); + + let metadata = ZstdBuffersMetadata { + inner_encoding_id: view.inner_encoding_id.to_string(), + inner_metadata: view.inner_metadata.clone(), + uncompressed_sizes: view.uncompressed_sizes.clone(), + buffer_alignments: view.buffer_alignments.clone(), + child_dtypes, + child_lens, + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "ZstdBuffers plugin does not recognize serialized ID {serialized_id}" + ); + let metadata = ZstdBuffersMetadata::decode(metadata)?; + let compressed_buffers: Vec = buffers.to_vec(); + + // Children belong to inner encodings, and serialization doesn't + // preserve their dtypes and values. Check dtypes are recovered from + // metadata. + vortex_ensure_eq!(metadata.child_dtypes.len(), children.len()); + vortex_ensure_eq!(metadata.child_lens.len(), children.len()); + + let slots: ArraySlots = (0..children.len()) + .map(|i| { + let child_dtype = DType::from_proto(&metadata.child_dtypes[i], session)?; + let child_len = usize::try_from(metadata.child_lens[i])?; + children.get(i, &child_dtype, child_len).map(Some) + }) + .collect::>>()? + .into(); + + let data = ZstdBuffersData { + inner_encoding_id: array_id_from_string(&metadata.inner_encoding_id), + inner_metadata: metadata.inner_metadata.clone(), + compressed_buffers, + uncompressed_sizes: metadata.uncompressed_sizes.clone(), + buffer_alignments: metadata.buffer_alignments.clone(), + }; + + data.validate()?; + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/vortex-array/src/array/foreign.rs b/vortex-array/src/array/foreign.rs index 4a665044f3f..abfec305497 100644 --- a/vortex-array/src/array/foreign.rs +++ b/vortex-array/src/array/foreign.rs @@ -9,7 +9,6 @@ use std::hash::Hasher; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_session::VortexSession; use crate::Array; use crate::ArrayRef; @@ -27,7 +26,6 @@ use crate::dtype::DType; use crate::executor::ExecutionCtx; use crate::hash::ArrayEq; use crate::hash::ArrayHash; -use crate::serde::ArrayChildren; use crate::validity::Validity; #[derive(Clone, Debug)] @@ -134,35 +132,6 @@ impl VTable for ForeignArray { .with_slots(array.slots().iter().cloned().collect())) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(array.metadata.clone())) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - let child_arrays = (0..children.len()) - .map(|idx| children.get(idx, dtype, len).map(Some)) - .collect::>()?; - - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - ForeignArrayData::new(metadata.to_vec(), buffers.to_vec()), - ) - .with_slots(child_arrays)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { format!("child[{idx}]") } diff --git a/vortex-array/src/array/plugin.rs b/vortex-array/src/array/plugin.rs index 5b4fca39854..8e90916684d 100644 --- a/vortex-array/src/array/plugin.rs +++ b/vortex-array/src/array/plugin.rs @@ -8,14 +8,10 @@ use std::sync::Arc; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_session::VortexSession; use crate::ArrayRef; -use crate::IntoArray; -use crate::array::Array; use crate::array::ArrayId; -use crate::array::VTable; use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::serde::ArrayChildren; @@ -161,47 +157,3 @@ impl Debug for dyn ArrayPlugin { f.debug_tuple("ArrayPlugin").field(&self.id()).finish() } } - -impl ArrayPlugin for V { - fn id(&self) -> ArrayId { - VTable::id(self) - } - - fn serialize( - &self, - array: &ArrayRef, - session: &VortexSession, - ) -> VortexResult> { - vortex_ensure!( - self.id() == array.encoding_id(), - "array plugin {} cannot serialize in-memory array {}", - self.id(), - array.encoding_id(), - ); - Ok(V::serialize(array.as_::(), session)? - .map(|metadata| ArraySerialization::from_array(self.id(), array, metadata))) - } - - fn deserialize( - &self, - parts: ArrayDeserialization<'_>, - session: &VortexSession, - ) -> VortexResult { - vortex_ensure!( - self.id() == parts.serialized_id, - "array plugin {} does not recognize serialized ID {}", - self.id(), - parts.serialized_id, - ); - Ok(Array::::try_from_parts(V::deserialize( - self, - parts.dtype, - parts.len, - parts.metadata, - parts.buffers, - parts.children, - session, - )?)? - .into_array()) - } -} diff --git a/vortex-array/src/array/vtable/mod.rs b/vortex-array/src/array/vtable/mod.rs index 3078a6ae9f5..75db622ae50 100644 --- a/vortex-array/src/array/vtable/mod.rs +++ b/vortex-array/src/array/vtable/mod.rs @@ -5,7 +5,7 @@ //! //! A Vortex array encoding is implemented by a small static vtable type plus an associated //! `TypedArrayData` value stored in each array instance. The vtable owns behavior such as -//! validation, serialization, execution, child traversal, scalar access, and validity access. +//! validation, execution, child traversal, scalar access, and validity access. //! //! The public [`ArrayRef`] API performs common precondition checks before calling //! into these traits. Implementations should focus on encoding-specific work and uphold the @@ -26,7 +26,6 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use crate::Array; use crate::ArrayRef; @@ -47,7 +46,6 @@ use crate::hash::ArrayEq; use crate::hash::ArrayHash; use crate::patches::Patches; use crate::scalar::ScalarValue; -use crate::serde::ArrayChildren; use crate::validity::Validity; /// The array [`VTable`] encapsulates logic for an Array type within Vortex. @@ -55,8 +53,9 @@ use crate::validity::Validity; /// The logic is split across several "VTable" traits to enable easier code organization than /// simply lumping everything into a single trait. /// -/// From this [`VTable`] trait, we derive implementations for the sealed `DynArrayData` trait and the -/// public [`ArrayPlugin`] registry trait. +/// From this [`VTable`] trait, we derive implementations for the sealed `DynArrayData` trait. +/// Serialization is not part of the vtable: encodings that support serde implement +/// [`ArrayPlugin`] separately and register it with the session. /// /// The functions defined in these vtable traits will typically document their pre- and /// post-conditions. The pre-conditions are validated inside the `DynArrayData` and [`ArrayRef`] @@ -115,31 +114,6 @@ pub trait VTable: 'static + Clone + Sized + Send + Sync + Debug { buffers: &[BufferHandle], ) -> VortexResult>; - /// Serialize encoding metadata into a byte buffer for IPC or file storage. - /// - /// Return `None` if the array cannot be serialized by this encoding. Buffers and children are - /// serialized separately through [`buffer`](Self::buffer), [`nbuffers`](Self::nbuffers), and - /// child traversal. - fn serialize( - array: ArrayView<'_, Self>, - session: &VortexSession, - ) -> VortexResult>>; - - /// Deserialize an array from serialized metadata, buffers, and children. - /// - /// The returned [`ArrayParts`] are still validated by the generic adapter. - /// Deserializers should use the provided `session` to resolve plugin-owned metadata instead of - /// relying on global state. - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult>; - /// Writes the array's logical values into a canonical builder. /// /// The default implementation executes the full array to [`Canonical`] and appends that result. diff --git a/vortex-array/src/arrays/bool/vtable/mod.rs b/vortex-array/src/arrays/bool/vtable/mod.rs index 393414924b3..b36fdfd2e83 100644 --- a/vortex-array/src/arrays/bool/vtable/mod.rs +++ b/vortex-array/src/arrays/bool/vtable/mod.rs @@ -4,8 +4,6 @@ use std::hash::Hash; use std::hash::Hasher; -use prost::Message; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -26,10 +24,9 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::BoolBuilder; use crate::dtype::DType; -use crate::serde::ArrayChildren; -use crate::validity::Validity; mod kernel; mod operations; +mod plugin; mod validity; use vortex_session::registry::CachedId; @@ -114,20 +111,6 @@ impl VTable for Bool { ) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - let offset = array.meta.offset(); - assert!(offset < 8, "Offset must be <8, got {offset}"); - Ok(Some( - BoolMetadata { - offset: u32::try_from(offset).vortex_expect("checked"), - } - .encode_to_vec(), - )) - } - fn validate( &self, data: &BoolData, @@ -159,35 +142,6 @@ impl VTable for Bool { Ok(()) } - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - let metadata = BoolMetadata::decode(metadata)?; - if buffers.len() != 1 { - vortex_bail!("Expected 1 buffer, got {}", buffers.len()); - } - - let validity = if children.is_empty() { - Validity::from(dtype.nullability()) - } else if children.len() == 1 { - let validity = children.get(0, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("Expected 0 or 1 child, got {}", children.len()); - }; - - let buffer = buffers[0].clone(); - let slots = BoolData::make_slots(&validity, len); - let data = BoolData::try_new_from_handle(buffer, metadata.offset as usize, len, validity)?; - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { BoolSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/bool/vtable/plugin.rs b/vortex-array/src/arrays/bool/vtable/plugin.rs new file mode 100644 index 00000000000..fd174e9467d --- /dev/null +++ b/vortex-array/src/arrays/bool/vtable/plugin.rs @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use prost::Message; +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_session::VortexSession; + +use super::BoolMetadata; +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::VTable; +use crate::arrays::Bool; +use crate::arrays::bool::BoolData; +use crate::validity::Validity; + +impl ArrayPlugin for Bool { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array + .as_opt::() + .ok_or_else(|| vortex_err!("Bool plugin cannot serialize {}", array.encoding_id()))?; + let offset = view.meta.offset(); + assert!(offset < 8, "Offset must be <8, got {offset}"); + let metadata = BoolMetadata { + offset: u32::try_from(offset).vortex_expect("checked"), + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Bool plugin does not recognize serialized ID {serialized_id}" + ); + let metadata = BoolMetadata::decode(metadata)?; + if buffers.len() != 1 { + vortex_bail!("Expected 1 buffer, got {}", buffers.len()); + } + + let validity = if children.is_empty() { + Validity::from(dtype.nullability()) + } else if children.len() == 1 { + let validity = children.get(0, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + vortex_bail!("Expected 0 or 1 child, got {}", children.len()); + }; + + let buffer = buffers[0].clone(); + let slots = BoolData::make_slots(&validity, len); + let data = BoolData::try_new_from_handle(buffer, metadata.offset as usize, len, validity)?; + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index 59e89f031e1..a8b510c39e1 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -4,14 +4,10 @@ use std::hash::Hasher; use itertools::Itertools; -use smallvec::SmallVec; 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::ArrayEq; @@ -22,14 +18,12 @@ use crate::EqMode; use crate::ExecutionCtx; use crate::ExecutionResult; use crate::IntoArray; -use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayParts; use crate::array::ArrayView; use crate::array::VTable; use crate::array::with_empty_buffers; -use crate::arrays::PrimitiveArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::chunked::ChunkedData; use crate::arrays::chunked::array::ChunkedSlots; @@ -40,9 +34,9 @@ use crate::builders::ArrayBuilder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; -use crate::serde::ArrayChildren; mod canonical; mod operations; +mod plugin; mod validity; /// A [`Chunked`]-encoded Vortex array. @@ -164,76 +158,6 @@ impl VTable for Chunked { with_empty_buffers(self, array, buffers) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - if !metadata.is_empty() { - vortex_bail!( - "ChunkedArray expects empty metadata, got {} bytes", - metadata.len() - ); - } - if children.is_empty() { - vortex_bail!("Chunked array needs at least one child"); - } - - let nchunks = children.len() - 1; - let chunk_offsets = children.get( - ChunkedSlots::CHUNK_OFFSETS, - &DType::Primitive(PType::U64, Nullability::NonNullable), - nchunks + 1, - )?; - let mut ctx = session.create_execution_ctx(); - let chunk_offsets_buf = chunk_offsets - .clone() - .execute::(&mut ctx)? - .to_buffer::(); - let chunk_offsets_usize = chunk_offsets_buf - .iter() - .copied() - .map(|offset| { - usize::try_from(offset) - .map_err(|_| vortex_err!("chunk offset {offset} exceeds usize range")) - }) - .collect::>>()?; - let mut slots = SmallVec::with_capacity(children.len()); - slots.push(Some(chunk_offsets)); - for (idx, (start, end)) in chunk_offsets_usize - .iter() - .copied() - .tuple_windows() - .enumerate() - { - let chunk_len = end - start; - slots.push(Some(children.get( - idx + ChunkedSlots::CHUNKS_OFFSET, - dtype, - chunk_len, - )?)); - } - - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - ChunkedData::new(chunk_offsets_usize), - ) - .with_slots(slots)) - } - fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, diff --git a/vortex-array/src/arrays/chunked/vtable/plugin.rs b/vortex-array/src/arrays/chunked/vtable/plugin.rs new file mode 100644 index 00000000000..1138defc885 --- /dev/null +++ b/vortex-array/src/arrays/chunked/vtable/plugin.rs @@ -0,0 +1,125 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use itertools::Itertools; +use smallvec::SmallVec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::VTable; +use crate::arrays::Chunked; +use crate::arrays::PrimitiveArray; +use crate::arrays::chunked::ChunkedData; +use crate::arrays::chunked::array::ChunkedSlots; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; + +impl ArrayPlugin for Chunked { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + array.is::(), + "Chunked plugin cannot serialize {}", + array.encoding_id() + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers: _, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Chunked plugin does not recognize serialized ID {serialized_id}" + ); + if !metadata.is_empty() { + vortex_bail!( + "ChunkedArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + if children.is_empty() { + vortex_bail!("Chunked array needs at least one child"); + } + + let nchunks = children.len() - 1; + let chunk_offsets = children.get( + ChunkedSlots::CHUNK_OFFSETS, + &DType::Primitive(PType::U64, Nullability::NonNullable), + nchunks + 1, + )?; + let mut ctx = session.create_execution_ctx(); + let chunk_offsets_buf = chunk_offsets + .clone() + .execute::(&mut ctx)? + .to_buffer::(); + let chunk_offsets_usize = chunk_offsets_buf + .iter() + .copied() + .map(|offset| { + usize::try_from(offset) + .map_err(|_| vortex_err!("chunk offset {offset} exceeds usize range")) + }) + .collect::>>()?; + let mut slots = SmallVec::with_capacity(children.len()); + slots.push(Some(chunk_offsets)); + for (idx, (start, end)) in chunk_offsets_usize + .iter() + .copied() + .tuple_windows() + .enumerate() + { + let chunk_len = end - start; + slots.push(Some(children.get( + idx + ChunkedSlots::CHUNKS_OFFSET, + dtype, + chunk_len, + )?)); + } + + Ok(Array::try_from_parts( + ArrayParts::new( + self.clone(), + dtype.clone(), + len, + ChunkedData::new(chunk_offsets_usize), + ) + .with_slots(slots), + )? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index c0796a0e7d0..af98ad33a3c 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -10,7 +10,6 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::ArrayEq; @@ -49,11 +48,10 @@ use crate::match_each_native_ptype; use crate::match_each_varbin_builder; use crate::scalar::DecimalValue; use crate::scalar::ListScalar; -use crate::scalar::Scalar; use crate::scalar::ScalarValue; -use crate::serde::ArrayChildren; pub(crate) mod canonical; mod operations; +mod plugin; mod validity; /// A [`Constant`]-encoded Vortex array. @@ -131,45 +129,6 @@ impl VTable for Constant { vortex_panic!("ConstantArray slot_name index {idx} out of bounds") } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - // HACK: Because the scalar is stored in the buffers, we do not need to serialize the - // metadata at all. - Ok(Some(vec![])) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - _metadata: &[u8], - - buffers: &[BufferHandle], - _children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - vortex_ensure!( - buffers.len() == 1, - "Expected 1 buffer, got {}", - buffers.len() - ); - - let buffer = buffers[0].clone().try_to_host_sync()?; - let bytes: &[u8] = buffer.as_ref(); - - let scalar_value = ScalarValue::from_proto_bytes(bytes, dtype, session)?; - let scalar = Scalar::try_new(dtype.clone(), scalar_value)?; - - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - ConstantData::new(scalar), - )) - } - fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, diff --git a/vortex-array/src/arrays/constant/vtable/plugin.rs b/vortex-array/src/arrays/constant/vtable/plugin.rs new file mode 100644 index 00000000000..2e1757cae65 --- /dev/null +++ b/vortex-array/src/arrays/constant/vtable/plugin.rs @@ -0,0 +1,83 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_session::VortexSession; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::VTable; +use crate::arrays::Constant; +use crate::arrays::constant::ConstantData; +use crate::scalar::Scalar; +use crate::scalar::ScalarValue; + +impl ArrayPlugin for Constant { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + array.is::(), + "Constant plugin cannot serialize {}", + array.encoding_id() + ); + // HACK: Because the scalar is stored in the buffers, we do not need to serialize the + // metadata at all. + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata: _, + buffers, + children: _, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Constant plugin does not recognize serialized ID {serialized_id}" + ); + vortex_ensure!( + buffers.len() == 1, + "Expected 1 buffer, got {}", + buffers.len() + ); + + let buffer = buffers[0].clone().try_to_host_sync()?; + let bytes: &[u8] = buffer.as_ref(); + + let scalar_value = ScalarValue::from_proto_bytes(bytes, dtype, session)?; + let scalar = Scalar::try_new(dtype.clone(), scalar_value)?; + + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + ConstantData::new(scalar), + ))? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index 5f2e087e983..662011c7c55 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -3,8 +3,6 @@ use std::hash::Hasher; -use prost::Message; -use vortex_buffer::Alignment; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -24,11 +22,9 @@ use crate::builders::ArrayBuilder; use crate::builders::DecimalBuilder; use crate::dtype::DType; use crate::dtype::DecimalType; -use crate::dtype::NativeDecimalType; -use crate::match_each_decimal_value_type; -use crate::serde::ArrayChildren; mod kernel; mod operations; +mod plugin; mod validity; use std::hash::Hash; @@ -104,18 +100,6 @@ impl VTable for Decimal { ) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - DecimalMetadata { - values_type: array.values_type() as i32, - } - .encode_to_vec(), - )) - } - fn validate( &self, data: &DecimalData, @@ -148,37 +132,6 @@ impl VTable for Decimal { Ok(()) } - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - let metadata = DecimalMetadata::decode(metadata)?; - let values = fixed_width::single_buffer(buffers)?; - - let validity = fixed_width::deserialize_validity(dtype.nullability(), len, children)?; - - let Some(decimal_dtype) = dtype.as_decimal_opt() else { - vortex_bail!("Expected Decimal dtype, got {:?}", dtype) - }; - - let slots = DecimalData::make_slots(&validity, len); - let data = match_each_decimal_value_type!(metadata.values_type(), |D| { - // Check and reinterpret-cast the buffer - vortex_ensure!( - values.is_aligned_to(Alignment::of::()), - "DecimalArray buffer not aligned for values type {:?}", - D::DECIMAL_TYPE - ); - DecimalData::try_new_handle(values, metadata.values_type(), *decimal_dtype) - })?; - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { DecimalSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/decimal/vtable/plugin.rs b/vortex-array/src/arrays/decimal/vtable/plugin.rs new file mode 100644 index 00000000000..65798864288 --- /dev/null +++ b/vortex-array/src/arrays/decimal/vtable/plugin.rs @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use prost::Message; +use vortex_buffer::Alignment; +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::DecimalMetadata; +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::VTable; +use crate::arrays::Decimal; +use crate::arrays::decimal::DecimalData; +use crate::arrays::fixed_width::vtable as fixed_width; +use crate::dtype::NativeDecimalType; +use crate::match_each_decimal_value_type; + +impl ArrayPlugin for Decimal { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_opt::().ok_or_else(|| { + vortex_err!("Decimal plugin cannot serialize {}", array.encoding_id()) + })?; + let metadata = DecimalMetadata { + values_type: view.values_type() as i32, + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Decimal plugin does not recognize serialized ID {serialized_id}" + ); + let metadata = DecimalMetadata::decode(metadata)?; + let values = fixed_width::single_buffer(buffers)?; + + let validity = fixed_width::deserialize_validity(dtype.nullability(), len, children)?; + + let Some(decimal_dtype) = dtype.as_decimal_opt() else { + vortex_bail!("Expected Decimal dtype, got {:?}", dtype) + }; + + let slots = DecimalData::make_slots(&validity, len); + let data = match_each_decimal_value_type!(metadata.values_type(), |D| { + // Check and reinterpret-cast the buffer + vortex_ensure!( + values.is_aligned_to(Alignment::of::()), + "DecimalArray buffer not aligned for values type {:?}", + D::DECIMAL_TYPE + ); + DecimalData::try_new_handle(values, metadata.values_type(), *decimal_dtype) + })?; + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index 2eca82b5368..a4db25a69e2 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/vortex-array/src/arrays/dict/vtable/mod.rs @@ -4,12 +4,8 @@ use std::hash::Hasher; use num_traits::AsPrimitive; -use prost::Message; -use smallvec::smallvec; 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_mask::AllOr; use vortex_mask::Mask; @@ -17,7 +13,6 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use super::DictData; -use super::DictMetadata; use super::DictOwnedExt; use super::DictParts; use super::array::DictSlots; @@ -39,7 +34,6 @@ use crate::array::with_empty_buffers; use crate::arrays::ConstantArray; use crate::arrays::Primitive; use crate::arrays::VarBinView; -use crate::arrays::dict::DictArrayExt; use crate::arrays::dict::DictArraySlotsExt; use crate::arrays::dict::compute::rules::PARENT_RULES; use crate::arrays::dict::execute::take_canonical; @@ -48,19 +42,17 @@ use crate::builders::ArrayBuilder; use crate::builders::VarBinBuilder; use crate::builders::VarBinViewBuilder; use crate::dtype::DType; -use crate::dtype::Nullability; use crate::dtype::OffsetBuilderPType; -use crate::dtype::PType; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; use crate::match_each_integer_ptype; use crate::match_each_varbin_builder; use crate::require_child; use crate::scalar::Scalar; -use crate::serde::ArrayChildren; mod kernel; mod operations; +mod plugin; mod validity; /// A [`Dict`]-encoded Vortex array. @@ -135,59 +127,6 @@ impl VTable for Dict { with_empty_buffers(self, array, buffers) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - DictMetadata { - codes_ptype: PType::try_from(array.codes().dtype())? as i32, - values_len: u32::try_from(array.values().len()).map_err(|_| { - vortex_err!( - "Dictionary values size {} overflowed u32", - array.values().len() - ) - })?, - is_nullable_codes: Some(array.codes().dtype().is_nullable()), - all_values_referenced: Some(array.has_all_values_referenced()), - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - let metadata = DictMetadata::decode(metadata)?; - if children.len() != 2 { - vortex_bail!( - "Expected 2 children for dict encoding, found {}", - children.len() - ) - } - let codes_nullable = metadata - .is_nullable_codes - .map(Nullability::from) - // If no `is_nullable_codes` metadata use the nullability of the values - // (and whole array) as before. - .unwrap_or_else(|| dtype.nullability()); - let codes_dtype = DType::Primitive(metadata.codes_ptype(), codes_nullable); - let codes = children.get(0, &codes_dtype, len)?; - let values = children.get(1, dtype, metadata.values_len as usize)?; - let all_values_referenced = metadata.all_values_referenced.unwrap_or(false); - - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, unsafe { - DictData::new_unchecked().set_all_values_referenced(all_values_referenced) - }) - .with_slots(smallvec![Some(codes), Some(values)])) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { DictSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/dict/vtable/plugin.rs b/vortex-array/src/arrays/dict/vtable/plugin.rs new file mode 100644 index 00000000000..1c463eedde5 --- /dev/null +++ b/vortex-array/src/arrays/dict/vtable/plugin.rs @@ -0,0 +1,105 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use prost::Message; +use smallvec::smallvec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::VTable; +use crate::arrays::Dict; +use crate::arrays::dict::DictArrayExt; +use crate::arrays::dict::DictArraySlotsExt; +use crate::arrays::dict::DictData; +use crate::arrays::dict::DictMetadata; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; + +impl ArrayPlugin for Dict { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array + .as_opt::() + .ok_or_else(|| vortex_err!("Dict plugin cannot serialize {}", array.encoding_id()))?; + let metadata = DictMetadata { + codes_ptype: PType::try_from(view.codes().dtype())? as i32, + values_len: u32::try_from(view.values().len()).map_err(|_| { + vortex_err!( + "Dictionary values size {} overflowed u32", + view.values().len() + ) + })?, + is_nullable_codes: Some(view.codes().dtype().is_nullable()), + all_values_referenced: Some(view.has_all_values_referenced()), + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers: _, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Dict plugin does not recognize serialized ID {serialized_id}" + ); + let metadata = DictMetadata::decode(metadata)?; + if children.len() != 2 { + vortex_bail!( + "Expected 2 children for dict encoding, found {}", + children.len() + ) + } + let codes_nullable = metadata + .is_nullable_codes + .map(Nullability::from) + // If no `is_nullable_codes` metadata use the nullability of the values + // (and whole array) as before. + .unwrap_or_else(|| dtype.nullability()); + let codes_dtype = DType::Primitive(metadata.codes_ptype(), codes_nullable); + let codes = children.get(0, &codes_dtype, len)?; + let values = children.get(1, dtype, metadata.values_len as usize)?; + let all_values_referenced = metadata.all_values_referenced.unwrap_or(false); + + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, unsafe { + DictData::new_unchecked().set_all_values_referenced(all_values_referenced) + }) + .with_slots(smallvec![Some(codes), Some(values)]), + )? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/extension/vtable/mod.rs b/vortex-array/src/arrays/extension/vtable/mod.rs index 2937a2b98b9..bce38c6df8f 100644 --- a/vortex-array/src/arrays/extension/vtable/mod.rs +++ b/vortex-array/src/arrays/extension/vtable/mod.rs @@ -28,10 +28,10 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::ExtensionBuilder; use crate::dtype::DType; -use crate::serde::ArrayChildren; mod kernel; mod operations; +mod plugin; mod validity; /// An extension array that wraps another array with additional type information. @@ -144,42 +144,6 @@ impl VTable for Extension { with_empty_buffers(self, array, buffers) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - if !metadata.is_empty() { - vortex_bail!( - "ExtensionArray expects empty metadata, got {} bytes", - metadata.len() - ); - } - 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( - ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData) - .with_slots(ExtensionSlots { storage }.into_slots()), - ) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { ExtensionSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/extension/vtable/plugin.rs b/vortex-array/src/arrays/extension/vtable/plugin.rs new file mode 100644 index 00000000000..3d0c45bad0a --- /dev/null +++ b/vortex-array/src/arrays/extension/vtable/plugin.rs @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_session::VortexSession; + +use crate::ArrayRef; +use crate::EmptyArrayData; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::VTable; +use crate::arrays::Extension; +use crate::arrays::extension::array::ExtensionSlots; +use crate::dtype::DType; + +impl ArrayPlugin for Extension { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + array.is::(), + "Extension plugin cannot serialize {}", + array.encoding_id() + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers: _, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Extension plugin does not recognize serialized ID {serialized_id}" + ); + if !metadata.is_empty() { + vortex_bail!( + "ExtensionArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + 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(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData) + .with_slots(ExtensionSlots { storage }.into_slots()), + )? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/filter/vtable.rs b/vortex-array/src/arrays/filter/vtable.rs index f5a61ba79d7..a133efa7e27 100644 --- a/vortex-array/src/arrays/filter/vtable.rs +++ b/vortex-array/src/arrays/filter/vtable.rs @@ -5,12 +5,10 @@ use std::hash::Hasher; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use vortex_mask::Mask; use vortex_mask::MaskValuesRef; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::AnyCanonical; @@ -42,7 +40,6 @@ use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; use crate::require_child; use crate::scalar::Scalar; -use crate::serde::ArrayChildren; use crate::validity::Validity; /// A [`Filter`]-encoded Vortex array. @@ -131,27 +128,6 @@ impl VTable for Filter { FilterSlots::NAMES[idx].to_string() } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - // TODO(joe): make this configurable - vortex_bail!("Filter array is not serializable") - } - - fn deserialize( - &self, - _dtype: &DType, - _len: usize, - _metadata: &[u8], - - _buffers: &[BufferHandle], - _children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_bail!("Filter array is not serializable") - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { // Match the mask once. A zero-length mask is both all true and all false, so check the // empty output before the unfiltered one. diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs index 656a08b0f66..2814b7a4a0c 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs @@ -32,10 +32,9 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::FixedSizeListBuilder; use crate::dtype::DType; -use crate::serde::ArrayChildren; -use crate::validity::Validity; mod kernel; mod operations; +mod plugin; mod validity; /// A [`FixedSizeList`]-encoded Vortex array. @@ -98,13 +97,6 @@ impl VTable for FixedSizeList { PARENT_RULES.evaluate(array, parent, child_idx) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - fn validate( &self, data: &FixedSizeListData, @@ -147,54 +139,6 @@ impl VTable for FixedSizeList { Ok(()) } - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - if !metadata.is_empty() { - vortex_bail!( - "FixedSizeListArray expects empty metadata, got {} bytes", - metadata.len() - ); - } - vortex_ensure!( - buffers.is_empty(), - "`FixedSizeList::build` expects no buffers" - ); - - let DType::FixedSizeList(element_dtype, list_size, _) = &dtype else { - vortex_bail!("Expected `DType::FixedSizeList`, got {:?}", dtype); - }; - - let validity = { - if children.len() > 2 { - vortex_bail!("`FixedSizeList::build` method expected 1 or 2 children") - } - - if children.len() == 2 { - let validity = children.get(1, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - debug_assert_eq!(children.len(), 1); - Validity::from(dtype.nullability()) - } - }; - - let num_elements = len * (*list_size as usize); - let elements = children.get(0, element_dtype.as_ref(), num_elements)?; - - let data = - FixedSizeListData::try_build(elements.clone(), *list_size, validity.clone(), len)?; - let slots = FixedSizeListData::make_slots(&elements, &validity, len); - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { FixedSizeListSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/plugin.rs b/vortex-array/src/arrays/fixed_size_list/vtable/plugin.rs new file mode 100644 index 00000000000..b31e3f7f972 --- /dev/null +++ b/vortex-array/src/arrays/fixed_size_list/vtable/plugin.rs @@ -0,0 +1,102 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_session::VortexSession; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::VTable; +use crate::arrays::FixedSizeList; +use crate::arrays::fixed_size_list::FixedSizeListData; +use crate::dtype::DType; +use crate::validity::Validity; + +impl ArrayPlugin for FixedSizeList { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + array.is::(), + "FixedSizeList plugin cannot serialize {}", + array.encoding_id() + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "FixedSizeList plugin does not recognize serialized ID {serialized_id}" + ); + if !metadata.is_empty() { + vortex_bail!( + "FixedSizeListArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + vortex_ensure!( + buffers.is_empty(), + "`FixedSizeList::build` expects no buffers" + ); + + let DType::FixedSizeList(element_dtype, list_size, _) = &dtype else { + vortex_bail!("Expected `DType::FixedSizeList`, got {:?}", dtype); + }; + + let validity = { + if children.len() > 2 { + vortex_bail!("`FixedSizeList::build` method expected 1 or 2 children") + } + + if children.len() == 2 { + let validity = children.get(1, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + debug_assert_eq!(children.len(), 1); + Validity::from(dtype.nullability()) + } + }; + + let num_elements = len * (*list_size as usize); + let elements = children.get(0, element_dtype.as_ref(), num_elements)?; + + let data = + FixedSizeListData::try_build(elements.clone(), *list_size, validity.clone(), len)?; + let slots = FixedSizeListData::make_slots(&elements, &validity, len); + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/interleave/mod.rs b/vortex-array/src/arrays/interleave/mod.rs index fb649cf6e19..f79d0f0f78f 100644 --- a/vortex-array/src/arrays/interleave/mod.rs +++ b/vortex-array/src/arrays/interleave/mod.rs @@ -50,7 +50,6 @@ 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::ArrayEq; @@ -75,7 +74,6 @@ use crate::dtype::DType; use crate::dtype::Nullability; use crate::executor::ExecutionResult; use crate::scalar::Scalar; -use crate::serde::ArrayChildren; use crate::validity::Validity; /// An [`Interleave`]-encoded Vortex array. See the [module docs](self) for the specification. @@ -364,25 +362,6 @@ impl VTable for Interleave { } } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - vortex_bail!("Interleave array is not serializable") - } - - fn deserialize( - &self, - _dtype: &DType, - _len: usize, - _metadata: &[u8], - _buffers: &[BufferHandle], - _children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_bail!("Interleave array is not serializable") - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { execute::execute(array, ctx) } diff --git a/vortex-array/src/arrays/list/vtable/mod.rs b/vortex-array/src/arrays/list/vtable/mod.rs index c55e7050351..5542a19a052 100644 --- a/vortex-array/src/arrays/list/vtable/mod.rs +++ b/vortex-array/src/arrays/list/vtable/mod.rs @@ -4,13 +4,11 @@ use std::hash::Hasher; use std::sync::Arc; -use prost::Message; use vortex_error::VortexExpect; 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::ArrayEq; @@ -26,7 +24,6 @@ use crate::array::ArrayParts; use crate::array::ArrayView; use crate::array::VTable; use crate::array::with_empty_buffers; -use crate::arrays::list::ListArraySlotsExt; use crate::arrays::list::ListData; use crate::arrays::list::ListSlots; use crate::arrays::list::compute::rules::PARENT_RULES; @@ -34,12 +31,10 @@ use crate::arrays::listview::list_view_from_list; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::dtype::DType; -use crate::dtype::Nullability; use crate::dtype::PType; use crate::match_each_list_builder; -use crate::serde::ArrayChildren; -use crate::validity::Validity; mod operations; +mod plugin; mod validity; /// A [`List`]-encoded Vortex array. pub type ListArray = Array; @@ -100,19 +95,6 @@ impl VTable for List { PARENT_RULES.evaluate(array, parent, child_idx) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - ListMetadata { - elements_len: array.elements().len() as u64, - offset_ptype: PType::try_from(array.offsets().dtype())? as i32, - } - .encode_to_vec(), - )) - } - fn validate( &self, _data: &ListData, @@ -150,46 +132,6 @@ impl VTable for List { Ok(()) } - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - let metadata = ListMetadata::decode(metadata)?; - let validity = if children.len() == 2 { - Validity::from(dtype.nullability()) - } else if children.len() == 3 { - let validity = children.get(2, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("Expected 2 or 3 children, got {}", children.len()); - }; - - let DType::List(element_dtype, _) = &dtype else { - vortex_bail!("Expected List dtype, got {:?}", dtype); - }; - let elements = children.get( - 0, - element_dtype.as_ref(), - usize::try_from(metadata.elements_len)?, - )?; - - let offsets = children.get( - 1, - &DType::Primitive(metadata.offset_ptype(), Nullability::NonNullable), - len + 1, - )?; - - let data = ListData::try_build(elements.clone(), offsets.clone(), validity.clone())?; - let slots = ListData::make_slots(&elements, &offsets, &validity, len); - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { ListSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/list/vtable/plugin.rs b/vortex-array/src/arrays/list/vtable/plugin.rs new file mode 100644 index 00000000000..f8a018a011e --- /dev/null +++ b/vortex-array/src/arrays/list/vtable/plugin.rs @@ -0,0 +1,103 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use prost::Message; +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::ListMetadata; +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::VTable; +use crate::arrays::List; +use crate::arrays::list::ListArraySlotsExt; +use crate::arrays::list::ListData; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::validity::Validity; + +impl ArrayPlugin for List { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array + .as_opt::() + .ok_or_else(|| vortex_err!("List plugin cannot serialize {}", array.encoding_id()))?; + let metadata = ListMetadata { + elements_len: view.elements().len() as u64, + offset_ptype: PType::try_from(view.offsets().dtype())? as i32, + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers: _, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "List plugin does not recognize serialized ID {serialized_id}" + ); + let metadata = ListMetadata::decode(metadata)?; + let validity = if children.len() == 2 { + Validity::from(dtype.nullability()) + } else if children.len() == 3 { + let validity = children.get(2, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + vortex_bail!("Expected 2 or 3 children, got {}", children.len()); + }; + + let DType::List(element_dtype, _) = &dtype else { + vortex_bail!("Expected List dtype, got {:?}", dtype); + }; + let elements = children.get( + 0, + element_dtype.as_ref(), + usize::try_from(metadata.elements_len)?, + )?; + + let offsets = children.get( + 1, + &DType::Primitive(metadata.offset_ptype(), Nullability::NonNullable), + len + 1, + )?; + + let data = ListData::try_build(elements.clone(), offsets.clone(), validity.clone())?; + let slots = ListData::make_slots(&elements, &offsets, &validity, len); + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/listview/vtable/mod.rs b/vortex-array/src/arrays/listview/vtable/mod.rs index 806e7a6563c..0fdc7117ae2 100644 --- a/vortex-array/src/arrays/listview/vtable/mod.rs +++ b/vortex-array/src/arrays/listview/vtable/mod.rs @@ -5,7 +5,6 @@ use std::hash::Hash; use std::hash::Hasher; use std::sync::Arc; -use prost::Message; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -26,20 +25,17 @@ use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; use crate::array::with_empty_buffers; -use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::ListViewData; use crate::arrays::listview::ListViewSlots; use crate::arrays::listview::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::dtype::DType; -use crate::dtype::Nullability; use crate::dtype::PType; use crate::match_each_list_builder; -use crate::serde::ArrayChildren; -use crate::validity::Validity; mod kernel; mod operations; +mod plugin; mod validity; /// A [`ListView`]-encoded Vortex array. pub type ListViewArray = Array; @@ -103,20 +99,6 @@ impl VTable for ListView { with_empty_buffers(self, array, buffers) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - ListViewMetadata { - elements_len: array.elements().len() as u64, - offset_ptype: PType::try_from(array.offsets().dtype())? as i32, - size_ptype: PType::try_from(array.sizes().dtype())? as i32, - } - .encode_to_vec(), - )) - } - fn validate( &self, _data: &ListViewData, @@ -157,65 +139,6 @@ impl VTable for ListView { Ok(()) } - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - let metadata = ListViewMetadata::decode(metadata)?; - vortex_ensure!( - buffers.is_empty(), - "`ListViewArray::build` expects no buffers" - ); - - let DType::List(element_dtype, _) = dtype else { - vortex_bail!("Expected List dtype, got {:?}", dtype); - }; - - let validity = if children.len() == 3 { - Validity::from(dtype.nullability()) - } else if children.len() == 4 { - let validity = children.get(3, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!( - "`ListViewArray::build` expects 3 or 4 children, got {}", - children.len() - ); - }; - - // Get elements with the correct length from metadata. - let elements = children.get( - 0, - element_dtype.as_ref(), - usize::try_from(metadata.elements_len)?, - )?; - - // Get offsets with proper type from metadata. - let offsets = children.get( - 1, - &DType::Primitive(metadata.offset_ptype(), Nullability::NonNullable), - len, - )?; - - // Get sizes with proper type from metadata. - let sizes = children.get( - 2, - &DType::Primitive(metadata.size_ptype(), Nullability::NonNullable), - len, - )?; - - ListViewData::validate(&elements, &offsets, &sizes, &validity)?; - let data = ListViewData::try_new()?; - let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len); - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { ListViewSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/listview/vtable/plugin.rs b/vortex-array/src/arrays/listview/vtable/plugin.rs new file mode 100644 index 00000000000..f4af2b7798f --- /dev/null +++ b/vortex-array/src/arrays/listview/vtable/plugin.rs @@ -0,0 +1,123 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use prost::Message; +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::ListViewMetadata; +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::VTable; +use crate::arrays::ListView; +use crate::arrays::listview::ListViewArraySlotsExt; +use crate::arrays::listview::ListViewData; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::validity::Validity; + +impl ArrayPlugin for ListView { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_opt::().ok_or_else(|| { + vortex_err!("ListView plugin cannot serialize {}", array.encoding_id()) + })?; + let metadata = ListViewMetadata { + elements_len: view.elements().len() as u64, + offset_ptype: PType::try_from(view.offsets().dtype())? as i32, + size_ptype: PType::try_from(view.sizes().dtype())? as i32, + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "ListView plugin does not recognize serialized ID {serialized_id}" + ); + let metadata = ListViewMetadata::decode(metadata)?; + vortex_ensure!( + buffers.is_empty(), + "`ListViewArray::build` expects no buffers" + ); + + let DType::List(element_dtype, _) = dtype else { + vortex_bail!("Expected List dtype, got {:?}", dtype); + }; + + let validity = if children.len() == 3 { + Validity::from(dtype.nullability()) + } else if children.len() == 4 { + let validity = children.get(3, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + vortex_bail!( + "`ListViewArray::build` expects 3 or 4 children, got {}", + children.len() + ); + }; + + // Get elements with the correct length from metadata. + let elements = children.get( + 0, + element_dtype.as_ref(), + usize::try_from(metadata.elements_len)?, + )?; + + // Get offsets with proper type from metadata. + let offsets = children.get( + 1, + &DType::Primitive(metadata.offset_ptype(), Nullability::NonNullable), + len, + )?; + + // Get sizes with proper type from metadata. + let sizes = children.get( + 2, + &DType::Primitive(metadata.size_ptype(), Nullability::NonNullable), + len, + )?; + + ListViewData::validate(&elements, &offsets, &sizes, &validity)?; + let data = ListViewData::try_new()?; + let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len); + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/map/vtable/mod.rs b/vortex-array/src/arrays/map/vtable/mod.rs index f11e3e72a89..eb913844803 100644 --- a/vortex-array/src/arrays/map/vtable/mod.rs +++ b/vortex-array/src/arrays/map/vtable/mod.rs @@ -18,7 +18,6 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::array::ValidityVTableFromChild; use crate::array::with_empty_buffers; -use crate::arrays::ListView; use crate::arrays::map::MapData; use crate::arrays::map::MapSlots; use crate::arrays::map::MapSlotsView; @@ -28,10 +27,10 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::dtype::DType; use crate::match_each_map_builder; -use crate::serde::ArrayChildren; mod kernel; mod operations; +mod plugin; mod validity; /// A [`Map`]-encoded Vortex array. @@ -100,53 +99,6 @@ impl VTable for Map { with_empty_buffers(self, array, buffers) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - if !metadata.is_empty() { - vortex_bail!( - "MapArray expects empty metadata, got {} bytes", - metadata.len() - ); - } - vortex_ensure!(buffers.is_empty(), "MapArray expects no buffers"); - - let DType::Map(map_dtype, nullability) = dtype else { - vortex_bail!("Expected map dtype, got {dtype}"); - }; - vortex_ensure!( - children.len() == MapSlots::COUNT, - "MapArray expected {} child, found {}", - MapSlots::COUNT, - children.len() - ); - - let expected_entries_dtype = - DType::List(std::sync::Arc::new(map_dtype.entries_dtype()), *nullability); - let entries = children.get(MapSlots::ENTRIES, &expected_entries_dtype, len)?; - vortex_ensure!( - entries.is::(), - "MapArray entries must use vortex.listview encoding, got {}", - entries.encoding_id() - ); - - let slots = MapData::make_slots(entries); - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, MapData).with_slots(slots)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { MapSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/map/vtable/plugin.rs b/vortex-array/src/arrays/map/vtable/plugin.rs new file mode 100644 index 00000000000..02d41573145 --- /dev/null +++ b/vortex-array/src/arrays/map/vtable/plugin.rs @@ -0,0 +1,96 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_session::VortexSession; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::VTable; +use crate::arrays::ListView; +use crate::arrays::Map; +use crate::arrays::map::MapData; +use crate::arrays::map::MapSlots; +use crate::dtype::DType; + +impl ArrayPlugin for Map { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + array.is::(), + "Map plugin cannot serialize {}", + array.encoding_id() + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Map plugin does not recognize serialized ID {serialized_id}" + ); + if !metadata.is_empty() { + vortex_bail!( + "MapArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + vortex_ensure!(buffers.is_empty(), "MapArray expects no buffers"); + + let DType::Map(map_dtype, nullability) = dtype else { + vortex_bail!("Expected map dtype, got {dtype}"); + }; + vortex_ensure!( + children.len() == MapSlots::COUNT, + "MapArray expected {} child, found {}", + MapSlots::COUNT, + children.len() + ); + + let expected_entries_dtype = + DType::List(std::sync::Arc::new(map_dtype.entries_dtype()), *nullability); + let entries = children.get(MapSlots::ENTRIES, &expected_entries_dtype, len)?; + vortex_ensure!( + entries.is::(), + "MapArray entries must use vortex.listview encoding, got {}", + entries.encoding_id() + ); + + let slots = MapData::make_slots(entries); + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, MapData).with_slots(slots), + )? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/masked/vtable/mod.rs b/vortex-array/src/arrays/masked/vtable/mod.rs index ccadee27e98..6e57df24d54 100644 --- a/vortex-array/src/arrays/masked/vtable/mod.rs +++ b/vortex-array/src/arrays/masked/vtable/mod.rs @@ -2,17 +2,15 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod canonical; mod operations; +mod plugin; mod validity; use std::hash::Hasher; -use smallvec::smallvec; use vortex_error::VortexExpect; 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::AnyCanonical; @@ -28,7 +26,6 @@ use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; -use crate::array::validity_to_child; use crate::array::with_empty_buffers; use crate::arrays::ConstantArray; use crate::arrays::masked::MaskedArrayExt; @@ -44,8 +41,6 @@ use crate::executor::ExecutionResult; use crate::legacy_session; use crate::require_child; use crate::scalar::Scalar; -use crate::serde::ArrayChildren; -use crate::validity::Validity; /// A [`Masked`]-encoded Vortex array. pub type MaskedArray = Array; @@ -120,59 +115,6 @@ impl VTable for Masked { with_empty_buffers(self, array, buffers) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - - #[allow(clippy::disallowed_methods)] - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - if !metadata.is_empty() { - vortex_bail!( - "MaskedArray expects empty metadata, got {} bytes", - metadata.len() - ); - } - if !buffers.is_empty() { - vortex_bail!("Expected 0 buffer, got {}", buffers.len()); - } - - vortex_ensure!( - children.len() == 1 || children.len() == 2, - "`MaskedArray::build` expects 1 or 2 children, got {}", - children.len() - ); - - let child = children.get(0, &dtype.as_nonnullable(), len)?; - - let validity = if children.len() == 2 { - let validity = children.get(1, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - Validity::from(dtype.nullability()) - }; - - let validity_slot = validity_to_child(&validity, len); - let data = MaskedData::try_new( - len, - child.all_valid(&mut legacy_session().create_execution_ctx())?, - validity, - )?; - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data) - .with_slots(smallvec![Some(child), validity_slot])) - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { let array = require_child!(array, array.child(), MaskedSlots::CHILD => AnyCanonical); diff --git a/vortex-array/src/arrays/masked/vtable/plugin.rs b/vortex-array/src/arrays/masked/vtable/plugin.rs new file mode 100644 index 00000000000..9acff199536 --- /dev/null +++ b/vortex-array/src/arrays/masked/vtable/plugin.rs @@ -0,0 +1,103 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use smallvec::smallvec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_session::VortexSession; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::VTable; +use crate::array::validity_to_child; +use crate::arrays::Masked; +use crate::arrays::masked::MaskedData; +use crate::legacy_session; +use crate::validity::Validity; + +impl ArrayPlugin for Masked { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + array.is::(), + "Masked plugin cannot serialize {}", + array.encoding_id() + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + #[allow(clippy::disallowed_methods)] + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Masked plugin does not recognize serialized ID {serialized_id}" + ); + if !metadata.is_empty() { + vortex_bail!( + "MaskedArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + if !buffers.is_empty() { + vortex_bail!("Expected 0 buffer, got {}", buffers.len()); + } + + vortex_ensure!( + children.len() == 1 || children.len() == 2, + "`MaskedArray::build` expects 1 or 2 children, got {}", + children.len() + ); + + let child = children.get(0, &dtype.as_nonnullable(), len)?; + + let validity = if children.len() == 2 { + let validity = children.get(1, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + Validity::from(dtype.nullability()) + }; + + let validity_slot = validity_to_child(&validity, len); + let data = MaskedData::try_new( + len, + child.all_valid(&mut legacy_session().create_execution_ctx())?, + validity, + )?; + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data) + .with_slots(smallvec![Some(child), validity_slot]), + )? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/null/mod.rs b/vortex-array/src/arrays/null/mod.rs index b4d4536c9cf..585308d0975 100644 --- a/vortex-array/src/arrays/null/mod.rs +++ b/vortex-array/src/arrays/null/mod.rs @@ -5,7 +5,6 @@ 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::ArrayRef; @@ -26,10 +25,10 @@ use crate::builders::ArrayBuilder; use crate::builders::NullBuilder; use crate::dtype::DType; use crate::scalar::Scalar; -use crate::serde::ArrayChildren; use crate::validity::Validity; pub(crate) mod compute; +mod plugin; /// A [`Null`]-encoded Vortex array. pub type NullArray = Array; @@ -80,36 +79,6 @@ impl VTable for Null { vortex_panic!("NullArray slot_name index {idx} out of bounds") } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - - _buffers: &[BufferHandle], - _children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_ensure!( - metadata.is_empty(), - "NullArray expects empty metadata, got {} bytes", - metadata.len() - ); - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - EmptyArrayData, - )) - } - fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, diff --git a/vortex-array/src/arrays/null/plugin.rs b/vortex-array/src/arrays/null/plugin.rs new file mode 100644 index 00000000000..753d4eb696b --- /dev/null +++ b/vortex-array/src/arrays/null/plugin.rs @@ -0,0 +1,72 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_session::VortexSession; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::EmptyArrayData; +use crate::array::VTable; +use crate::arrays::Null; + +impl ArrayPlugin for Null { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + array.is::(), + "Null plugin cannot serialize {}", + array.encoding_id() + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers: _, + children: _, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Null plugin does not recognize serialized ID {serialized_id}" + ); + vortex_ensure!( + metadata.is_empty(), + "NullArray expects empty metadata, got {} bytes", + metadata.len() + ); + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + EmptyArrayData, + ))? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs index e7d0adae0dd..56be70d40b0 100644 --- a/vortex-array/src/arrays/patched/vtable/mod.rs +++ b/vortex-array/src/arrays/patched/vtable/mod.rs @@ -1,12 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use prost::Message; - use crate::ArrayEq; use crate::ArrayHash; mod kernels; mod operations; +mod plugin; mod slice; use std::hash::Hash; @@ -47,10 +46,8 @@ use crate::builders::ArrayBuilder; use crate::builders::PrimitiveBuilder; use crate::dtype::DType; use crate::dtype::NativePType; -use crate::dtype::PType; use crate::match_each_native_ptype; use crate::require_child; -use crate::serde::ArrayChildren; /// A [`Patched`]-encoded Vortex array. pub type PatchedArray = Array; @@ -138,54 +135,6 @@ impl VTable for Patched { with_empty_buffers(self, array, buffers) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - PatchedMetadata { - n_patches: u32::try_from(array.patch_indices().len())?, - n_lanes: u32::try_from(array.n_lanes())?, - offset: u32::try_from(array.offset())?, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - let metadata = PatchedMetadata::decode(metadata)?; - let n_patches = metadata.n_patches as usize; - let n_lanes = metadata.n_lanes as usize; - let offset = metadata.offset as usize; - - // n_chunks should correspond to the chunk in the `inner`. - // After slicing when offset > 0, there may be additional chunks. - let n_chunks = (len + offset).div_ceil(1024); - - let inner = children.get(0, dtype, len)?; - let lane_offsets = children.get(1, PType::U32.into(), n_chunks * n_lanes + 1)?; - let indices = children.get(2, PType::U16.into(), n_patches)?; - let values = children.get(3, dtype, n_patches)?; - - let data = PatchedData { n_lanes, offset }; - let slots = PatchedSlots { - inner, - lane_offsets, - patch_indices: indices, - patch_values: values, - } - .into_slots(); - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, diff --git a/vortex-array/src/arrays/patched/vtable/plugin.rs b/vortex-array/src/arrays/patched/vtable/plugin.rs new file mode 100644 index 00000000000..6bbde6b309a --- /dev/null +++ b/vortex-array/src/arrays/patched/vtable/plugin.rs @@ -0,0 +1,97 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use prost::Message; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use super::PatchedMetadata; +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::VTable; +use crate::arrays::Patched; +use crate::arrays::patched::PatchedArrayExt; +use crate::arrays::patched::PatchedArraySlotsExt; +use crate::arrays::patched::PatchedData; +use crate::arrays::patched::PatchedSlots; +use crate::dtype::PType; + +impl ArrayPlugin for Patched { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_opt::().ok_or_else(|| { + vortex_err!("Patched plugin cannot serialize {}", array.encoding_id()) + })?; + let metadata = PatchedMetadata { + n_patches: u32::try_from(view.patch_indices().len())?, + n_lanes: u32::try_from(view.n_lanes())?, + offset: u32::try_from(view.offset())?, + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers: _, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Patched plugin does not recognize serialized ID {serialized_id}" + ); + let metadata = PatchedMetadata::decode(metadata)?; + let n_patches = metadata.n_patches as usize; + let n_lanes = metadata.n_lanes as usize; + let offset = metadata.offset as usize; + + // n_chunks should correspond to the chunk in the `inner`. + // After slicing when offset > 0, there may be additional chunks. + let n_chunks = (len + offset).div_ceil(1024); + + let inner = children.get(0, dtype, len)?; + let lane_offsets = children.get(1, PType::U32.into(), n_chunks * n_lanes + 1)?; + let indices = children.get(2, PType::U16.into(), n_patches)?; + let values = children.get(3, dtype, n_patches)?; + + let data = PatchedData { n_lanes, offset }; + let slots = PatchedSlots { + inner, + lane_offsets, + patch_indices: indices, + patch_values: values, + } + .into_slots(); + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/piecewise_sequence/tests.rs b/vortex-array/src/arrays/piecewise_sequence/tests.rs index a087404fee2..f1fb83f3257 100644 --- a/vortex-array/src/arrays/piecewise_sequence/tests.rs +++ b/vortex-array/src/arrays/piecewise_sequence/tests.rs @@ -171,7 +171,7 @@ fn serialization_is_not_supported() -> VortexResult<()> { .unwrap_err(); assert!( err.to_string() - .contains("Array vortex.piecewise-sequence does not support serialization"), + .contains("Array vortex.piecewise-sequence is not registered for serialization"), "{err}" ); Ok(()) diff --git a/vortex-array/src/arrays/piecewise_sequence/vtable.rs b/vortex-array/src/arrays/piecewise_sequence/vtable.rs index d39bd6eccac..bcc38255956 100644 --- a/vortex-array/src/arrays/piecewise_sequence/vtable.rs +++ b/vortex-array/src/arrays/piecewise_sequence/vtable.rs @@ -7,7 +7,6 @@ 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::ArrayParts; @@ -35,7 +34,6 @@ use crate::dtype::PType; use crate::dtype::UnsignedPType; use crate::match_each_unsigned_integer_ptype; use crate::scalar::Scalar; -use crate::serde::ArrayChildren; use crate::validity::Validity; /// A [`PiecewiseSequence`]-encoded Vortex index array. @@ -109,25 +107,6 @@ impl VTable for PiecewiseSequence { PiecewiseSequenceSlots::NAMES[idx].to_string() } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(None) - } - - fn deserialize( - &self, - _dtype: &DType, - _len: usize, - _metadata: &[u8], - _buffers: &[BufferHandle], - _children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_bail!("PiecewiseSequenceArray is not serializable") - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { let (starts, lengths, multipliers) = execute_index_arrays(array.as_view(), ctx)?; diff --git a/vortex-array/src/arrays/primitive/vtable/mod.rs b/vortex-array/src/arrays/primitive/vtable/mod.rs index 45256a7a498..b8ea41c507e 100644 --- a/vortex-array/src/arrays/primitive/vtable/mod.rs +++ b/vortex-array/src/arrays/primitive/vtable/mod.rs @@ -18,16 +18,14 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::PrimitiveBuilder; use crate::dtype::DType; -use crate::dtype::PType; use crate::match_each_native_ptype; -use crate::serde::ArrayChildren; mod kernel; mod operations; +mod plugin; mod validity; use std::hash::Hasher; -use vortex_buffer::Alignment; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -93,13 +91,6 @@ impl VTable for Primitive { ) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - fn validate( &self, data: &PrimitiveData, @@ -130,49 +121,6 @@ impl VTable for Primitive { Ok(()) } - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - if !metadata.is_empty() { - vortex_bail!( - "PrimitiveArray expects empty metadata, got {} bytes", - metadata.len() - ); - } - let buffer = fixed_width::single_buffer(buffers)?; - - let validity = fixed_width::deserialize_validity(dtype.nullability(), len, children)?; - - let ptype = PType::try_from(dtype)?; - - vortex_ensure!( - buffer.is_aligned_to(Alignment::new(ptype.byte_width())), - "Misaligned buffer cannot be used to build PrimitiveArray of {ptype}" - ); - - if buffer.len() != ptype.byte_width() * len { - vortex_bail!( - "Buffer length {} does not match expected length {} for {}, {}", - buffer.len(), - ptype.byte_width() * len, - ptype.byte_width(), - len, - ); - } - - // SAFETY: the buffer length and alignment are checked above. - let slots = PrimitiveData::make_slots(&validity, len); - let data = unsafe { PrimitiveData::new_unchecked_from_handle(buffer, ptype, validity) }; - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { PrimitiveSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/primitive/vtable/plugin.rs b/vortex-array/src/arrays/primitive/vtable/plugin.rs new file mode 100644 index 00000000000..4213451f4dd --- /dev/null +++ b/vortex-array/src/arrays/primitive/vtable/plugin.rs @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::Alignment; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_session::VortexSession; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::VTable; +use crate::arrays::Primitive; +use crate::arrays::fixed_width::vtable as fixed_width; +use crate::arrays::primitive::PrimitiveData; +use crate::dtype::PType; + +impl ArrayPlugin for Primitive { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + array.is::(), + "Primitive plugin cannot serialize {}", + array.encoding_id() + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Primitive plugin does not recognize serialized ID {serialized_id}" + ); + if !metadata.is_empty() { + vortex_bail!( + "PrimitiveArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + let buffer = fixed_width::single_buffer(buffers)?; + + let validity = fixed_width::deserialize_validity(dtype.nullability(), len, children)?; + + let ptype = PType::try_from(dtype)?; + + vortex_ensure!( + buffer.is_aligned_to(Alignment::new(ptype.byte_width())), + "Misaligned buffer cannot be used to build PrimitiveArray of {ptype}" + ); + + if buffer.len() != ptype.byte_width() * len { + vortex_bail!( + "Buffer length {} does not match expected length {} for {}, {}", + buffer.len(), + ptype.byte_width() * len, + ptype.byte_width(), + len, + ); + } + + // SAFETY: the buffer length and alignment are checked above. + let slots = PrimitiveData::make_slots(&validity, len); + let data = unsafe { PrimitiveData::new_unchecked_from_handle(buffer, ptype, validity) }; + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs index c84a36a3169..d88ed0a6bd4 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs @@ -11,10 +11,8 @@ use std::ops::Deref; use itertools::Itertools; 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::ArrayEq; @@ -45,7 +43,6 @@ use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::VecExecutionArgs; -use crate::serde::ArrayChildren; /// A [`ScalarFn`]-encoded Vortex array. pub type ScalarFnArray = Array; @@ -138,26 +135,6 @@ impl VTable for ScalarFn { with_empty_buffers(self, array, buffers) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - // Not supported - Ok(None) - } - - fn deserialize( - &self, - _dtype: &DType, - _len: usize, - _metadata: &[u8], - _buffers: &[BufferHandle], - _children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_bail!("Deserialization of ScalarFnVTable metadata is not supported"); - } - fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { array .scalar_fn() diff --git a/vortex-array/src/arrays/shared/vtable.rs b/vortex-array/src/arrays/shared/vtable.rs index b1a9a4d7b15..64f6a313724 100644 --- a/vortex-array/src/arrays/shared/vtable.rs +++ b/vortex-array/src/arrays/shared/vtable.rs @@ -6,7 +6,6 @@ use std::hash::Hasher; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::ArrayEq; @@ -98,26 +97,6 @@ impl VTable for Shared { SharedSlots::NAMES[idx].to_string() } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - vortex_error::vortex_bail!("Shared array is not serializable") - } - - fn deserialize( - &self, - _dtype: &DType, - _len: usize, - _metadata: &[u8], - - _buffers: &[BufferHandle], - _children: &dyn crate::serde::ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_error::vortex_bail!("Shared array is not serializable") - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { array .get_or_compute(|source| source.clone().execute::(ctx)) diff --git a/vortex-array/src/arrays/slice/vtable.rs b/vortex-array/src/arrays/slice/vtable.rs index c09d81ee4f1..6f53b31398f 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -9,10 +9,8 @@ use std::ops::Range; use vortex_error::VortexExpect; 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::AnyCanonical; @@ -38,7 +36,6 @@ use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; use crate::require_child; use crate::scalar::Scalar; -use crate::serde::ArrayChildren; use crate::validity::Validity; /// A [`Slice`]-encoded Vortex array. @@ -128,27 +125,6 @@ impl VTable for Slice { SliceSlots::NAMES[idx].to_string() } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - // TODO(joe): make this configurable - vortex_bail!("Slice array is not serializable") - } - - fn deserialize( - &self, - _dtype: &DType, - _len: usize, - _metadata: &[u8], - - _buffers: &[BufferHandle], - _children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_bail!("Slice array is not serializable") - } - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { let array = require_child!(array, array.child(), SliceSlots::CHILD => AnyCanonical); diff --git a/vortex-array/src/arrays/struct_/vtable/mod.rs b/vortex-array/src/arrays/struct_/vtable/mod.rs index c1639501288..a05671981f5 100644 --- a/vortex-array/src/arrays/struct_/vtable/mod.rs +++ b/vortex-array/src/arrays/struct_/vtable/mod.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_panic; @@ -18,16 +17,14 @@ use crate::array::VTable; use crate::array::child_to_validity; use crate::array::with_empty_buffers; use crate::arrays::struct_::array::StructSlots; -use crate::arrays::struct_::array::struct_slots_with_capacity; use crate::arrays::struct_::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::StructBuilder; use crate::dtype::DType; -use crate::serde::ArrayChildren; -use crate::validity::Validity; mod kernel; mod operations; +mod plugin; mod validity; use vortex_session::registry::CachedId; @@ -131,62 +128,6 @@ impl VTable for Struct { with_empty_buffers(self, array, buffers) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - if !metadata.is_empty() { - vortex_bail!( - "StructArray expects empty metadata, got {} bytes", - metadata.len() - ); - } - let DType::Struct(struct_dtype, nullability) = dtype else { - vortex_bail!("Expected struct dtype, found {:?}", dtype) - }; - - let (validity, non_data_children) = if children.len() == struct_dtype.nfields() { - (Validity::from(*nullability), 0_usize) - } else if children.len() == struct_dtype.nfields() + 1 { - let validity = children.get(0, &Validity::DTYPE, len)?; - (Validity::Array(validity), 1_usize) - } else { - vortex_bail!( - "Expected {} or {} children, found {}", - struct_dtype.nfields(), - struct_dtype.nfields() + 1, - children.len() - ); - }; - - let mut slots = struct_slots_with_capacity(&validity, len, struct_dtype.nfields()); - for i in 0..struct_dtype.nfields() { - let child_dtype = struct_dtype - .field_by_index(i) - .vortex_expect("no out of bounds"); - slots.push(Some(children.get( - non_data_children + i, - &child_dtype, - len, - )?)); - } - - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots(slots)) - } - fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { if idx == StructSlots::VALIDITY { "validity".to_string() diff --git a/vortex-array/src/arrays/struct_/vtable/plugin.rs b/vortex-array/src/arrays/struct_/vtable/plugin.rs new file mode 100644 index 00000000000..d3b18cbf9c4 --- /dev/null +++ b/vortex-array/src/arrays/struct_/vtable/plugin.rs @@ -0,0 +1,105 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +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::IntoArray; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::EmptyArrayData; +use crate::array::VTable; +use crate::arrays::Struct; +use crate::arrays::struct_::array::struct_slots_with_capacity; +use crate::dtype::DType; +use crate::validity::Validity; + +impl ArrayPlugin for Struct { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + array.is::(), + "Struct plugin cannot serialize {}", + array.encoding_id() + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers: _, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Struct plugin does not recognize serialized ID {serialized_id}" + ); + if !metadata.is_empty() { + vortex_bail!( + "StructArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + let DType::Struct(struct_dtype, nullability) = dtype else { + vortex_bail!("Expected struct dtype, found {:?}", dtype) + }; + + let (validity, non_data_children) = if children.len() == struct_dtype.nfields() { + (Validity::from(*nullability), 0_usize) + } else if children.len() == struct_dtype.nfields() + 1 { + let validity = children.get(0, &Validity::DTYPE, len)?; + (Validity::Array(validity), 1_usize) + } else { + vortex_bail!( + "Expected {} or {} children, found {}", + struct_dtype.nfields(), + struct_dtype.nfields() + 1, + children.len() + ); + }; + + let mut slots = struct_slots_with_capacity(&validity, len, struct_dtype.nfields()); + for i in 0..struct_dtype.nfields() { + let child_dtype = struct_dtype + .field_by_index(i) + .vortex_expect("no out of bounds"); + slots.push(Some(children.get( + non_data_children + i, + &child_dtype, + len, + )?)); + } + + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots(slots), + )? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/union/vtable/mod.rs b/vortex-array/src/arrays/union/vtable/mod.rs index 1f29b950b0a..86adcc48c15 100644 --- a/vortex-array/src/arrays/union/vtable/mod.rs +++ b/vortex-array/src/arrays/union/vtable/mod.rs @@ -2,12 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_error::vortex_ensure; -use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::ArrayRef; @@ -22,15 +18,13 @@ use crate::array::VTable; use crate::array::with_empty_buffers; use crate::arrays::union::UnionArrayExt; use crate::arrays::union::UnionSlots; -use crate::arrays::union::array::make_union_parts; use crate::arrays::union::compute::rules::PARENT_RULES; -use crate::arrays::union::union_type_ids_dtype; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::dtype::DType; -use crate::serde::ArrayChildren; mod operations; +mod plugin; mod validate; mod validity; @@ -100,53 +94,6 @@ impl VTable for Union { with_empty_buffers(self, array, buffers) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_ensure!(metadata.is_empty(), "UnionArray expects empty metadata"); - vortex_ensure!(buffers.is_empty(), "UnionArray expects no buffers"); - let DType::Union(variants, nullability) = dtype else { - vortex_bail!("Expected union dtype, found {dtype}") - }; - vortex_ensure_eq!( - children.len(), - UnionSlots::CHILDREN_OFFSET + variants.len(), - "UnionArray expected {} children, found {}", - UnionSlots::CHILDREN_OFFSET + variants.len(), - children.len() - ); - - let type_ids = children.get( - UnionSlots::TYPE_IDS, - &union_type_ids_dtype(*nullability), - len, - )?; - let sparse_children = variants - .variants() - .enumerate() - .map(|(index, dtype)| children.get(UnionSlots::CHILDREN_OFFSET + index, &dtype, len)) - .collect::>>()?; - - Ok(make_union_parts( - type_ids, - variants.clone(), - sparse_children, - )) - } - fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { if idx == UnionSlots::TYPE_IDS { "type_ids".to_string() diff --git a/vortex-array/src/arrays/union/vtable/plugin.rs b/vortex-array/src/arrays/union/vtable/plugin.rs new file mode 100644 index 00000000000..659948e7f83 --- /dev/null +++ b/vortex-array/src/arrays/union/vtable/plugin.rs @@ -0,0 +1,94 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::VTable; +use crate::arrays::Union; +use crate::arrays::union::UnionSlots; +use crate::arrays::union::array::make_union_parts; +use crate::arrays::union::union_type_ids_dtype; +use crate::dtype::DType; + +impl ArrayPlugin for Union { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + array.is::(), + "Union plugin cannot serialize {}", + array.encoding_id() + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Union plugin does not recognize serialized ID {serialized_id}" + ); + vortex_ensure!(metadata.is_empty(), "UnionArray expects empty metadata"); + vortex_ensure!(buffers.is_empty(), "UnionArray expects no buffers"); + let DType::Union(variants, nullability) = dtype else { + vortex_bail!("Expected union dtype, found {dtype}") + }; + vortex_ensure_eq!( + children.len(), + UnionSlots::CHILDREN_OFFSET + variants.len(), + "UnionArray expected {} children, found {}", + UnionSlots::CHILDREN_OFFSET + variants.len(), + children.len() + ); + + let type_ids = children.get( + UnionSlots::TYPE_IDS, + &union_type_ids_dtype(*nullability), + len, + )?; + let sparse_children = variants + .variants() + .enumerate() + .map(|(index, dtype)| children.get(UnionSlots::CHILDREN_OFFSET + index, &dtype, len)) + .collect::>>()?; + + Ok(Array::try_from_parts(make_union_parts( + type_ids, + variants.clone(), + sparse_children, + ))? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/varbin/vtable/mod.rs b/vortex-array/src/arrays/varbin/vtable/mod.rs index 397b0f8084b..8c30ae6a147 100644 --- a/vortex-array/src/arrays/varbin/vtable/mod.rs +++ b/vortex-array/src/arrays/varbin/vtable/mod.rs @@ -3,7 +3,6 @@ use std::hash::Hasher; -use prost::Message; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -22,22 +21,19 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::arrays::PrimitiveArray; use crate::arrays::varbin::VarBinArrayExt; -use crate::arrays::varbin::VarBinArraySlotsExt; use crate::arrays::varbin::VarBinData; use crate::arrays::varbin::VarBinSlots; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::VarBinViewBuilder; use crate::dtype::DType; -use crate::dtype::Nullability; use crate::dtype::PType; use crate::match_each_integer_ptype; use crate::match_each_varbin_builder; -use crate::serde::ArrayChildren; -use crate::validity::Validity; pub(crate) mod canonical; mod kernel; mod operations; +mod plugin; mod validity; use canonical::varbin_to_canonical; @@ -148,54 +144,6 @@ impl VTable for VarBin { ) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - VarBinMetadata { - offsets_ptype: PType::try_from(array.offsets().dtype()) - .vortex_expect("Must be a valid PType") as i32, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - let metadata = VarBinMetadata::decode(metadata)?; - let validity = if children.len() == 1 { - Validity::from(dtype.nullability()) - } else if children.len() == 2 { - let validity = children.get(1, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("Expected 1 or 2 children, got {}", children.len()); - }; - - let offsets = children.get( - 0, - &DType::Primitive(metadata.offsets_ptype(), Nullability::NonNullable), - len + 1, - )?; - - if buffers.len() != 1 { - vortex_bail!("Expected 1 buffer, got {}", buffers.len()); - } - let bytes = buffers[0].clone().try_to_host_sync()?; - - let data = VarBinData::try_build(offsets.clone(), bytes, dtype.clone(), validity.clone())?; - let slots = VarBinData::make_slots(offsets, &validity, len); - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { VarBinSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/varbin/vtable/plugin.rs b/vortex-array/src/arrays/varbin/vtable/plugin.rs new file mode 100644 index 00000000000..b2d4a7b1d2b --- /dev/null +++ b/vortex-array/src/arrays/varbin/vtable/plugin.rs @@ -0,0 +1,100 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use prost::Message; +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_session::VortexSession; + +use super::VarBinMetadata; +use crate::ArrayParts; +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::VTable; +use crate::arrays::VarBin; +use crate::arrays::varbin::VarBinArraySlotsExt; +use crate::arrays::varbin::VarBinData; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::validity::Validity; + +impl ArrayPlugin for VarBin { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array + .as_opt::() + .ok_or_else(|| vortex_err!("VarBin plugin cannot serialize {}", array.encoding_id()))?; + let metadata = VarBinMetadata { + offsets_ptype: PType::try_from(view.offsets().dtype()) + .vortex_expect("Must be a valid PType") as i32, + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "VarBin plugin does not recognize serialized ID {serialized_id}" + ); + let metadata = VarBinMetadata::decode(metadata)?; + let validity = if children.len() == 1 { + Validity::from(dtype.nullability()) + } else if children.len() == 2 { + let validity = children.get(1, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + vortex_bail!("Expected 1 or 2 children, got {}", children.len()); + }; + + let offsets = children.get( + 0, + &DType::Primitive(metadata.offsets_ptype(), Nullability::NonNullable), + len + 1, + )?; + + if buffers.len() != 1 { + vortex_bail!("Expected 1 buffer, got {}", buffers.len()); + } + let bytes = buffers[0].clone().try_to_host_sync()?; + + let data = VarBinData::try_build(offsets.clone(), bytes, dtype.clone(), validity.clone())?; + let slots = VarBinData::make_slots(offsets, &validity, len); + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index 0d2963cd6b2..90e32e904ba 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -2,14 +2,11 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::hash::Hasher; -use std::mem::size_of; use std::sync::Arc; -use vortex_buffer::Buffer; 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; @@ -19,12 +16,10 @@ use crate::ArrayRef; use crate::EqMode; use crate::ExecutionCtx; use crate::ExecutionResult; -use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; -use crate::arrays::varbinview::BinaryView; use crate::arrays::varbinview::VarBinViewData; use crate::arrays::varbinview::array::VarBinViewSlots; use crate::arrays::varbinview::compute::rules::PARENT_RULES; @@ -35,10 +30,9 @@ use crate::dtype::DType; use crate::hash::ArrayEq; use crate::hash::ArrayHash; use crate::match_each_varbin_builder; -use crate::serde::ArrayChildren; -use crate::validity::Validity; mod kernel; mod operations; +mod plugin; mod validity; /// A [`VarBinView`]-encoded Vortex array. pub type VarBinViewArray = Array; @@ -154,83 +148,6 @@ impl VTable for VarBinView { ) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - if !metadata.is_empty() { - vortex_bail!( - "VarBinViewArray expects empty metadata, got {} bytes", - metadata.len() - ); - } - let Some((views_handle, data_handles)) = buffers.split_last() else { - vortex_bail!("Expected at least 1 buffer, got 0"); - }; - - let validity = if children.is_empty() { - Validity::from(dtype.nullability()) - } else if children.len() == 1 { - let validity = children.get(0, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("Expected 0 or 1 children, got {}", children.len()); - }; - - let views_nbytes = views_handle.len(); - let expected_views_nbytes = len - .checked_mul(size_of::()) - .ok_or_else(|| vortex_err!("views byte length overflow for len={len}"))?; - if views_nbytes != expected_views_nbytes { - vortex_bail!( - "Expected views buffer length {} bytes, got {} bytes", - expected_views_nbytes, - views_nbytes - ); - } - - // If any buffer is on device, skip host validation and use try_new_handle. - if buffers.iter().any(|b| b.is_on_device()) { - let data = VarBinViewData::try_new_handle( - views_handle.clone(), - Arc::from(data_handles.to_vec()), - dtype.clone(), - validity.clone(), - )?; - let slots = VarBinViewData::make_slots(&validity, len); - return Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)); - } - - let data_buffers = data_handles - .iter() - .map(|b| b.as_host().clone()) - .collect::>(); - let views = Buffer::::from_byte_buffer(views_handle.clone().as_host().clone()); - - let data = VarBinViewData::try_new( - views, - Arc::from(data_buffers), - dtype.clone(), - validity.clone(), - &mut session.create_execution_ctx(), - )?; - let slots = VarBinViewData::make_slots(&validity, len); - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { VarBinViewSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/varbinview/vtable/plugin.rs b/vortex-array/src/arrays/varbinview/vtable/plugin.rs new file mode 100644 index 00000000000..fbd63822fd5 --- /dev/null +++ b/vortex-array/src/arrays/varbinview/vtable/plugin.rs @@ -0,0 +1,133 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::mem::size_of; +use std::sync::Arc; + +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use crate::ArrayParts; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::VTable; +use crate::arrays::VarBinView; +use crate::arrays::varbinview::BinaryView; +use crate::arrays::varbinview::VarBinViewData; +use crate::validity::Validity; + +impl ArrayPlugin for VarBinView { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + array.is::(), + "VarBinView plugin cannot serialize {}", + array.encoding_id() + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "VarBinView plugin does not recognize serialized ID {serialized_id}" + ); + if !metadata.is_empty() { + vortex_bail!( + "VarBinViewArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + let Some((views_handle, data_handles)) = buffers.split_last() else { + vortex_bail!("Expected at least 1 buffer, got 0"); + }; + + let validity = if children.is_empty() { + Validity::from(dtype.nullability()) + } else if children.len() == 1 { + let validity = children.get(0, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + vortex_bail!("Expected 0 or 1 children, got {}", children.len()); + }; + + let views_nbytes = views_handle.len(); + let expected_views_nbytes = len + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("views byte length overflow for len={len}"))?; + if views_nbytes != expected_views_nbytes { + vortex_bail!( + "Expected views buffer length {} bytes, got {} bytes", + expected_views_nbytes, + views_nbytes + ); + } + + // If any buffer is on device, skip host validation and use try_new_handle. + if buffers.iter().any(|b| b.is_on_device()) { + let data = VarBinViewData::try_new_handle( + views_handle.clone(), + Arc::from(data_handles.to_vec()), + dtype.clone(), + validity.clone(), + )?; + let slots = VarBinViewData::make_slots(&validity, len); + return Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()); + } + + let data_buffers = data_handles + .iter() + .map(|b| b.as_host().clone()) + .collect::>(); + let views = Buffer::::from_byte_buffer(views_handle.clone().as_host().clone()); + + let data = VarBinViewData::try_new( + views, + Arc::from(data_buffers), + dtype.clone(), + validity.clone(), + &mut session.create_execution_ctx(), + )?; + let slots = VarBinViewData::make_slots(&validity, len); + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/variant/vtable/mod.rs b/vortex-array/src/arrays/variant/vtable/mod.rs index 837bd32a7bf..d7e0efa3997 100644 --- a/vortex-array/src/arrays/variant/vtable/mod.rs +++ b/vortex-array/src/arrays/variant/vtable/mod.rs @@ -3,9 +3,9 @@ mod kernel; mod operations; +mod plugin; mod validity; -use prost::Message; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -35,7 +35,6 @@ use crate::dtype::StructFields; use crate::proto::dtype as pb; use crate::scalar::Scalar; use crate::scalar::ScalarValue; -use crate::serde::ArrayChildren; /// A [`Variant`]-encoded Vortex array. pub type VariantArray = Array; @@ -132,62 +131,6 @@ impl VTable for Variant { with_empty_buffers(self, array, buffers) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - let shredded_dtype = array.slots()[VariantSlots::SHREDDED] - .as_ref() - .map(|shredded| shredded.dtype().try_into()) - .transpose()?; - Ok(Some( - VariantMetadataProto { shredded_dtype }.encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - vortex_ensure!( - buffers.is_empty(), - "VariantArray expects 0 buffers, got {}", - buffers.len() - ); - let proto = VariantMetadataProto::decode(metadata)?; - let shredded_dtype = proto - .shredded_dtype - .as_ref() - .map(|dtype| DType::from_proto(dtype, session)) - .transpose()?; - vortex_ensure!(matches!(dtype, DType::Variant(_)), "Expected Variant DType"); - let expected_children = 1 + usize::from(shredded_dtype.is_some()); - vortex_ensure!( - children.len() == expected_children, - "Expected {} children, got {}", - expected_children, - children.len(), - ); - let core_storage = children.get(0, dtype, len)?; - let shredded = shredded_dtype - .map(|dtype| children.get(1, &dtype, len)) - .transpose()?; - Ok( - ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots( - VariantSlots { - core_storage, - shredded, - } - .into_slots(), - ), - ) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { match VariantSlots::NAMES.get(idx) { Some(name) => (*name).to_string(), diff --git a/vortex-array/src/arrays/variant/vtable/plugin.rs b/vortex-array/src/arrays/variant/vtable/plugin.rs new file mode 100644 index 00000000000..c59ce700735 --- /dev/null +++ b/vortex-array/src/arrays/variant/vtable/plugin.rs @@ -0,0 +1,101 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use prost::Message; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use super::VariantMetadataProto; +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::EmptyArrayData; +use crate::array::VTable; +use crate::arrays::Variant; +use crate::arrays::variant::VariantSlots; +use crate::dtype::DType; + +impl ArrayPlugin for Variant { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_opt::().ok_or_else(|| { + vortex_err!("Variant plugin cannot serialize {}", array.encoding_id()) + })?; + let shredded_dtype = view.slots()[VariantSlots::SHREDDED] + .as_ref() + .map(|shredded| shredded.dtype().try_into()) + .transpose()?; + let metadata = VariantMetadataProto { shredded_dtype }.encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + serialized_id, + dtype, + len, + metadata, + buffers, + children, + } = parts; + vortex_ensure!( + serialized_id == VTable::id(self), + "Variant plugin does not recognize serialized ID {serialized_id}" + ); + vortex_ensure!( + buffers.is_empty(), + "VariantArray expects 0 buffers, got {}", + buffers.len() + ); + let proto = VariantMetadataProto::decode(metadata)?; + let shredded_dtype = proto + .shredded_dtype + .as_ref() + .map(|dtype| DType::from_proto(dtype, session)) + .transpose()?; + vortex_ensure!(matches!(dtype, DType::Variant(_)), "Expected Variant DType"); + let expected_children = 1 + usize::from(shredded_dtype.is_some()); + vortex_ensure!( + children.len() == expected_children, + "Expected {} children, got {}", + expected_children, + children.len(), + ); + let core_storage = children.get(0, dtype, len)?; + let shredded = shredded_dtype + .map(|dtype| children.get(1, &dtype, len)) + .transpose()?; + Ok(Array::try_from_parts( + ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots( + VariantSlots { + core_storage, + shredded, + } + .into_slots(), + ), + )? + .into_array()) + } +} diff --git a/vortex-array/src/serde.rs b/vortex-array/src/serde.rs index 21db181d52e..da7c6c5402f 100644 --- a/vortex-array/src/serde.rs +++ b/vortex-array/src/serde.rs @@ -777,7 +777,6 @@ mod tests { use vortex_session::registry::CachedId; use super::*; - use crate::Array; use crate::ArrayPlugin; use crate::ArraySerialization; use crate::ArrayVTable; @@ -849,16 +848,18 @@ mod tests { "old primitive wire ID cannot represent length {}", parts.len, ); - Ok(Array::::try_from_parts(ArrayVTable::deserialize( + ArrayPlugin::deserialize( &Primitive, - parts.dtype, - parts.len, - parts.metadata, - parts.buffers, - parts.children, + ArrayDeserialization::new( + old_primitive_id(), + parts.dtype, + parts.len, + parts.metadata, + parts.buffers, + parts.children, + ), session, - )?)? - .into_array()) + ) } } diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index e9d0d36a55f..c1eb037de83 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -30,7 +30,6 @@ use crate::arrays::ListView; use crate::arrays::Map; use crate::arrays::Masked; use crate::arrays::Null; -use crate::arrays::PiecewiseSequence; use crate::arrays::Primitive; use crate::arrays::Struct; use crate::arrays::Union; @@ -105,7 +104,6 @@ impl Default for ArraySession { this.register(Dict); this.register(List); this.register(Masked); - this.register(PiecewiseSequence); this.register(VarBin); this diff --git a/vortex-array/src/test_harness/trace/tests.rs b/vortex-array/src/test_harness/trace/tests.rs index bfe99c3857b..0c469da13b2 100644 --- a/vortex-array/src/test_harness/trace/tests.rs +++ b/vortex-array/src/test_harness/trace/tests.rs @@ -57,7 +57,6 @@ use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::like::Like; use crate::scalar_fn::fns::like::LikeOptions; use crate::scalar_fn::fns::operators::Operator; -use crate::serde::ArrayChildren; use crate::session::ArraySession; use crate::test_harness::trace::TraceOptions; use crate::test_harness::trace::TraceResolution; @@ -176,25 +175,6 @@ impl VTable for StackParent { None } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(None) - } - - fn deserialize( - &self, - _dtype: &DType, - _len: usize, - _metadata: &[u8], - _buffers: &[BufferHandle], - _children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_bail!("StackParent cannot be deserialized") - } - fn with_buffers( &self, array: ArrayView<'_, Self>, @@ -285,25 +265,6 @@ impl VTable for StackChild { None } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(None) - } - - fn deserialize( - &self, - _dtype: &DType, - _len: usize, - _metadata: &[u8], - _buffers: &[BufferHandle], - _children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_bail!("StackChild cannot be deserialized") - } - fn with_buffers( &self, array: ArrayView<'_, Self>, diff --git a/vortex-python/src/arrays/py/vtable.rs b/vortex-python/src/arrays/py/vtable.rs index 5697d5a7574..5b1b45fb86d 100644 --- a/vortex-python/src/arrays/py/vtable.rs +++ b/vortex-python/src/arrays/py/vtable.rs @@ -18,19 +18,18 @@ use vortex::array::OperationsVTable; use vortex::array::VTable; use vortex::array::ValidityVTable; use vortex::array::buffer::BufferHandle; -use vortex::array::serde::ArrayChildren; use vortex::array::validity::Validity; use vortex::array::with_empty_buffers; use vortex::dtype::DType; use vortex::error::VortexResult; -use vortex::error::vortex_bail; use vortex::error::vortex_ensure; use vortex::error::vortex_panic; use vortex::scalar::Scalar; -use vortex::session::VortexSession; use crate::arrays::py::PythonArray; +mod plugin; + /// Wrapper struct encapsulating a Python encoding. #[derive(Debug, Clone)] pub struct PythonVTable { @@ -92,26 +91,6 @@ impl VTable for PythonVTable { with_empty_buffers(self, array, buffers) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(None) - } - - fn deserialize( - &self, - _dtype: &DType, - _len: usize, - bytes: &[u8], - _buffers: &[BufferHandle], - _children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - _ = bytes; - vortex_bail!("PythonArray deserialization is not supported"); - } - fn slot_name(_array: ArrayView<'_, Self>, _idx: usize) -> String { vortex_panic!("PythonArray has no slots") } diff --git a/vortex-python/src/arrays/py/vtable/plugin.rs b/vortex-python/src/arrays/py/vtable/plugin.rs new file mode 100644 index 00000000000..7782dbfce74 --- /dev/null +++ b/vortex-python/src/arrays/py/vtable/plugin.rs @@ -0,0 +1,36 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex::array::ArrayDeserialization; +use vortex::array::ArrayId; +use vortex::array::ArrayPlugin; +use vortex::array::ArrayRef; +use vortex::array::ArraySerialization; +use vortex::array::VTable; +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::session::VortexSession; + +use super::PythonVTable; + +impl ArrayPlugin for PythonVTable { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + _array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + Ok(None) + } + + fn deserialize( + &self, + _parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_bail!("PythonArray deserialization is not supported"); + } +}