diff --git a/encodings/alp/src/alp/array.rs b/encodings/alp/src/alp/array.rs index 3eae9b59580..275174a7e82 100644 --- a/encodings/alp/src/alp/array.rs +++ b/encodings/alp/src/alp/array.rs @@ -9,11 +9,14 @@ use std::hash::Hasher; use prost::Message; use vortex_array::Array; +use vortex_array::ArrayDeserialization; use vortex_array::ArrayEq; use vortex_array::ArrayHash; 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::ArrayView; use vortex_array::EqMode; @@ -34,7 +37,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; @@ -113,33 +115,75 @@ impl VTable for ALP { vortex_array::vtable::with_empty_buffers(self, array, buffers) } - fn serialize( + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + ALPSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let array = require_child!(array, array.encoded(), ALPSlots::ENCODED => Primitive); + require_patches!( + array, + ALPSlots::PATCH_INDICES, + ALPSlots::PATCH_VALUES, + ALPSlots::PATCH_CHUNK_OFFSETS + ); + + Ok(ExecutionResult::done( + execute_decompress(array, ctx)?.into_array(), + )) + } + + fn reduce_parent( array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + RULES.evaluate(array, parent, child_idx) + } +} + +impl ArrayPlugin for ALP { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, _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(), - )) + ) -> VortexResult> { + let view = array.as_::(); + 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( + ArrayPlugin::id(self), + array, + metadata, + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + children, + .. + } = parts; + let metadata = ALPMetadata::decode(metadata)?; let encoded_ptype = match &dtype { DType::Primitive(PType::F32, n) => DType::Primitive(PType::I32, *n), @@ -177,39 +221,14 @@ impl VTable for ALP { }, patches, ); - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots, - )) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - ALPSlots::NAMES[idx].to_string() - } - - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - let array = require_child!(array, array.encoded(), ALPSlots::ENCODED => Primitive); - require_patches!( - array, - ALPSlots::PATCH_INDICES, - ALPSlots::PATCH_VALUES, - ALPSlots::PATCH_CHUNK_OFFSETS - ); - - Ok(ExecutionResult::done( - execute_decompress(array, ctx)?.into_array(), - )) - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - RULES.evaluate(array, parent, child_idx) + ))? + .into_array()) } } diff --git a/encodings/alp/src/alp/mod.rs b/encodings/alp/src/alp/mod.rs index 66c30bba019..28acf814586 100644 --- a/encodings/alp/src/alp/mod.rs +++ b/encodings/alp/src/alp/mod.rs @@ -30,7 +30,8 @@ mod tests { use prost::Message; use rstest::rstest; - use vortex_array::Array; + use vortex_array::ArrayDeserialization; + use vortex_array::ArrayPlugin; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; @@ -39,7 +40,6 @@ mod tests { use vortex_array::dtype::PType; use vortex_array::patches::PatchesMetadata; use vortex_array::test_harness::check_metadata; - use vortex_array::vtable::VTable; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -81,16 +81,18 @@ mod tests { let encoded = alp_encode(values.as_view(), Some(Exponents { e: 0, f: 0 }), &mut ctx)?; let sliced = ::slice(encoded.as_view(), start..4096, &mut ctx)? .expect("ALP slices"); - let metadata = ALP::serialize(sliced.as_::(), &session)?.expect("ALP serializes"); - let parts = ALP.deserialize( - sliced.dtype(), - sliced.len(), - &metadata, - &[], - &sliced.children(), + let serialization = ALP.serialize(&sliced, &session)?.expect("ALP serializes"); + let read = ALP.deserialize( + ArrayDeserialization::new( + serialization.serialized_id, + sliced.dtype(), + sliced.len(), + &serialization.metadata, + &[], + &sliced.children(), + ), &session, )?; - let read = Array::::try_from_parts(parts)?.into_array(); let expected = values.into_array().slice(start..4096)?; let patch_index = ((start / 1024 + 1) * 1024).next_multiple_of(7) - start; diff --git a/encodings/alp/src/alp/plugin.rs b/encodings/alp/src/alp/plugin.rs index c331adb4be9..27e3af827b4 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; @@ -18,7 +17,6 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::Patched; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; -use vortex_error::vortex_err; use vortex_session::VortexSession; use crate::ALP; @@ -38,14 +36,6 @@ impl ArrayPlugin for ALPPatchedPlugin { ArrayVTable::id(&ALP) } - fn read_ids(&self) -> Vec { - vec![self.id()] - } - - fn write_ids(&self) -> Vec { - vec![self.id()] - } - fn serialize( &self, array: &ArrayRef, @@ -65,16 +55,7 @@ impl ArrayPlugin for ALPPatchedPlugin { self.id(), "ALP plugin does not recognize 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 = ALP.deserialize(parts, session)?.downcast::(); // 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 e6eeff1d1fa..add68dd3690 100644 --- a/encodings/alp/src/alp_rd/array.rs +++ b/encodings/alp/src/alp_rd/array.rs @@ -11,11 +11,14 @@ use ::alp::MAX_DICT_SIZE; use itertools::Itertools; use prost::Message; use vortex_array::Array; +use vortex_array::ArrayDeserialization; use vortex_array::ArrayEq; use vortex_array::ArrayHash; 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::ArrayView; use vortex_array::EqMode; @@ -37,7 +40,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; @@ -152,40 +154,137 @@ impl VTable for ALPRD { vortex_array::vtable::with_empty_buffers(self, array, buffers) } - fn serialize( + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + ALPRDSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let array = require_child!(array, array.left_parts(), ALPRDSlots::LEFT_PARTS => Primitive); + let array = + require_child!(array, array.right_parts(), ALPRDSlots::RIGHT_PARTS => Primitive); + require_patches!( + array, + ALPRDSlots::PATCH_INDICES, + ALPRDSlots::PATCH_VALUES, + ALPRDSlots::PATCH_CHUNK_OFFSETS + ); + + let dtype = array.dtype().clone(); + let right_bit_width = array.right_bit_width(); + let ALPRDDataParts { + left_parts, + right_parts, + left_parts_dictionary, + left_parts_patches, + } = ALPRDArrayOwnedExt::into_data_parts(array); + let ptype = dtype.as_ptype(); + + let left_parts = left_parts + .try_downcast::() + .ok() + .vortex_expect("ALPRD execute: left_parts is primitive"); + let right_parts = right_parts + .try_downcast::() + .ok() + .vortex_expect("ALPRD execute: right_parts is primitive"); + + // Decode the left_parts using our builtin dictionary. + let left_parts_dict = left_parts_dictionary; + let validity = left_parts + .as_ref() + .validity()? + .execute_mask(left_parts.as_ref().len(), ctx)?; + + let mut left_parts = left_parts.into_buffer_mut::(); + check_left_parts_codes(left_parts.as_mut_slice(), &validity, left_parts_dict.len())?; + + let decoded_array = if ptype == PType::F32 { + PrimitiveArray::new( + alp_rd_decode::( + left_parts, + &left_parts_dict, + right_bit_width, + right_parts.into_buffer_mut::(), + left_parts_patches, + ctx, + )?, + Validity::from_mask(validity, dtype.nullability()), + ) + } else { + PrimitiveArray::new( + alp_rd_decode::( + left_parts, + &left_parts_dict, + right_bit_width, + right_parts.into_buffer_mut::(), + left_parts_patches, + ctx, + )?, + Validity::from_mask(validity, dtype.nullability()), + ) + }; + + Ok(ExecutionResult::done(decoded_array.into_array())) + } + + fn reduce_parent( array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + RULES.evaluate(array, parent, child_idx) + } +} + +impl ArrayPlugin for ALPRD { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, _session: &VortexSession, - ) -> VortexResult>> { - let dict = array + ) -> VortexResult> { + let view = array.as_::(); + let dict = view .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(), - )) + 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( + ArrayPlugin::id(self), + array, + metadata, + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + children, + .. + } = parts; + let metadata = ALPRDMetadata::decode(metadata)?; let left_parts_ptype = metadata.get_left_parts_ptype()?; if children.len() < 2 { @@ -254,94 +353,14 @@ impl VTable for ALPRD { })?, left_parts_patches, ); - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots, - )) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - ALPRDSlots::NAMES[idx].to_string() - } - - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - let array = require_child!(array, array.left_parts(), ALPRDSlots::LEFT_PARTS => Primitive); - let array = - require_child!(array, array.right_parts(), ALPRDSlots::RIGHT_PARTS => Primitive); - require_patches!( - array, - ALPRDSlots::PATCH_INDICES, - ALPRDSlots::PATCH_VALUES, - ALPRDSlots::PATCH_CHUNK_OFFSETS - ); - - let dtype = array.dtype().clone(); - let right_bit_width = array.right_bit_width(); - let ALPRDDataParts { - left_parts, - right_parts, - left_parts_dictionary, - left_parts_patches, - } = ALPRDArrayOwnedExt::into_data_parts(array); - let ptype = dtype.as_ptype(); - - let left_parts = left_parts - .try_downcast::() - .ok() - .vortex_expect("ALPRD execute: left_parts is primitive"); - let right_parts = right_parts - .try_downcast::() - .ok() - .vortex_expect("ALPRD execute: right_parts is primitive"); - - // Decode the left_parts using our builtin dictionary. - let left_parts_dict = left_parts_dictionary; - let validity = left_parts - .as_ref() - .validity()? - .execute_mask(left_parts.as_ref().len(), ctx)?; - - let mut left_parts = left_parts.into_buffer_mut::(); - check_left_parts_codes(left_parts.as_mut_slice(), &validity, left_parts_dict.len())?; - - let decoded_array = if ptype == PType::F32 { - PrimitiveArray::new( - alp_rd_decode::( - left_parts, - &left_parts_dict, - right_bit_width, - right_parts.into_buffer_mut::(), - left_parts_patches, - ctx, - )?, - Validity::from_mask(validity, dtype.nullability()), - ) - } else { - PrimitiveArray::new( - alp_rd_decode::( - left_parts, - &left_parts_dict, - right_bit_width, - right_parts.into_buffer_mut::(), - left_parts_patches, - ctx, - )?, - Validity::from_mask(validity, dtype.nullability()), - ) - }; - - Ok(ExecutionResult::done(decoded_array.into_array())) - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - RULES.evaluate(array, parent, child_idx) + ))? + .into_array()) } } @@ -637,6 +656,7 @@ mod test { use prost::Message; use rstest::rstest; use vortex_array::ArrayDeserialization; + use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -723,11 +743,14 @@ mod test { .encode_to_vec(); let children: Vec = Vec::new(); let result = ALPRD.deserialize( - &DType::Primitive(PType::F64, Nullability::NonNullable), - 0, - &metadata, - &[], - &children, + ArrayDeserialization::new( + VTable::id(&ALPRD), + &DType::Primitive(PType::F64, Nullability::NonNullable), + 0, + &metadata, + &[], + &children, + ), &SESSION, ); assert!(result.is_err_and(|err| { @@ -752,10 +775,9 @@ mod test { metadata: &ALPRDMetadata, children: Vec, ) -> VortexResult { - vortex_array::ArrayPlugin::deserialize( - &ALPRD, + ALPRD.deserialize( ArrayDeserialization::new( - ALPRD.id(), + VTable::id(&ALPRD), dtype, len, &metadata.encode_to_vec(), diff --git a/encodings/bytebool/src/array.rs b/encodings/bytebool/src/array.rs index 88007ecfd01..0d27ab55cbc 100644 --- a/encodings/bytebool/src/array.rs +++ b/encodings/bytebool/src/array.rs @@ -7,11 +7,14 @@ use std::fmt::Formatter; use std::hash::Hasher; use vortex_array::Array; +use vortex_array::ArrayDeserialization; use vortex_array::ArrayEq; use vortex_array::ArrayHash; 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::ArrayView; use vortex_array::EqMode; @@ -25,7 +28,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; @@ -113,22 +115,59 @@ impl VTable for ByteBool { )) } + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + ByteBoolSlots::NAMES[idx].to_string() + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + crate::rules::RULES.evaluate(array, parent, child_idx) + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + // convert truthy values to set/unset bits + let boolean_buffer = BitBufferMut::from(array.truthy_bytes()).freeze(); + let validity = array.validity()?; + Ok(ExecutionResult::done( + BoolArray::new(boolean_buffer, validity).into_array(), + )) + } +} + +impl ArrayPlugin for ByteBool { + fn id(&self) -> ArrayId { + VTable::id(self) + } + fn serialize( - _array: ArrayView<'_, Self>, + &self, + array: &ArrayRef, _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) + ) -> VortexResult> { + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + vec![], + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + if !metadata.is_empty() { vortex_bail!( "ByteBoolArray expects empty metadata, got {} bytes", @@ -151,34 +190,14 @@ impl VTable for ByteBool { let data = ByteBoolData::new(buffer); let slots = ByteBoolData::make_slots(&validity, len); - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots, - )) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - ByteBoolSlots::NAMES[idx].to_string() - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - crate::rules::RULES.evaluate(array, parent, child_idx) - } - - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { - // convert truthy values to set/unset bits - let boolean_buffer = BitBufferMut::from(array.truthy_bytes()).freeze(); - let validity = array.validity()?; - Ok(ExecutionResult::done( - BoolArray::new(boolean_buffer, validity).into_array(), - )) + ))? + .into_array()) } } diff --git a/encodings/datetime-parts/src/array.rs b/encodings/datetime-parts/src/array.rs index 4416327acaf..f18dc97fb86 100644 --- a/encodings/datetime-parts/src/array.rs +++ b/encodings/datetime-parts/src/array.rs @@ -9,11 +9,14 @@ use std::hash::Hasher; use prost::Message; use vortex_array::AnyCanonical; use vortex_array::Array; +use vortex_array::ArrayDeserialization; use vortex_array::ArrayEq; use vortex_array::ArrayHash; use vortex_array::ArrayId; use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; use vortex_array::ArrayView; use vortex_array::EqMode; use vortex_array::ExecutionCtx; @@ -27,7 +30,6 @@ 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; @@ -130,29 +132,71 @@ impl VTable for DateTimeParts { vortex_array::vtable::with_empty_buffers(self, array, buffers) } - fn serialize( + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + DateTimePartsSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let array = require_child!(array, array.days(), DateTimePartsSlots::DAYS => Primitive); + let array = + require_child!(array, array.seconds(), DateTimePartsSlots::SECONDS => AnyCanonical); + let array = require_child!(array, array.subseconds(), DateTimePartsSlots::SUBSECONDS => AnyCanonical); + + let dtype = array.dtype().clone(); + let parts = array.into_parts(); + + Ok(ExecutionResult::done( + decode_to_temporal(parts, &dtype, ctx)?.into_array(), + )) + } + + fn reduce_parent( array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } +} + +impl ArrayPlugin for DateTimeParts { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, _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(), - )) + ) -> VortexResult> { + let view = array.as_::(); + 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( + ArrayPlugin::id(self), + array, + metadata, + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + children, + .. + } = parts; + let metadata = DateTimePartsMetadata::decode(metadata)?; if children.len() != 3 { vortex_bail!( @@ -179,39 +223,14 @@ impl VTable for DateTimeParts { let slots = smallvec![Some(days), Some(seconds), Some(subseconds)]; let data = DateTimePartsData {}; - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots, - )) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - DateTimePartsSlots::NAMES[idx].to_string() - } - - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - let array = require_child!(array, array.days(), DateTimePartsSlots::DAYS => Primitive); - let array = - require_child!(array, array.seconds(), DateTimePartsSlots::SECONDS => AnyCanonical); - let array = require_child!(array, array.subseconds(), DateTimePartsSlots::SUBSECONDS => AnyCanonical); - - let dtype = array.dtype().clone(); - let parts = array.into_parts(); - - Ok(ExecutionResult::done( - decode_to_temporal(parts, &dtype, ctx)?.into_array(), - )) - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - PARENT_RULES.evaluate(array, parent, child_idx) + ))? + .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 796789d5799..b9f148afc5a 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs @@ -26,7 +26,6 @@ use vortex_array::dtype::DecimalDType; use vortex_array::require_child; use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; -use vortex_array::serde::ArrayChildren; use vortex_array::vtable::OperationsVTable; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityChild; @@ -37,7 +36,6 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use super::MAX_LOWER_PARTS; use super::assemble::assemble_decimal; @@ -254,25 +252,6 @@ impl VTable for DecimalByteParts { vortex_array::vtable::with_empty_buffers(self, array, buffers) } - 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!("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/plugin/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/mod.rs index c91c1e3c21a..a4718dbc6c4 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 @@ -45,8 +45,8 @@ pub fn decimal_byte_parts_v2_id() -> ArrayId { /// 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. +/// Register this plugin, or call [`crate::initialize`], to enable both formats. +/// [`DecimalByteParts`] does not implement [`ArrayPlugin`], so it cannot be registered on its own. #[derive(Clone, Debug)] pub struct DecimalBytePartsPlugin; 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 f74c3bf1c55..88d5eb3dc94 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 @@ -248,38 +248,6 @@ fn serialization_requires_v2_permission() -> VortexResult<()> { Ok(()) } -#[test] -fn bare_vtable_refuses_serde() -> VortexResult<()> { - let session = VortexSession::empty(); - 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.is_none(), - "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() } diff --git a/encodings/elias-fano/src/unary/tests.rs b/encodings/elias-fano/src/unary/tests.rs index d4a7f16ab0a..6eb37d669ad 100644 --- a/encodings/elias-fano/src/unary/tests.rs +++ b/encodings/elias-fano/src/unary/tests.rs @@ -278,10 +278,8 @@ fn test_try_new_rejects() -> VortexResult<()> { Ok(()) } -#[rstest] -#[case::vtable(false)] -#[case::plugin(true)] -fn test_deserialize_rejects(#[case] plugin: bool) -> VortexResult<()> { +#[test] +fn test_deserialize_rejects() -> VortexResult<()> { let dtype = DType::Primitive(PType::U64, NonNullable); let metadata = |num_ones: u64| { UnaryMetadata { @@ -293,26 +291,12 @@ fn test_deserialize_rejects(#[case] plugin: bool) -> VortexResult<()> { // `[1, 0]` takes the bits `10 0`. let bits = BufferHandle::new_host(BitBuffer::from(vec![true, false, false]).inner().clone()); let no_children: [ArrayRef; 0] = []; - // Through the vtable alone, and through the array plugin, the path readers run. + // Through the array plugin, the path readers run. let deserialize = |metadata: &[u8], buffers: &[BufferHandle], children: &[ArrayRef]| { - if plugin { - ArrayPlugin::deserialize( - &Unary, - ArrayDeserialization::new( - VTable::id(&Unary), - &dtype, - 2, - metadata, - buffers, - &children, - ), - &SESSION, - ) - } else { - VTable::deserialize(&Unary, &dtype, 2, metadata, buffers, &children, &SESSION) - .and_then(Array::try_from_parts) - .map(IntoArray::into_array) - } + Unary.deserialize( + ArrayDeserialization::new(VTable::id(&Unary), &dtype, 2, metadata, buffers, &children), + &SESSION, + ) }; let one = slice::from_ref(&bits); diff --git a/encodings/elias-fano/src/unary/vtable/mod.rs b/encodings/elias-fano/src/unary/vtable/mod.rs index cfbc28fdb13..ef8ee9d30c2 100644 --- a/encodings/elias-fano/src/unary/vtable/mod.rs +++ b/encodings/elias-fano/src/unary/vtable/mod.rs @@ -7,11 +7,14 @@ use std::hash::Hasher; use num_traits::AsPrimitive; use prost::Message; use vortex_array::Array; +use vortex_array::ArrayDeserialization; use vortex_array::ArrayEq; use vortex_array::ArrayHash; 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::ArrayView; use vortex_array::EqMode; @@ -24,7 +27,6 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::PType; use vortex_array::match_each_unsigned_integer_ptype; -use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_buffer::BitBuffer; @@ -134,28 +136,54 @@ impl VTable for Unary { )) } + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + // Validity, once added, can take an optional slot 0, as in Bool. See `Unary`. + vortex_panic!("UnaryArray has no slots, got index {idx}") + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(unary_decode(&array).into_array())) + } +} + +impl ArrayPlugin for Unary { + fn id(&self) -> ArrayId { + VTable::id(self) + } + fn serialize( - array: ArrayView<'_, Self>, + &self, + array: &ArrayRef, _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - UnaryMetadata { - num_ones: num_ones_of(array)?, - offset: u32::try_from(array.meta.offset())?, - } - .encode_to_vec(), - )) + ) -> VortexResult> { + let view = array.as_::(); + let metadata = UnaryMetadata { + num_ones: num_ones_of(view)?, + offset: u32::try_from(view.meta.offset())?, + } + .encode_to_vec(); + + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + metadata, + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + vortex_ensure_eq!(buffers.len(), 1); vortex_ensure_eq!(children.len(), 0); let metadata = UnaryMetadata::decode(metadata)?; @@ -165,22 +193,14 @@ impl VTable for Unary { usize::try_from(metadata.offset)?, len, )?; - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, ArraySlots::new(), - )) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - // Validity, once added, can take an optional slot 0, as in Bool. See `Unary`. - vortex_panic!("UnaryArray has no slots, got index {idx}") - } - - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { - Ok(ExecutionResult::done(unary_decode(&array).into_array())) + ))? + .into_array()) } } diff --git a/encodings/fastlanes/src/bitpacking/plugin/mod.rs b/encodings/fastlanes/src/bitpacking/plugin/mod.rs index aec1375c341..dc470fba704 100644 --- a/encodings/fastlanes/src/bitpacking/plugin/mod.rs +++ b/encodings/fastlanes/src/bitpacking/plugin/mod.rs @@ -54,8 +54,8 @@ pub fn bitpacked_v2_id() -> ArrayId { /// its metadata. Arrays whose blocks have their own bit widths serialize as /// `fastlanes.bitpacked.v2`, which stores the block offsets as a child. /// -/// Register this plugin, or call [`crate::initialize`], to enable serde. Direct registration of -/// [`BitPacked`] does not support serde. +/// Register this plugin, or call [`crate::initialize`], to enable serde. [`BitPacked`] does not +/// implement [`ArrayPlugin`], so it cannot be registered on its own. #[derive(Clone, Debug)] pub struct BitPackedPlugin; diff --git a/encodings/fastlanes/src/bitpacking/plugin/tests.rs b/encodings/fastlanes/src/bitpacking/plugin/tests.rs index 80a2309bd54..d28fd31bb20 100644 --- a/encodings/fastlanes/src/bitpacking/plugin/tests.rs +++ b/encodings/fastlanes/src/bitpacking/plugin/tests.rs @@ -242,16 +242,6 @@ fn serde_roundtrip( Ok(()) } -#[test] -fn vtable_serde_requires_plugin() -> VortexResult<()> { - let values = PrimitiveArray::from_iter([1u8, 2, 3]).into_array(); - let array = BitPackedData::encode(&values, 2, &mut SESSION.create_execution_ctx())?; - let session = VortexSession::empty(); - session.arrays().register(BitPacked); - assert!(session.array_serialize(&array.into_array())?.is_none()); - Ok(()) -} - /// Values whose 1024-value blocks need 2 to 6 bits. fn drifting() -> PrimitiveArray { PrimitiveArray::from_iter((0..5000u32).map(|i| i % (4 << (i / 1024)))) diff --git a/encodings/fastlanes/src/bitpacking/vtable/mod.rs b/encodings/fastlanes/src/bitpacking/vtable/mod.rs index e4ff6b061e9..02d9ec47b81 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/mod.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/mod.rs @@ -27,7 +27,6 @@ use vortex_array::patches::PatchesData; use vortex_array::require_opt_child; 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; @@ -166,25 +165,6 @@ impl VTable for BitPacked { )) } - 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!("BitPacked deserialization requires BitPackedPlugin") - } - fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, diff --git a/encodings/fastlanes/src/delta/vtable/mod.rs b/encodings/fastlanes/src/delta/vtable/mod.rs index 607c7153e1c..f891fc51e22 100644 --- a/encodings/fastlanes/src/delta/vtable/mod.rs +++ b/encodings/fastlanes/src/delta/vtable/mod.rs @@ -6,11 +6,14 @@ use std::hash::Hasher; use prost::Message; use vortex_array::Array; +use vortex_array::ArrayDeserialization; use vortex_array::ArrayEq; use vortex_array::ArrayHash; use vortex_array::ArrayId; use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; use vortex_array::ArrayView; use vortex_array::EqMode; use vortex_array::ExecutionCtx; @@ -22,7 +25,6 @@ use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::PType; use vortex_array::require_child; -use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::validity_to_child; @@ -134,28 +136,53 @@ impl VTable for Delta { DeltaSlots::NAMES[idx].to_string() } + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let array = require_child!(array, array.bases(), DeltaSlots::BASES => Primitive); + let array = require_child!(array, array.deltas(), DeltaSlots::DELTAS => Primitive); + Ok(ExecutionResult::done( + delta_decompress(&array, ctx)?.into_array(), + )) + } +} + +impl ArrayPlugin for Delta { + fn id(&self) -> ArrayId { + VTable::id(self) + } + fn serialize( - array: ArrayView<'_, Self>, + &self, + array: &ArrayRef, _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - DeltaMetadata { - deltas_len: array.deltas().len() as u64, - offset: array.offset() as u32, - } - .encode_to_vec(), - )) + ) -> VortexResult> { + let view = array.as_::(); + let metadata = DeltaMetadata { + deltas_len: view.deltas().len() as u64, + offset: view.offset() as u32, + } + .encode_to_vec(); + + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + metadata, + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + vortex_ensure!( buffers.is_empty(), "DeltaArray expects 0 buffers, got {}", @@ -193,21 +220,14 @@ impl VTable for Delta { validity_child, } .into_slots(); - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots, - )) - } - - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - let array = require_child!(array, array.bases(), DeltaSlots::BASES => Primitive); - let array = require_child!(array, array.deltas(), DeltaSlots::DELTAS => Primitive); - Ok(ExecutionResult::done( - delta_decompress(&array, ctx)?.into_array(), - )) + ))? + .into_array()) } } diff --git a/encodings/fastlanes/src/for_/plugin/mod.rs b/encodings/fastlanes/src/for_/plugin/mod.rs index 0efbff59e83..9229ad6c198 100644 --- a/encodings/fastlanes/src/for_/plugin/mod.rs +++ b/encodings/fastlanes/src/for_/plugin/mod.rs @@ -42,8 +42,8 @@ pub fn for_v2_id() -> ArrayId { /// its metadata. Arrays whose chunks have different references serialize as /// `fastlanes.for.v2`, which stores them as a child. /// -/// Register this plugin, or call [`crate::initialize`], to enable serde. Direct registration of -/// [`FoR`] does not support serde. +/// Register this plugin, or call [`crate::initialize`], to enable serde. [`FoR`] does not +/// implement [`ArrayPlugin`], so it cannot be registered on its own. #[derive(Clone, Debug)] pub struct FoRPlugin; diff --git a/encodings/fastlanes/src/for_/plugin/tests.rs b/encodings/fastlanes/src/for_/plugin/tests.rs index ff59e094595..a1eaf2d5050 100644 --- a/encodings/fastlanes/src/for_/plugin/tests.rs +++ b/encodings/fastlanes/src/for_/plugin/tests.rs @@ -130,15 +130,3 @@ fn v2_rejects_malformed_parts() -> VortexResult<()> { ); Ok(()) } - -#[test] -fn vtable_serde_requires_plugin() -> VortexResult<()> { - let array = FoR::encode( - PrimitiveArray::from_iter([1u8, 2, 3]), - &mut SESSION.create_execution_ctx(), - )?; - let session = VortexSession::empty(); - session.arrays().register(FoR); - assert!(session.array_serialize(&array.into_array())?.is_none()); - Ok(()) -} diff --git a/encodings/fastlanes/src/for_/vtable/mod.rs b/encodings/fastlanes/src/for_/vtable/mod.rs index 0fd9776e806..e919e209836 100644 --- a/encodings/fastlanes/src/for_/vtable/mod.rs +++ b/encodings/fastlanes/src/for_/vtable/mod.rs @@ -23,7 +23,6 @@ use vortex_array::dtype::DType; use vortex_array::require_child; use vortex_array::require_opt_child; use vortex_array::scalar::Scalar; -use vortex_array::serde::ArrayChildren; use vortex_array::smallvec::smallvec; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTableFromChild; @@ -128,25 +127,6 @@ impl VTable for FoR { FoRSlots::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!("FoR deserialization requires FoRPlugin") - } - fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, diff --git a/encodings/fastlanes/src/rle/vtable/mod.rs b/encodings/fastlanes/src/rle/vtable/mod.rs index a06c615807e..14a059868cf 100644 --- a/encodings/fastlanes/src/rle/vtable/mod.rs +++ b/encodings/fastlanes/src/rle/vtable/mod.rs @@ -6,11 +6,14 @@ use std::hash::Hasher; use prost::Message; use vortex_array::Array; +use vortex_array::ArrayDeserialization; use vortex_array::ArrayEq; use vortex_array::ArrayHash; use vortex_array::ArrayId; use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; use vortex_array::ArrayView; use vortex_array::EqMode; use vortex_array::ExecutionCtx; @@ -21,7 +24,6 @@ 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; @@ -155,33 +157,55 @@ impl VTable for RLE { RLESlots::NAMES[idx].to_string() } + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done( + rle_decompress(&array, ctx)?.into_array(), + )) + } +} + +impl ArrayPlugin for RLE { + fn id(&self) -> ArrayId { + VTable::id(self) + } + fn serialize( - array: ArrayView<'_, Self>, + &self, + array: &ArrayRef, _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(), - )) + ) -> VortexResult> { + let view = array.as_::(); + 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( + ArrayPlugin::id(self), + array, + metadata, + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + let metadata = RLEMetadata::decode(metadata)?; let indices_ptype = metadata.get_indices_ptype()?; let values_idx_offsets_ptype = metadata.get_values_idx_offsets_ptype()?; @@ -215,19 +239,14 @@ impl VTable for RLE { } .into_slots(); let data = RLEData::try_new(metadata.offset as usize)?; - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots, - )) - } - - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - Ok(ExecutionResult::done( - rle_decompress(&array, ctx)?.into_array(), - )) + ))? + .into_array()) } } @@ -349,6 +368,8 @@ fn validate_parts( mod tests { use prost::Message; use rstest::rstest; + use vortex_array::ArrayDeserialization; + use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -397,11 +418,14 @@ mod tests { let metadata = with_ptype(11).encode_to_vec(); let children: Vec = Vec::new(); let result = RLE.deserialize( - &DType::Primitive(PType::U32, Nullability::NonNullable), - 0, - &metadata, - &[], - &children, + ArrayDeserialization::new( + VTable::id(&RLE), + &DType::Primitive(PType::U32, Nullability::NonNullable), + 0, + &metadata, + &[], + &children, + ), &SESSION, ); assert!(result.is_err_and(|err| err.to_string().contains(expected))); diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index ae8d909b052..2d37cb9c442 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -15,16 +15,20 @@ use fsst::Symbol; use num_traits::AsPrimitive; use prost::Message as _; use vortex_array::Array; +use vortex_array::ArrayDeserialization; use vortex_array::ArrayEq; use vortex_array::ArrayHash; 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::ArrayView; use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; +use vortex_array::IntoArray; use vortex_array::TypedArrayRef; use vortex_array::VortexSessionExecute; use vortex_array::array_slots; @@ -187,60 +191,116 @@ impl VTable for FSST { )) } - fn serialize( + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + FSSTSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let array = require_child!( + array, + array.uncompressed_lengths(), + FSSTSlots::UNCOMPRESSED_LENGTHS => Primitive + ); + let array = require_child!( + array, + array.codes_offsets(), + FSSTSlots::CODES_OFFSETS => Primitive + ); + require_validity!(array, FSSTSlots::CODES_VALIDITY); + canonicalize_fsst(array.as_view(), ctx).map(ExecutionResult::done) + } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if let Some(result) = + match_each_varbin_builder!(builder, |builder| append_to_varbin(array, builder, ctx)) + { + return result; + } + + // The two arms here are every builder a `Utf8`/`Binary` dtype has: all four + // `VarBinBuilder` widths above, and `VarBinViewBuilder` below. There is deliberately no + // canonicalize-then-append fallback — it would decode to a `VarBinView` only for + // `VarBinView::append_to_builder` to reject the same remainder. + let Some(builder) = builder.as_any_mut().downcast_mut::() else { + vortex_bail!("append_to_builder for FSST requires a variable-binary builder") + }; + + // Decompress the whole block of data into a new buffer, which the builder adopts as a + // data buffer with views built over it in place. + let validity = array + .array() + .validity()? + .execute_mask(array.array().len(), ctx)?; + let (uncompressed_bytes, uncompressed_lens) = fsst_decode_bytes(array, ctx)?; + match_each_integer_ptype!(uncompressed_lens.ptype(), |P| { + builder.append_buffer_with_lengths( + uncompressed_bytes.freeze(), + uncompressed_lens.as_slice::

(), + &validity, + ) + }); + Ok(()) + } + + fn reduce_parent( array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + RULES.evaluate(array, parent, child_idx) + } +} + +impl ArrayPlugin for FSST { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, _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(), - )) + ) -> VortexResult> { + let view = array.as_::(); + 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( + ArrayPlugin::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. fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + 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( + let legacy = Self::deserialize_legacy( self, dtype, len, @@ -249,7 +309,8 @@ impl VTable for FSST { &symbol_lengths, children, &mut ctx, - ); + )?; + return Ok(Array::::try_from_parts(legacy)?.into_array()); } if buffers.len() == 3 { @@ -299,13 +360,14 @@ impl VTable for FSST { } .into_slots(); let data = FSSTData::try_new(symbols, symbol_lengths, codes_bytes, len)?; - return Ok(ArrayParts::new( + return Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots, - )); + ))? + .into_array()); } vortex_bail!( @@ -313,69 +375,6 @@ impl VTable for FSST { buffers.len() ); } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - FSSTSlots::NAMES[idx].to_string() - } - - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - let array = require_child!( - array, - array.uncompressed_lengths(), - FSSTSlots::UNCOMPRESSED_LENGTHS => Primitive - ); - let array = require_child!( - array, - array.codes_offsets(), - FSSTSlots::CODES_OFFSETS => Primitive - ); - require_validity!(array, FSSTSlots::CODES_VALIDITY); - canonicalize_fsst(array.as_view(), ctx).map(ExecutionResult::done) - } - - fn append_to_builder( - array: ArrayView<'_, Self>, - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - if let Some(result) = - match_each_varbin_builder!(builder, |builder| append_to_varbin(array, builder, ctx)) - { - return result; - } - - // The two arms here are every builder a `Utf8`/`Binary` dtype has: all four - // `VarBinBuilder` widths above, and `VarBinViewBuilder` below. There is deliberately no - // canonicalize-then-append fallback — it would decode to a `VarBinView` only for - // `VarBinView::append_to_builder` to reject the same remainder. - let Some(builder) = builder.as_any_mut().downcast_mut::() else { - vortex_bail!("append_to_builder for FSST requires a variable-binary builder") - }; - - // Decompress the whole block of data into a new buffer, which the builder adopts as a - // data buffer with views built over it in place. - let validity = array - .array() - .validity()? - .execute_mask(array.array().len(), ctx)?; - let (uncompressed_bytes, uncompressed_lens) = fsst_decode_bytes(array, ctx)?; - match_each_integer_ptype!(uncompressed_lens.ptype(), |P| { - builder.append_buffer_with_lengths( - uncompressed_bytes.freeze(), - uncompressed_lens.as_slice::

(), - &validity, - ) - }); - Ok(()) - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - RULES.evaluate(array, parent, child_idx) - } } /// Decompresses the code stream straight into `builder`'s byte storage. diff --git a/encodings/onpair/src/array.rs b/encodings/onpair/src/array.rs index 8b45fc13438..b813ca31137 100644 --- a/encodings/onpair/src/array.rs +++ b/encodings/onpair/src/array.rs @@ -16,15 +16,19 @@ use onpair::Dictionary; use onpair::DictionaryStorage; use prost::Message as _; use vortex_array::Array; +use vortex_array::ArrayDeserialization; use vortex_array::ArrayEq; use vortex_array::ArrayHash; use vortex_array::ArrayId; use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; use vortex_array::ArrayView; use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; +use vortex_array::IntoArray; use vortex_array::array_slots; use vortex_array::arrays::Primitive; use vortex_array::buffer::BufferHandle; @@ -39,7 +43,6 @@ use vortex_array::match_each_integer_ptype; use vortex_array::match_each_varbin_builder; use vortex_array::require_child; use vortex_array::require_validity; -use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; @@ -510,104 +513,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!("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!("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, - slots, - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { OnPairSlots::NAMES[idx].to_string() } @@ -684,6 +589,122 @@ impl VTable for OnPair { } } +impl ArrayPlugin for OnPair { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_::(); + 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( + ArrayPlugin::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + + if buffers.len() != 1 { + vortex_bail!("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!("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, + slots, + ))? + .into_array()) + } +} + /// Decodes the code stream straight into `builder`'s byte storage. /// /// The offsets are the running sum of the uncompressed lengths the array already stores, so the diff --git a/encodings/parquet-variant/src/vtable.rs b/encodings/parquet-variant/src/vtable.rs index b7b8105a2ab..a352b9baadd 100644 --- a/encodings/parquet-variant/src/vtable.rs +++ b/encodings/parquet-variant/src/vtable.rs @@ -3,9 +3,12 @@ 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::ArrayView; use vortex_array::EmptyArrayData; use vortex_array::ExecutionCtx; @@ -16,7 +19,6 @@ 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; @@ -176,33 +178,64 @@ impl VTable for ParquetVariant { ParquetVariantSlots::NAMES[idx].to_string() } + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let shredded = array + .typed_value() + .cloned() + .map(|typed_value| { + logical_shredded_from_parquet_typed_value(array.metadata(), typed_value, ctx) + }) + .transpose()?; + let core_storage = core_storage_without_typed_value(&array)?; + Ok(ExecutionResult::done( + VariantArray::try_new(core_storage, shredded)?.into_array(), + )) + } +} + +impl ArrayPlugin for ParquetVariant { + fn id(&self) -> ArrayId { + VTable::id(self) + } + fn serialize( - array: ArrayView<'_, Self>, + &self, + array: &ArrayRef, _session: &VortexSession, - ) -> VortexResult>> { - let typed_value_dtype = array + ) -> VortexResult> { + let view = array.as_::(); + let typed_value_dtype = view .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(), - )) + 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( + ArrayPlugin::id(self), + array, + metadata, + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + vortex_ensure!( buffers.is_empty(), "ParquetVariantArray expects 0 buffers, got {}", @@ -265,27 +298,14 @@ impl VTable for ParquetVariant { typed_value, } .into_slots(); - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, EmptyArrayData, slots, - )) - } - - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - let shredded = array - .typed_value() - .cloned() - .map(|typed_value| { - logical_shredded_from_parquet_typed_value(array.metadata(), typed_value, ctx) - }) - .transpose()?; - let core_storage = core_storage_without_typed_value(&array)?; - Ok(ExecutionResult::done( - VariantArray::try_new(core_storage, shredded)?.into_array(), - )) + ))? + .into_array()) } } diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index f7cc8c0a261..9c1caf2714e 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -19,11 +19,14 @@ use pco::wrapped::FileCompressor; use pco::wrapped::FileDecompressor; use prost::Message; use vortex_array::Array; +use vortex_array::ArrayDeserialization; use vortex_array::ArrayEq; use vortex_array::ArrayHash; use vortex_array::ArrayId; use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; use vortex_array::ArrayView; use vortex_array::EqMode; use vortex_array::ExecutionCtx; @@ -39,7 +42,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; @@ -194,26 +196,67 @@ impl VTable for Pco { )) } - fn serialize( + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + PcoSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let unsliced_validity = array.unsliced_validity(); + Ok(ExecutionResult::done( + array + .data() + .decompress(&unsliced_validity, ctx)? + .into_array(), + )) + } + + fn reduce_parent( array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + crate::rules::RULES.evaluate(array, parent, child_idx) + } +} + +impl ArrayPlugin for Pco { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, _session: &VortexSession, - ) -> VortexResult>> { + ) -> VortexResult> { + let view = array.as_::(); vortex_ensure!( - array.slice_start() == 0 && array.slice_stop() == array.unsliced_n_rows(), + view.slice_start() == 0 && view.slice_stop() == view.unsliced_n_rows(), "cannot serialize sliced PCO arrays; materialize the slice first" ); - Ok(Some(array.metadata.clone().encode_to_vec())) + let metadata = view.metadata.clone().encode_to_vec(); + + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + metadata, + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + let metadata = PcoMetadata::decode(metadata)?; let validity = if children.is_empty() { Validity::from(dtype.nullability()) @@ -249,35 +292,14 @@ impl VTable for Pco { // publishing the array. let data = unsafe { PcoData::new_unchecked(chunk_metas, pages, dtype.as_ptype(), metadata, len) }; - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots, - )) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - PcoSlots::NAMES[idx].to_string() - } - - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - let unsliced_validity = array.unsliced_validity(); - Ok(ExecutionResult::done( - array - .data() - .decompress(&unsliced_validity, ctx)? - .into_array(), - )) - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - crate::rules::RULES.evaluate(array, parent, child_idx) + ))? + .into_array()) } } diff --git a/encodings/sequence/src/array.rs b/encodings/sequence/src/array.rs index 22438e35494..31a3b49cfd4 100644 --- a/encodings/sequence/src/array.rs +++ b/encodings/sequence/src/array.rs @@ -10,16 +10,20 @@ use num_traits::AsPrimitive; use prost::Message; use smallvec::smallvec; use vortex_array::Array; +use vortex_array::ArrayDeserialization; use vortex_array::ArrayEq; use vortex_array::ArrayHash; 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::ArrayView; use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; +use vortex_array::IntoArray; use vortex_array::ProbeState; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; @@ -35,7 +39,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; @@ -337,27 +340,61 @@ impl VTable for Sequence { vortex_array::vtable::with_empty_buffers(self, array, buffers) } - fn serialize( + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + SLOT_NAMES[idx].to_string() + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + sequence_decompress(&array).map(ExecutionResult::done) + } + + fn reduce_parent( array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + RULES.evaluate(array, parent, child_idx) + } +} + +impl ArrayPlugin for Sequence { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, _session: &VortexSession, - ) -> VortexResult>> { + ) -> VortexResult> { + let view = array.as_::(); let metadata = SequenceMetadata { - base: Some((&array.base()).into()), - multiplier: Some((&array.multiplier()).into()), - }; + base: Some((&view.base()).into()), + multiplier: Some((&view.multiplier()).into()), + } + .encode_to_vec(); - Ok(Some(metadata.encode_to_vec())) + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + metadata, + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + vortex_ensure!( buffers.is_empty(), "SequenceArray expects 0 buffers, got {}", @@ -408,29 +445,14 @@ impl VTable for Sequence { let data = SequenceData::try_new(base, multiplier, *output_ptype, dtype.nullability(), len)?; - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, ArraySlots::new(), - )) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() - } - - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { - sequence_decompress(&array).map(ExecutionResult::done) - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - RULES.evaluate(array, parent, child_idx) + ))? + .into_array()) } } diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index ecf00243f15..61014600e99 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -10,11 +10,14 @@ use std::hash::Hasher; use prost::Message as _; use vortex_array::AnyCanonical; use vortex_array::Array; +use vortex_array::ArrayDeserialization; use vortex_array::ArrayEq; use vortex_array::ArrayHash; 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::ArrayView; use vortex_array::Canonical; @@ -44,7 +47,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; @@ -91,32 +93,32 @@ pub fn initialize(session: &VortexSession) { let aggregate_fns = session.aggregate_fns(); aggregate_fns.register_aggregate_kernel( - Sparse.id(), + VTable::id(&Sparse), Some(IsConstant.id()), &compute::is_constant::SparseIsConstantKernel, ); aggregate_fns.register_aggregate_kernel( - Sparse.id(), + VTable::id(&Sparse), Some(Sum.id()), &compute::sum::SparseSumKernel, ); aggregate_fns.register_aggregate_kernel( - Sparse.id(), + VTable::id(&Sparse), Some(SumV2.id()), &compute::sum::SparseSumKernel, ); aggregate_fns.register_aggregate_kernel( - Sparse.id(), + VTable::id(&Sparse), Some(MinMax.id()), &compute::min_max::SparseMinMaxKernel, ); aggregate_fns.register_aggregate_kernel( - Sparse.id(), + VTable::id(&Sparse), Some(NullCount.id()), &compute::null_count::SparseNullCountKernel, ); aggregate_fns.register_aggregate_kernel( - Sparse.id(), + VTable::id(&Sparse), Some(NanCount.id()), &compute::nan_count::SparseNanCountKernel, ); @@ -250,66 +252,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"); - - 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, - &mut session.create_execution_ctx(), - )?; - let slots = SparseData::make_slots(&patches); - let data = SparseData::from_patches(&patches, fill_value)?; - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { SparseSlots::NAMES[idx].to_string() } @@ -400,6 +342,84 @@ impl VTable for Sparse { } } +impl ArrayPlugin for Sparse { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_::(); + let patches = view.patches().to_metadata(view.len(), view.dtype())?; + // Note that we DO NOT serialize the fill value since that is stored in the buffers. + let metadata = SparseMetadata { patches }.encode_to_vec(); + + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + + 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"); + + 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, + &mut session.create_execution_ctx(), + )?; + 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, + slots, + ))? + .into_array()) + } +} + const PATCH_SLOTS: PatchSlotIndices = PatchSlotIndices { indices: SparseSlots::PATCH_INDICES, values: SparseSlots::PATCH_VALUES, diff --git a/encodings/zigzag/src/array.rs b/encodings/zigzag/src/array.rs index 101da38b521..324ba31382c 100644 --- a/encodings/zigzag/src/array.rs +++ b/encodings/zigzag/src/array.rs @@ -6,11 +6,14 @@ use std::fmt::Formatter; use std::hash::Hasher; use vortex_array::Array; +use vortex_array::ArrayDeserialization; use vortex_array::ArrayEq; use vortex_array::ArrayHash; use vortex_array::ArrayId; use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; use vortex_array::ArrayView; use vortex_array::EqMode; use vortex_array::ExecutionCtx; @@ -24,7 +27,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; @@ -92,22 +94,55 @@ impl VTable for ZigZag { vortex_array::vtable::with_empty_buffers(self, array, buffers) } + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + ZigZagSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done( + zigzag_decode(array.encoded().clone().execute(ctx)?).into_array(), + )) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + RULES.evaluate(array, parent, child_idx) + } +} + +impl ArrayPlugin for ZigZag { + fn id(&self) -> ArrayId { + VTable::id(self) + } + fn serialize( - _array: ArrayView<'_, Self>, + &self, + array: &ArrayRef, _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) + ) -> VortexResult> { + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + vec![], + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + children, + .. + } = parts; + if !metadata.is_empty() { vortex_bail!( "ZigZagArray expects empty metadata, got {} bytes", @@ -124,31 +159,14 @@ impl VTable for ZigZag { let encoded = children.get(0, &encoded_type, len)?; let slots = smallvec![Some(encoded.clone())]; let data = ZigZagData::try_new(encoded.dtype())?; - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots, - )) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - ZigZagSlots::NAMES[idx].to_string() - } - - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - Ok(ExecutionResult::done( - zigzag_decode(array.encoded().clone().execute(ctx)?).into_array(), - )) - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - RULES.evaluate(array, parent, child_idx) + ))? + .into_array()) } } diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index 78034a5fd83..b05dc89c381 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -14,11 +14,14 @@ use itertools::Itertools as _; use num_traits::AsPrimitive; use prost::Message as _; use vortex_array::Array; +use vortex_array::ArrayDeserialization; use vortex_array::ArrayEq; use vortex_array::ArrayHash; use vortex_array::ArrayId; use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; use vortex_array::ArrayView; use vortex_array::Canonical; use vortex_array::EqMode; @@ -39,7 +42,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; @@ -217,22 +219,86 @@ impl VTable for Zstd { )) } - fn serialize( + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + ZstdSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let unsliced_validity = child_to_validity( + array.as_ref().slots()[ZstdSlots::VALIDITY].as_ref(), + array.dtype().nullability(), + ); + array + .data() + .decompress(array.dtype(), &unsliced_validity, ctx)? + .execute::(ctx) + .map(ExecutionResult::done) + } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if let Some(result) = + match_each_varbin_builder!(builder, |builder| append_to_varbin(array, builder, ctx)) + { + return result; + } + if let Some(builder) = builder.as_any_mut().downcast_mut::() { + return append_to_varbinview(array, builder, ctx); + } + array + .array() + .clone() + .execute::(ctx)? + .into_array() + .append_to_builder(builder, ctx) + } + + fn reduce_parent( array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + crate::rules::RULES.evaluate(array, parent, child_idx) + } +} + +impl ArrayPlugin for Zstd { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(array.metadata.clone().encode_to_vec())) + ) -> VortexResult> { + let view = array.as_::(); + let metadata = view.metadata.clone().encode_to_vec(); + + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + metadata, + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + let metadata = ZstdMetadata::decode(metadata)?; let validity = if children.is_empty() { Validity::from(dtype.nullability()) @@ -265,58 +331,14 @@ impl VTable for Zstd { let slots = smallvec![validity_to_child(&validity, len)]; let data = ZstdData::new(dictionary_buffer, compressed_buffers, metadata, len); - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots, - )) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - ZstdSlots::NAMES[idx].to_string() - } - - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - let unsliced_validity = child_to_validity( - array.as_ref().slots()[ZstdSlots::VALIDITY].as_ref(), - array.dtype().nullability(), - ); - array - .data() - .decompress(array.dtype(), &unsliced_validity, ctx)? - .execute::(ctx) - .map(ExecutionResult::done) - } - - fn append_to_builder( - array: ArrayView<'_, Self>, - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - if let Some(result) = - match_each_varbin_builder!(builder, |builder| append_to_varbin(array, builder, ctx)) - { - return result; - } - if let Some(builder) = builder.as_any_mut().downcast_mut::() { - return append_to_varbinview(array, builder, ctx); - } - array - .array() - .clone() - .execute::(ctx)? - .into_array() - .append_to_builder(builder, ctx) - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - crate::rules::RULES.evaluate(array, parent, child_idx) + ))? + .into_array()) } } diff --git a/encodings/zstd/src/zstd_buffers.rs b/encodings/zstd/src/zstd_buffers.rs index 7619ce9ec6e..07bdfb11308 100644 --- a/encodings/zstd/src/zstd_buffers.rs +++ b/encodings/zstd/src/zstd_buffers.rs @@ -15,16 +15,18 @@ use vortex_array::ArrayEq; use vortex_array::ArrayHash; 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::ArrayView; use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; +use vortex_array::IntoArray; 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; @@ -531,41 +533,68 @@ impl VTable for ZstdBuffers { format!("child_{idx}") } + // with_slots handles child replacement via the slots mechanism + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let session = ctx.session(); + let inner_array = ZstdBuffers::decompress_and_build_inner(&array, session)?; + inner_array + .execute::(ctx) + .map(ExecutionResult::done) + } +} + +impl ArrayPlugin for ZstdBuffers { + fn id(&self) -> ArrayId { + VTable::id(self) + } + fn serialize( - array: ArrayView<'_, Self>, + &self, + array: &ArrayRef, _session: &VortexSession, - ) -> VortexResult>> { - let children: Vec<&ArrayRef> = array.slots().iter().flatten().collect(); + ) -> VortexResult> { + let view = array.as_::(); + 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(); - 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, - inner_child_count: Some(u32::try_from(array.inner_child_count)?), - validity_child: array.validity_slot.map(u32::try_from).transpose()?, - } - .encode_to_vec(), - )) + 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, + inner_child_count: Some(u32::try_from(view.inner_child_count)?), + validity_child: view.validity_slot.map(u32::try_from).transpose()?, + } + .encode_to_vec(); + + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + metadata, + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + let metadata = ZstdBuffersMetadata::decode(metadata)?; let compressed_buffers: Vec = buffers.to_vec(); @@ -616,23 +645,14 @@ impl VTable for ZstdBuffers { None => {} } - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots.into(), - )) - } - - // with_slots handles child replacement via the slots mechanism - - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - let session = ctx.session(); - let inner_array = ZstdBuffers::decompress_and_build_inner(&array, session)?; - inner_array - .execute::(ctx) - .map(ExecutionResult::done) + ))? + .into_array()) } } @@ -860,24 +880,27 @@ mod tests { let input = make_nullable_primitive_array(); let compressed = ZstdBuffers::compress(&input, 3, &session)?; - let mut metadata = ZstdBuffersMetadata::decode( - ZstdBuffers::serialize(compressed.as_view(), &session)? - .vortex_expect("metadata") - .as_slice(), - )?; + let serialization = ZstdBuffers + .serialize(compressed.as_array(), &session)? + .vortex_expect("ZstdBuffers always serializes"); + let mut metadata = ZstdBuffersMetadata::decode(serialization.metadata.as_slice())?; metadata.inner_child_count = None; metadata.validity_child = None; let children: Vec = compressed.slots().iter().flatten().cloned().collect(); - let parts = ZstdBuffers.deserialize( - input.dtype(), - input.len(), - &metadata.encode_to_vec(), - &compressed.data().compressed_buffers, - &children.as_slice(), - &session, - )?; - let decoded = Array::try_from_parts(parts)?; + let decoded = ZstdBuffers + .deserialize( + ArrayDeserialization::new( + VTable::id(&ZstdBuffers), + input.dtype(), + input.len(), + &metadata.encode_to_vec(), + &compressed.data().compressed_buffers, + &children.as_slice(), + ), + &session, + )? + .downcast::(); assert!(decoded.data().validity_slot.is_some()); let mut ctx = session.create_execution_ctx(); diff --git a/vortex-array/src/array/foreign.rs b/vortex-array/src/array/foreign.rs index 57a0d74ba45..418c04ea1ec 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 { )) } - 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()), - 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 b60559f8af9..87497f2b16a 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_eq; 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; @@ -108,6 +104,14 @@ impl<'a> ArrayDeserialization<'a> { /// /// A plugin can retain historical read IDs after it stops writing them. A reader that predates a /// newer ID has no registration for it and reports it as unknown. +/// +/// [`VTable`] does not serialize, so each encoding opts in to serde with its own implementation of +/// this trait. An encoding with one wire form usually implements this trait on its vtable type. An +/// encoding with several wire forms, or one whose wire form belongs to another encoding, uses a +/// separate plugin type. In-memory-only encodings, such as lazy compute arrays, do not implement +/// this trait at all, so they cannot be registered. +/// +/// [`VTable`]: crate::array::VTable pub trait ArrayPlugin: 'static + Send + Sync { /// Returns the ID of the in-memory encoding that this plugin writes. fn id(&self) -> ArrayId; @@ -117,20 +121,31 @@ pub trait ArrayPlugin: 'static + Send + Sync { /// [`ArraySession::register`](crate::session::ArraySession::register) makes this plugin the /// reader of each ID. An ID can stay in this list after the plugin stops writing it, so that /// old files stay readable. - fn read_ids(&self) -> Vec; + /// + /// By default, a plugin reads only its in-memory [`id`](Self::id). + fn read_ids(&self) -> Vec { + vec![self.id()] + } /// Returns the wire IDs that [`serialize`](Self::serialize) may return. /// /// An ID in this list does not have to be in [`read_ids`](Self::read_ids). A plugin can write /// its arrays in the wire form of another encoding. The plugin that owns that form reads it. /// The node must hold the logical values of the array under the dtype of the array. - fn write_ids(&self) -> Vec; + /// + /// By default, a plugin writes only its in-memory [`id`](Self::id). + fn write_ids(&self) -> Vec { + vec![self.id()] + } /// Serialize `array` to its wire representation. /// - /// This function is called only for arrays whose in-memory encoding matches [`id`](Self::id). - /// The returned ID must be declared by [`write_ids`](Self::write_ids). Return - /// `Ok(None)` when the array cannot be serialized. + /// Callers **must** pass an array whose in-memory encoding matches [`id`](Self::id). An + /// implementation may panic otherwise. The returned ID must be declared by + /// [`write_ids`](Self::write_ids). + /// + /// Return `Ok(None)` when the array cannot be serialized, and let the caller decide how to + /// report it. Return an error only when an array that can be serialized fails to serialize. fn serialize( &self, array: &ArrayRef, @@ -139,11 +154,18 @@ pub trait ArrayPlugin: 'static + Send + Sync { /// Deserialize one recognized wire representation into the current in-memory array. /// - /// `serialized_id` identifies the exact representation encountered on disk. The returned - /// array does not necessarily have to use this plugin's in-memory ID; this supports legacy - /// representations that are normalized into another current in-memory array. Implementations - /// must validate the contract of that exact ID rather than accepting every form understood by - /// the current in-memory representation under an older ID. + /// `serialized_id` identifies the exact representation encountered on disk. Callers **must** + /// pass an ID from [`read_ids`](Self::read_ids), which the session's registry guarantees. A + /// plugin that reads one ID can therefore ignore `serialized_id`. Passing another ID gives an + /// error or an incorrect array. + /// + /// The returned array does not necessarily have to use this plugin's in-memory ID; this + /// supports legacy representations that are normalized into another current in-memory array. + /// Implementations must validate the contract of that exact ID rather than accepting every + /// form understood by the current in-memory representation under an older ID. + /// + /// Use `session` to resolve plugin-owned metadata, such as extension dtypes, instead of global + /// state. fn deserialize( &self, parts: ArrayDeserialization<'_>, @@ -164,53 +186,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 read_ids(&self) -> Vec { - vec![self.id()] - } - - fn write_ids(&self) -> Vec { - vec![self.id()] - } - - fn serialize( - &self, - array: &ArrayRef, - session: &VortexSession, - ) -> VortexResult> { - vortex_ensure_eq!( - self.id(), - array.encoding_id(), - "array plugin cannot serialize an in-memory array of another encoding", - ); - 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_eq!( - self.id(), - parts.serialized_id, - "array plugin does not recognize the 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 39745b92fbd..3c663fcbbf9 100644 --- a/vortex-array/src/array/vtable/mod.rs +++ b/vortex-array/src/array/vtable/mod.rs @@ -5,7 +5,8 @@ //! //! 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. Serialization lives +//! in the separate [`ArrayPlugin`] trait, which only serializable encodings implement. //! //! 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 +27,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 +47,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 +54,8 @@ 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 an implementation of the sealed `DynArrayData` trait. A +/// [`VTable`] does not serialize: an encoding with a wire form also implements [`ArrayPlugin`]. /// /// 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,35 +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. - /// - /// Buffers and children are serialized separately through [`buffer`](Self::buffer), - /// [`nbuffers`](Self::nbuffers), and child traversal. - /// - /// Return `Ok(None)` if this encoding cannot be serialized. This includes in-memory-only - /// encodings and encodings whose wire format is owned by a separate [`ArrayPlugin`]. Do not - /// return an error for this case: the caller decides how to report an unserializable array. - /// Return an error only when an array that can be serialized fails to serialize. - 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 fd481570e0a..79e0bd471c0 100644 --- a/vortex-array/src/arrays/bool/vtable/mod.rs +++ b/vortex-array/src/arrays/bool/vtable/mod.rs @@ -13,9 +13,13 @@ use vortex_error::vortex_ensure_eq; use vortex_error::vortex_panic; use vortex_session::VortexSession; +use crate::ArrayDeserialization; +use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::ExecutionCtx; use crate::ExecutionResult; +use crate::IntoArray; use crate::array::Array; use crate::array::ArrayParts; use crate::array::ArrayView; @@ -27,7 +31,6 @@ 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; @@ -114,20 +117,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, @@ -158,41 +147,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, - slots, - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { BoolSlots::NAMES[idx].to_string() } @@ -221,6 +175,73 @@ impl VTable for Bool { } } +impl ArrayPlugin for Bool { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_::(); + 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( + ArrayPlugin::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + + 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, + slots, + ))? + .into_array()) + } +} + #[derive(Clone, Debug)] pub struct Bool; diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index af6709b3f2b..84ab7ca7d67 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -15,9 +15,12 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::ArrayDeserialization; use crate::ArrayEq; use crate::ArrayHash; +use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::Canonical; use crate::EqMode; use crate::ExecutionCtx; @@ -41,7 +44,6 @@ 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 validity; @@ -158,76 +160,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), - slots, - )) - } - fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, @@ -295,3 +227,89 @@ impl VTable for Chunked { PARENT_RULES.evaluate(array, parent, child_idx) } } + +impl ArrayPlugin for Chunked { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + children, + .. + } = parts; + + 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), + slots, + ))? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index 3b55071ffb8..209f9c42e84 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -13,10 +13,13 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::ArrayDeserialization; use crate::ArrayEq; use crate::ArrayHash; use crate::ArrayParts; +use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::ArraySlots; use crate::EqMode; use crate::ExecutionCtx; @@ -52,7 +55,6 @@ 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 validity; @@ -133,42 +135,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_eq!(buffers.len(), 1); - - 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), - ArraySlots::new(), - )) - } - fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, @@ -283,6 +249,56 @@ impl VTable for Constant { } } +impl ArrayPlugin for Constant { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + // HACK: Because the scalar is stored in the buffers, we do not need to serialize the + // metadata at all. + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + buffers, + .. + } = parts; + + vortex_ensure_eq!(buffers.len(), 1); + + 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), + ArraySlots::new(), + ))? + .into_array()) + } +} + /// Appends the constant list `array` as one run sharing a single copy of its elements. /// /// The list's elements materialize once, and diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index 9153448d4c0..934e84d0dad 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -3,8 +3,10 @@ //! The [`Decimal`] vtable and the plugin that owns its wire form. //! -//! The vtable does not serialize or deserialize. [`DecimalPlugin`] reads and writes the +//! The vtable does not implement [`ArrayPlugin`]. [`DecimalPlugin`] reads and writes the //! `vortex.decimal` wire form, including legacy nodes. +//! +//! [`ArrayPlugin`]: crate::array::ArrayPlugin mod kernel; mod operations; @@ -42,7 +44,6 @@ use crate::builders::DecimalBuilder; use crate::dtype::DType; use crate::hash::ArrayEq; use crate::hash::ArrayHash; -use crate::serde::ArrayChildren; /// A [`Decimal`]-encoded Vortex array. pub type DecimalArray = Array; @@ -103,13 +104,6 @@ impl VTable for Decimal { )) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(None) - } - fn validate( &self, data: &DecimalData, @@ -138,18 +132,6 @@ impl VTable for Decimal { Ok(()) } - fn deserialize( - &self, - _dtype: &DType, - _len: usize, - _metadata: &[u8], - _buffers: &[BufferHandle], - _children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_bail!("Decimal deserialization requires DecimalPlugin") - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { DecimalSlots::NAMES[idx].to_string() } @@ -180,7 +162,9 @@ impl VTable for Decimal { /// The vtable of [`DecimalArray`]. /// -/// This vtable does not serialize or deserialize. A session **must** register [`DecimalPlugin`] -/// in its place to read or write decimal arrays. +/// This vtable does not implement [`ArrayPlugin`]. A session **must** register [`DecimalPlugin`] +/// to read or write decimal arrays. +/// +/// [`ArrayPlugin`]: crate::array::ArrayPlugin #[derive(Clone, Debug)] pub struct Decimal; diff --git a/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs b/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs index 04925656780..8fc4bfbd29b 100644 --- a/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs @@ -66,7 +66,7 @@ impl DecimalMetadata { /// [`WidenArray::from_stored_decimal`], which scans its valid values on the host, so a legacy node /// in device memory fails to read. /// -/// The [`Decimal`] vtable does not serialize or deserialize, so the default [`ArraySession`] +/// The [`Decimal`] vtable does not implement [`ArrayPlugin`], so the default [`ArraySession`] /// registers this plugin in its place. /// /// [`DecimalArray`]: super::DecimalArray @@ -79,14 +79,6 @@ impl ArrayPlugin for DecimalPlugin { VTable::id(&Decimal) } - fn read_ids(&self) -> Vec { - vec![self.id()] - } - - fn write_ids(&self) -> Vec { - vec![self.id()] - } - fn serialize( &self, array: &ArrayRef, diff --git a/vortex-array/src/arrays/dict/array.rs b/vortex-array/src/arrays/dict/array.rs index 6c8a6f22a65..3e4832f4285 100644 --- a/vortex-array/src/arrays/dict/array.rs +++ b/vortex-array/src/arrays/dict/array.rs @@ -388,6 +388,8 @@ mod test { use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array::ArrayDeserialization; + use crate::array::ArrayPlugin; use crate::array::VTable; use crate::arrays::ChunkedArray; use crate::arrays::Dict; @@ -607,11 +609,14 @@ mod test { .encode_to_vec(); let children: Vec = Vec::new(); let result = Dict.deserialize( - &DType::Primitive(PType::I32, NonNullable), - 0, - &metadata, - &[], - &children, + ArrayDeserialization::new( + VTable::id(&Dict), + &DType::Primitive(PType::I32, NonNullable), + 0, + &metadata, + &[], + &children, + ), &SESSION, ); assert!(result.is_err_and(|err| { diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index 5b1b7e3c6d1..2e2df639a88 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/vortex-array/src/arrays/dict/vtable/mod.rs @@ -23,9 +23,12 @@ use super::DictParts; use super::array::DictSlots; use super::array::DictSlotsView; use crate::AnyCanonical; +use crate::ArrayDeserialization; use crate::ArrayEq; use crate::ArrayHash; +use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::Canonical; use crate::CanonicalView; use crate::EqMode; @@ -58,7 +61,6 @@ 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; @@ -136,63 +138,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)?; - let codes_ptype = metadata.get_codes_ptype()?; - 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(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) }, - smallvec![Some(codes), Some(values)], - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { DictSlots::NAMES[idx].to_string() } @@ -298,6 +243,80 @@ impl VTable for Dict { } } +impl ArrayPlugin for Dict { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_::(); + 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( + ArrayPlugin::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + children, + .. + } = parts; + + let metadata = DictMetadata::decode(metadata)?; + let codes_ptype = metadata.get_codes_ptype()?; + 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(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) }, + smallvec![Some(codes), Some(values)], + ))? + .into_array()) + } +} + /// Gathers the dictionary values straight into `builder` as views. /// /// The canonical route first takes the values to full logical length — allocating an intermediate diff --git a/vortex-array/src/arrays/extension/vtable/mod.rs b/vortex-array/src/arrays/extension/vtable/mod.rs index e8efe3a44f8..567f0848d90 100644 --- a/vortex-array/src/arrays/extension/vtable/mod.rs +++ b/vortex-array/src/arrays/extension/vtable/mod.rs @@ -10,10 +10,14 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::ArrayDeserialization; +use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::EmptyArrayData; use crate::ExecutionCtx; use crate::ExecutionResult; +use crate::IntoArray; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayParts; @@ -28,7 +32,6 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::ExtensionBuilder; use crate::dtype::DType; -use crate::serde::ArrayChildren; mod kernel; mod operations; @@ -138,45 +141,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, - ExtensionSlots { storage }.into_slots(), - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { ExtensionSlots::NAMES[idx].to_string() } @@ -208,3 +172,57 @@ impl VTable for Extension { PARENT_RULES.evaluate(array, parent, child_idx) } } + +impl ArrayPlugin for Extension { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + children, + .. + } = parts; + + 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, + 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 da5271b08bf..14f918cd50a 100644 --- a/vortex-array/src/arrays/filter/vtable.rs +++ b/vortex-array/src/arrays/filter/vtable.rs @@ -5,13 +5,11 @@ 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_ensure_eq; 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; @@ -43,7 +41,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. @@ -129,27 +126,6 @@ impl VTable for Filter { FilterSlots::NAMES[idx].to_string() } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - // TODO(joe): make this configurable - Ok(None) - } - - 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 0e140a47998..8a85a888896 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs @@ -14,13 +14,17 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::ArrayDeserialization; use crate::ArrayEq; use crate::ArrayHash; use crate::ArrayParts; +use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::EqMode; use crate::ExecutionCtx; use crate::ExecutionResult; +use crate::IntoArray; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; @@ -33,7 +37,6 @@ 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; @@ -99,13 +102,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, @@ -142,16 +138,57 @@ impl VTable for FixedSizeList { Ok(()) } - fn deserialize( + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + FixedSizeListSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(array)) + } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let Some(builder) = builder.as_any_mut().downcast_mut::() else { + vortex_bail!("append_to_builder for FixedSizeList requires a FixedSizeListBuilder"); + }; + builder.append_fixed_size_list_array(&array.into_owned(), ctx) + } +} + +impl ArrayPlugin for FixedSizeList { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + vec![], + ))) + } - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + if !metadata.is_empty() { vortex_bail!( "FixedSizeListArray expects empty metadata, got {} bytes", @@ -187,31 +224,13 @@ impl VTable for FixedSizeList { let data = FixedSizeListData::try_build(elements.clone(), *list_size, validity.clone(), len)?; let slots = FixedSizeListData::make_slots(&elements, &validity, len); - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots, - )) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - FixedSizeListSlots::NAMES[idx].to_string() - } - - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { - Ok(ExecutionResult::done(array)) - } - - fn append_to_builder( - array: ArrayView<'_, Self>, - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - let Some(builder) = builder.as_any_mut().downcast_mut::() else { - vortex_bail!("append_to_builder for FixedSizeList requires a FixedSizeListBuilder"); - }; - builder.append_fixed_size_list_array(&array.into_owned(), ctx) + ))? + .into_array()) } } diff --git a/vortex-array/src/arrays/interleave/mod.rs b/vortex-array/src/arrays/interleave/mod.rs index b9c5c7374ee..207e7b8aa37 100644 --- a/vortex-array/src/arrays/interleave/mod.rs +++ b/vortex-array/src/arrays/interleave/mod.rs @@ -48,7 +48,6 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; pub use self::array::InterleaveArray; @@ -73,7 +72,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; /// The [`Interleave`] encoding. See the [module docs](self). @@ -217,25 +215,6 @@ impl VTable for Interleave { InterleaveSlots::slot_name(idx) } - 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!("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 5aea3d244ba..7ecc9749ae6 100644 --- a/vortex-array/src/arrays/list/vtable/mod.rs +++ b/vortex-array/src/arrays/list/vtable/mod.rs @@ -14,9 +14,12 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::ArrayDeserialization; use crate::ArrayEq; use crate::ArrayHash; +use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::EqMode; use crate::ExecutionCtx; use crate::ExecutionResult; @@ -39,7 +42,6 @@ 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 validity; @@ -113,19 +115,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, @@ -156,16 +145,71 @@ impl VTable for List { Ok(()) } - fn deserialize( + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + ListSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done( + list_view_from_list(array, ctx)?.into_array(), + )) + } + + // The complexity comes from the expansion of `match_each_list_builder!`. + #[expect(clippy::cognitive_complexity)] + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + match match_each_list_builder!(&mut *builder, |b| b.append_list_array(array, ctx)) { + Some(result) => result, + None => vortex_bail!( + "cannot append a List array of dtype {} to a {} builder", + array.dtype(), + builder.dtype() + ), + } + } +} + +impl ArrayPlugin for List { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_::(); + 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( + ArrayPlugin::id(self), + array, + metadata, + ))) + } - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + children, + .. + } = parts; + let metadata = ListMetadata::decode(metadata)?; let offset_ptype = metadata.get_offset_ptype()?; let validity = if children.len() == 2 { @@ -199,40 +243,14 @@ impl VTable for List { &mut session.create_execution_ctx(), )?; let slots = ListData::make_slots(&elements, &offsets, &validity, len); - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots, - )) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - ListSlots::NAMES[idx].to_string() - } - - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - Ok(ExecutionResult::done( - list_view_from_list(array, ctx)?.into_array(), - )) - } - - // The complexity comes from the expansion of `match_each_list_builder!`. - #[expect(clippy::cognitive_complexity)] - fn append_to_builder( - array: ArrayView<'_, Self>, - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - match match_each_list_builder!(&mut *builder, |b| b.append_list_array(array, ctx)) { - Some(result) => result, - None => vortex_bail!( - "cannot append a List array of dtype {} to a {} builder", - array.dtype(), - builder.dtype() - ), - } + ))? + .into_array()) } } @@ -251,6 +269,8 @@ mod tests { use super::List; use super::ListMetadata; use crate::ArrayRef; + use crate::array::ArrayDeserialization; + use crate::array::ArrayPlugin; use crate::array::VTable; use crate::dtype::DType; use crate::dtype::Nullability; @@ -270,7 +290,10 @@ mod tests { } .encode_to_vec(); let children: Vec = Vec::new(); - let result = List.deserialize(&dtype, 0, &metadata, &[], &children, &SESSION); + let result = List.deserialize( + ArrayDeserialization::new(VTable::id(&List), &dtype, 0, &metadata, &[], &children), + &SESSION, + ); assert!(result.is_err_and(|err| { err.to_string() .contains("ListMetadata offset_ptype must be a known PType, got 11") diff --git a/vortex-array/src/arrays/listview/vtable/mod.rs b/vortex-array/src/arrays/listview/vtable/mod.rs index c2cac004a31..9e9985ac526 100644 --- a/vortex-array/src/arrays/listview/vtable/mod.rs +++ b/vortex-array/src/arrays/listview/vtable/mod.rs @@ -16,13 +16,17 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::ArrayDeserialization; use crate::ArrayEq; use crate::ArrayHash; use crate::ArrayParts; +use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::EqMode; use crate::ExecutionCtx; use crate::ExecutionResult; +use crate::IntoArray; use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayId; @@ -39,7 +43,6 @@ 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; @@ -126,20 +129,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, @@ -174,16 +163,79 @@ impl VTable for ListView { Ok(()) } - fn deserialize( + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + ListViewSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(array)) + } + + // The complexity comes from the expansion of `match_each_list_builder!`. + #[expect(clippy::cognitive_complexity)] + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + match match_each_list_builder!(&mut *builder, |b| b.append_listview_array(array, ctx)) { + Some(result) => result, + None => vortex_bail!( + "cannot append a ListView array of dtype {} to a {} builder", + array.dtype(), + builder.dtype() + ), + } + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } +} + +impl ArrayPlugin for ListView { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_::(); + 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(); - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + let metadata = ListViewMetadata::decode(metadata)?; let offset_ptype = metadata.get_offset_ptype()?; let size_ptype = metadata.get_size_ptype()?; @@ -238,46 +290,14 @@ impl VTable for ListView { )?; let data = ListViewData::try_new()?; let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len); - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots, - )) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - ListViewSlots::NAMES[idx].to_string() - } - - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { - Ok(ExecutionResult::done(array)) - } - - // The complexity comes from the expansion of `match_each_list_builder!`. - #[expect(clippy::cognitive_complexity)] - fn append_to_builder( - array: ArrayView<'_, Self>, - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - match match_each_list_builder!(&mut *builder, |b| b.append_listview_array(array, ctx)) { - Some(result) => result, - None => vortex_bail!( - "cannot append a ListView array of dtype {} to a {} builder", - array.dtype(), - builder.dtype() - ), - } - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - PARENT_RULES.evaluate(array, parent, child_idx) + ))? + .into_array()) } } @@ -294,6 +314,8 @@ mod tests { use super::ListView; use super::ListViewMetadata; use crate::ArrayRef; + use crate::array::ArrayDeserialization; + use crate::array::ArrayPlugin; use crate::array::VTable; use crate::dtype::DType; use crate::dtype::Nullability; @@ -323,7 +345,10 @@ mod tests { ); let metadata = with_ptype(11).encode_to_vec(); let children: Vec = Vec::new(); - let result = ListView.deserialize(&dtype, 0, &metadata, &[], &children, &SESSION); + let result = ListView.deserialize( + ArrayDeserialization::new(VTable::id(&ListView), &dtype, 0, &metadata, &[], &children), + &SESSION, + ); assert!(result.is_err_and(|err| err.to_string().contains(expected))); assert_eq!(get_ptype(&with_ptype(PType::U32 as i32))?, PType::U32); diff --git a/vortex-array/src/arrays/map/vtable/mod.rs b/vortex-array/src/arrays/map/vtable/mod.rs index 17980991d3a..cdae1d1479b 100644 --- a/vortex-array/src/arrays/map/vtable/mod.rs +++ b/vortex-array/src/arrays/map/vtable/mod.rs @@ -9,10 +9,14 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::ArrayDeserialization; use crate::ArrayParts; +use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::ExecutionCtx; use crate::ExecutionResult; +use crate::IntoArray; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; @@ -29,7 +33,6 @@ 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; @@ -96,22 +99,69 @@ impl VTable for Map { with_empty_buffers(self, array, buffers) } + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + MapSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(array)) + } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + match match_each_map_builder!(&mut *builder, |b| b.append_map_array(array, ctx)) { + Some(result) => result, + None => vortex_bail!( + "cannot append a Map array of dtype {} to a {} builder", + array.dtype(), + builder.dtype() + ), + } + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } +} + +impl ArrayPlugin for Map { + fn id(&self) -> ArrayId { + VTable::id(self) + } + fn serialize( - _array: ArrayView<'_, Self>, + &self, + array: &ArrayRef, _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) + ) -> VortexResult> { + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + vec![], + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + if !metadata.is_empty() { vortex_bail!( "MapArray expects empty metadata, got {} bytes", @@ -135,43 +185,13 @@ impl VTable for Map { ); let slots = MapData::make_slots(entries); - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, MapData, slots, - )) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - MapSlots::NAMES[idx].to_string() - } - - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { - Ok(ExecutionResult::done(array)) - } - - fn append_to_builder( - array: ArrayView<'_, Self>, - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - match match_each_map_builder!(&mut *builder, |b| b.append_map_array(array, ctx)) { - Some(result) => result, - None => vortex_bail!( - "cannot append a Map array of dtype {} to a {} builder", - array.dtype(), - builder.dtype() - ), - } - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - PARENT_RULES.evaluate(array, parent, child_idx) + ))? + .into_array()) } } diff --git a/vortex-array/src/arrays/masked/vtable/mod.rs b/vortex-array/src/arrays/masked/vtable/mod.rs index 507223a94a0..d679d0b93c9 100644 --- a/vortex-array/src/arrays/masked/vtable/mod.rs +++ b/vortex-array/src/arrays/masked/vtable/mod.rs @@ -16,10 +16,13 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::AnyCanonical; +use crate::ArrayDeserialization; use crate::ArrayEq; use crate::ArrayHash; use crate::ArrayParts; +use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::Canonical; use crate::EqMode; use crate::IntoArray; @@ -43,7 +46,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 [`Masked`]-encoded Vortex array. pub type MaskedArray = Array; @@ -115,23 +117,74 @@ impl VTable for Masked { with_empty_buffers(self, array, buffers) } + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let array = require_child!(array, array.child(), MaskedSlots::CHILD => AnyCanonical); + + let validity = array.masked_validity(); + + // Fast path: all masked means result is all nulls. + if validity.definitely_all_null() { + return Ok(ExecutionResult::done( + ConstantArray::new(Scalar::null(array.dtype().as_nullable()), array.len()) + .into_array(), + )); + } + + // NB: We intentionally do NOT have a fast path for `validity_mask.all_true()`. + // `MaskedArray`'s dtype is always `Nullable`, but the child has `NonNullable` `DType` (by + // invariant). Simply returning the child's canonical would cause a dtype mismatch. + // While we could manually convert the dtype, `mask_validity_canonical` is already O(1) for + // `AllTrue` masks (no data copying), so there's no benefit. + + let child = Canonical::from(array.child().as_::()); + Ok(ExecutionResult::done( + mask_validity_canonical(child, validity, ctx)?.into_array(), + )) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + MaskedSlots::NAMES[idx].to_string() + } +} + +impl ArrayPlugin for Masked { + fn id(&self) -> ArrayId { + VTable::id(self) + } + fn serialize( - _array: ArrayView<'_, Self>, + &self, + array: &ArrayRef, _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) + ) -> VortexResult> { + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + vec![], + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + if !metadata.is_empty() { vortex_bail!( "MaskedArray expects empty metadata, got {} bytes", @@ -163,49 +216,14 @@ impl VTable for Masked { child.all_valid(&mut session.create_execution_ctx())?, validity, )?; - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, smallvec![Some(child), validity_slot], - )) - } - - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - let array = require_child!(array, array.child(), MaskedSlots::CHILD => AnyCanonical); - - let validity = array.masked_validity(); - - // Fast path: all masked means result is all nulls. - if validity.definitely_all_null() { - return Ok(ExecutionResult::done( - ConstantArray::new(Scalar::null(array.dtype().as_nullable()), array.len()) - .into_array(), - )); - } - - // NB: We intentionally do NOT have a fast path for `validity_mask.all_true()`. - // `MaskedArray`'s dtype is always `Nullable`, but the child has `NonNullable` `DType` (by - // invariant). Simply returning the child's canonical would cause a dtype mismatch. - // While we could manually convert the dtype, `mask_validity_canonical` is already O(1) for - // `AllTrue` masks (no data copying), so there's no benefit. - - let child = Canonical::from(array.child().as_::()); - Ok(ExecutionResult::done( - mask_validity_canonical(child, validity, ctx)?.into_array(), - )) - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - PARENT_RULES.evaluate(array, parent, child_idx) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - MaskedSlots::NAMES[idx].to_string() + ))? + .into_array()) } } diff --git a/vortex-array/src/arrays/null/mod.rs b/vortex-array/src/arrays/null/mod.rs index e55212db65a..d523123afdc 100644 --- a/vortex-array/src/arrays/null/mod.rs +++ b/vortex-array/src/arrays/null/mod.rs @@ -9,11 +9,14 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::ArrayDeserialization; +use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::ArraySlots; -use crate::ArrayVTable; use crate::ExecutionCtx; use crate::ExecutionResult; +use crate::IntoArray; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayParts; @@ -33,7 +36,6 @@ use crate::builders::NullBuilder; use crate::dtype::DType; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar::Scalar; -use crate::serde::ArrayChildren; use crate::validity::Validity; pub(crate) mod compute; @@ -42,9 +44,11 @@ pub(crate) mod compute; pub type NullArray = Array; pub(crate) fn initialize(session: &VortexSession) { - session - .kernels() - .register_execute_parent_kernel(Dict.id(), Null, TakeExecuteAdaptor(Null)); + session.kernels().register_execute_parent_kernel( + VTable::id(&Dict), + Null, + TakeExecuteAdaptor(Null), + ); } impl VTable for Null { @@ -93,37 +97,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, - ArraySlots::new(), - )) - } - fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, @@ -149,6 +122,51 @@ impl VTable for Null { } } +impl ArrayPlugin for Null { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + .. + } = parts; + + 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, + ArraySlots::new(), + ))? + .into_array()) + } +} + /// A array where all values are null. /// /// This mirrors the Apache Arrow Null array encoding and provides an efficient representation diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs index 01c71d52997..f8a57a8fe5f 100644 --- a/vortex-array/src/arrays/patched/vtable/mod.rs +++ b/vortex-array/src/arrays/patched/vtable/mod.rs @@ -3,8 +3,11 @@ use prost::Message; +use crate::ArrayDeserialization; use crate::ArrayEq; use crate::ArrayHash; +use crate::ArrayPlugin; +use crate::ArraySerialization; mod kernels; mod operations; mod slice; @@ -50,7 +53,6 @@ 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,60 +140,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, - slots, - )) - } - fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, @@ -317,6 +265,77 @@ impl VTable for Patched { } } +impl ArrayPlugin for Patched { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_::(); + 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( + ArrayPlugin::id(self), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + children, + .. + } = parts; + + 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, + slots, + ))? + .into_array()) + } +} + /// Apply patches on top of the existing value types. fn apply_patches_primitive( output: &mut [V], diff --git a/vortex-array/src/arrays/piecewise_sequence/tests.rs b/vortex-array/src/arrays/piecewise_sequence/tests.rs index 8715c439b5c..d75673e016f 100644 --- a/vortex-array/src/arrays/piecewise_sequence/tests.rs +++ b/vortex-array/src/arrays/piecewise_sequence/tests.rs @@ -175,7 +175,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 db5b33a1081..0ab653fbf9a 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_eq; 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. @@ -105,25 +103,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 df4006626a2..b8869f7260c 100644 --- a/vortex-array/src/arrays/primitive/vtable/mod.rs +++ b/vortex-array/src/arrays/primitive/vtable/mod.rs @@ -6,10 +6,14 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; +use crate::ArrayDeserialization; use crate::ArrayParts; +use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::ExecutionCtx; use crate::ExecutionResult; +use crate::IntoArray; use crate::array::Array; use crate::array::ArrayView; use crate::array::VTable; @@ -21,7 +25,6 @@ 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 validity; @@ -97,13 +100,6 @@ impl VTable for Primitive { )) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - fn validate( &self, data: &PrimitiveData, @@ -132,16 +128,68 @@ impl VTable for Primitive { Ok(()) } - fn deserialize( + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + PrimitiveSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(array)) + } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + match_each_native_ptype!(array.ptype(), |P| { + if let Some(builder) = builder.as_any_mut().downcast_mut::>() { + return builder.append_primitive_array(&array.into_owned(), ctx); + } + }); + + vortex_bail!("append_to_builder for Primitive requires a matching PrimitiveBuilder"); + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + RULES.evaluate(array, parent, child_idx) + } +} + +impl ArrayPlugin for Primitive { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + vec![], + ))) + } - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + if !metadata.is_empty() { vortex_bail!( "PrimitiveArray expects empty metadata, got {} bytes", @@ -172,43 +220,14 @@ impl VTable for Primitive { // 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( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots, - )) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - PrimitiveSlots::NAMES[idx].to_string() - } - - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { - Ok(ExecutionResult::done(array)) - } - - fn append_to_builder( - array: ArrayView<'_, Self>, - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - match_each_native_ptype!(array.ptype(), |P| { - if let Some(builder) = builder.as_any_mut().downcast_mut::>() { - return builder.append_primitive_array(&array.into_owned(), ctx); - } - }); - - vortex_bail!("append_to_builder for Primitive requires a matching PrimitiveBuilder"); - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - RULES.evaluate(array, parent, child_idx) + ))? + .into_array()) } } diff --git a/vortex-array/src/arrays/runend/array.rs b/vortex-array/src/arrays/runend/array.rs index 3a4e5493387..ecfb79455c3 100644 --- a/vortex-array/src/arrays/runend/array.rs +++ b/vortex-array/src/arrays/runend/array.rs @@ -20,11 +20,14 @@ use vortex_session::registry::CachedId; use crate::AnyCanonical; use crate::Array; +use crate::ArrayDeserialization; use crate::ArrayEq; use crate::ArrayHash; use crate::ArrayId; use crate::ArrayParts; +use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::ArrayView; use crate::EqMode; use crate::ExecutionCtx; @@ -52,7 +55,6 @@ use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::require_child; -use crate::serde::ArrayChildren; use crate::validity::Validity; use crate::vtable::VTable; use crate::vtable::ValidityVTable; @@ -139,30 +141,64 @@ impl VTable for RunEnd { crate::vtable::with_empty_buffers(self, array, buffers) } - fn serialize( + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + RunEndSlots::NAMES[idx].to_string() + } + + fn reduce_parent( array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + RULES.evaluate(array, parent, child_idx) + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let array = require_child!(array, array.ends(), RunEndSlots::ENDS => Primitive); + let array = require_child!(array, array.values(), RunEndSlots::VALUES => AnyCanonical); + run_end_canonicalize(&array, ctx).map(ExecutionResult::done) + } +} + +impl ArrayPlugin for RunEnd { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, _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(), - )) + ) -> VortexResult> { + let view = array.as_::(); + 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( + ArrayPlugin::id(self), + array, + metadata, + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + children, + .. + } = parts; + let metadata = RunEndMetadata::decode(metadata)?; let ends_ptype = metadata.get_ends_ptype()?; let ends_dtype = DType::Primitive(ends_ptype, Nullability::NonNullable); @@ -180,31 +216,14 @@ impl VTable for RunEnd { )?; let slots = RunEndSlots { ends, values }.into_slots(); let data = RunEndData::new(offset); - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots, - )) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - RunEndSlots::NAMES[idx].to_string() - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - RULES.evaluate(array, parent, child_idx) - } - - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - let array = require_child!(array, array.ends(), RunEndSlots::ENDS => Primitive); - let array = require_child!(array, array.values(), RunEndSlots::VALUES => AnyCanonical); - run_end_canonicalize(&array, ctx).map(ExecutionResult::done) + ))? + .into_array()) } } @@ -596,6 +615,8 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array::ArrayDeserialization; + use crate::array::ArrayPlugin; use crate::arrays::DecimalArray; use crate::arrays::DictArray; use crate::arrays::ListArray; @@ -956,11 +977,14 @@ mod tests { .encode_to_vec(); let children: Vec = Vec::new(); let result = RunEnd.deserialize( - &DType::Primitive(PType::I32, Nullability::NonNullable), - 0, - &metadata, - &[], - &children, + ArrayDeserialization::new( + VTable::id(&RunEnd), + &DType::Primitive(PType::I32, Nullability::NonNullable), + 0, + &metadata, + &[], + &children, + ), &SESSION, ); assert!(result.is_err_and(|err| { diff --git a/vortex-array/src/arrays/scalar_fn/plugin.rs b/vortex-array/src/arrays/scalar_fn/plugin.rs index 43a5616bc98..4a4a7d3d5ab 100644 --- a/vortex-array/src/arrays/scalar_fn/plugin.rs +++ b/vortex-array/src/arrays/scalar_fn/plugin.rs @@ -62,14 +62,6 @@ impl ArrayPlugin for ScalarFnArrayPlugi self.0.id() } - fn read_ids(&self) -> Vec { - vec![self.id()] - } - - fn write_ids(&self) -> Vec { - vec![self.id()] - } - fn serialize( &self, array: &ArrayRef, diff --git a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs index d2827ff68c3..9e8df5e46ee 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs @@ -11,11 +11,9 @@ use std::ops::Deref; use itertools::Itertools; use validity::ValidityArray; 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_session::VortexSession; use crate::ArrayEq; use crate::ArrayHash; @@ -45,7 +43,6 @@ use crate::scalar_fn::ReduceNodeValidity; use crate::scalar_fn::ScalarFnExecutionStep; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::VecExecutionArgs; -use crate::serde::ArrayChildren; use crate::validity::Validity; /// A [`ScalarFn`]-encoded Vortex array. @@ -144,26 +141,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/scalar_fn/vtable/validity.rs b/vortex-array/src/arrays/scalar_fn/vtable/validity.rs index 71c5a2985b4..d0569217b4b 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/validity.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/validity.rs @@ -2,9 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::ArrayRef; @@ -24,7 +22,6 @@ use crate::dtype::Nullability; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; use crate::scalar::Scalar; -use crate::serde::ArrayChildren; use crate::validity::Validity; #[array_slots(ScalarFnValidity)] @@ -96,25 +93,6 @@ impl VTable for ScalarFnValidity { with_empty_buffers(self, array, buffers) } - 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!("ScalarFnValidity deserialize not supported"); - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { ValiditySlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/slice/vtable.rs b/vortex-array/src/arrays/slice/vtable.rs index 76f98c41b9f..241dea96031 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -9,12 +9,10 @@ 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_ensure_eq; use vortex_error::vortex_err; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::AnyCanonical; @@ -41,7 +39,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. @@ -129,27 +126,6 @@ impl VTable for Slice { SliceSlots::NAMES[idx].to_string() } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - // TODO(joe): make this configurable - Ok(None) - } - - 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 d258d82ad6e..a8d9fb08169 100644 --- a/vortex-array/src/arrays/struct_/vtable/mod.rs +++ b/vortex-array/src/arrays/struct_/vtable/mod.rs @@ -7,9 +7,13 @@ use vortex_error::vortex_bail; use vortex_error::vortex_panic; use vortex_session::VortexSession; +use crate::ArrayDeserialization; +use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::ExecutionCtx; use crate::ExecutionResult; +use crate::IntoArray; use crate::array::Array; use crate::array::ArrayParts; use crate::array::ArrayView; @@ -24,7 +28,6 @@ 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; @@ -131,23 +134,68 @@ impl VTable for Struct { with_empty_buffers(self, array, buffers) } + fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { + if idx == StructSlots::VALIDITY { + "validity".to_string() + } else { + array.dtype().as_struct_fields().names()[idx - StructSlots::FIELDS_OFFSET].to_string() + } + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(array)) + } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let Some(builder) = builder.as_any_mut().downcast_mut::() else { + vortex_bail!("append_to_builder for Struct requires a StructBuilder"); + }; + builder.append_struct_array(&array.into_owned(), ctx) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } +} + +impl ArrayPlugin for Struct { + fn id(&self) -> ArrayId { + VTable::id(self) + } + fn serialize( - _array: ArrayView<'_, Self>, + &self, + array: &ArrayRef, _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) + ) -> VortexResult> { + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + vec![], + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + children, + .. + } = parts; + if !metadata.is_empty() { vortex_bail!( "StructArray expects empty metadata, got {} bytes", @@ -184,44 +232,14 @@ impl VTable for Struct { )?)); } - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, EmptyArrayData, slots, - )) - } - - fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { - if idx == StructSlots::VALIDITY { - "validity".to_string() - } else { - array.dtype().as_struct_fields().names()[idx - StructSlots::FIELDS_OFFSET].to_string() - } - } - - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { - Ok(ExecutionResult::done(array)) - } - - fn append_to_builder( - array: ArrayView<'_, Self>, - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - let Some(builder) = builder.as_any_mut().downcast_mut::() else { - vortex_bail!("append_to_builder for Struct requires a StructBuilder"); - }; - builder.append_struct_array(&array.into_owned(), ctx) - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - PARENT_RULES.evaluate(array, parent, child_idx) + ))? + .into_array()) } } diff --git a/vortex-array/src/arrays/union/vtable/mod.rs b/vortex-array/src/arrays/union/vtable/mod.rs index 7800f9430d2..d9725d3221e 100644 --- a/vortex-array/src/arrays/union/vtable/mod.rs +++ b/vortex-array/src/arrays/union/vtable/mod.rs @@ -10,9 +10,13 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::ArrayDeserialization; +use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::ExecutionCtx; use crate::ExecutionResult; +use crate::IntoArray; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayParts; @@ -28,7 +32,6 @@ 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 validate; @@ -100,22 +103,66 @@ impl VTable for Union { with_empty_buffers(self, array, buffers) } - fn serialize( + fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { + if idx == UnionSlots::TYPE_IDS { + "type_ids".to_string() + } else { + array.variants().names()[idx - UnionSlots::CHILDREN_OFFSET].to_string() + } + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(array)) + } + + fn append_to_builder( _array: ArrayView<'_, Self>, + _builder: &mut dyn ArrayBuilder, + _ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + vortex_bail!("TODO(connor)[Union]: implement append_to_builder for Union arrays") + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } +} + +impl ArrayPlugin for Union { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) + ) -> VortexResult> { + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + vec![], + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + 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 { @@ -134,38 +181,11 @@ impl VTable for Union { .map(|(index, dtype)| children.get(UnionSlots::CHILDREN_OFFSET + index, &dtype, len)) .collect::>>()?; - Ok(make_union_parts( + Ok(Array::::try_from_parts(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() - } else { - array.variants().names()[idx - UnionSlots::CHILDREN_OFFSET].to_string() - } - } - - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { - Ok(ExecutionResult::done(array)) - } - - fn append_to_builder( - _array: ArrayView<'_, Self>, - _builder: &mut dyn ArrayBuilder, - _ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - vortex_bail!("TODO(connor)[Union]: implement append_to_builder for Union arrays") - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - PARENT_RULES.evaluate(array, parent, child_idx) + ))? + .into_array()) } } diff --git a/vortex-array/src/arrays/varbin/vtable/mod.rs b/vortex-array/src/arrays/varbin/vtable/mod.rs index e20611c4671..af537af6717 100644 --- a/vortex-array/src/arrays/varbin/vtable/mod.rs +++ b/vortex-array/src/arrays/varbin/vtable/mod.rs @@ -13,8 +13,11 @@ use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::registry::CachedId; +use crate::ArrayDeserialization; use crate::ArrayParts; +use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::ExecutionCtx; use crate::ExecutionResult; use crate::IntoArray; @@ -36,7 +39,6 @@ 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; @@ -155,28 +157,82 @@ impl VTable for VarBin { )) } - fn serialize( + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + VarBinSlots::NAMES[idx].to_string() + } + + fn reduce_parent( 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(), + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if let Some(result) = + match_each_varbin_builder!(builder, |builder| builder.append_varbin(array, ctx)) + { + return result; + } + + // The two arms here are every builder a `Utf8`/`Binary` dtype has: all four + // `VarBinBuilder` widths above, and `VarBinViewBuilder` below. + let Some(builder) = builder.as_any_mut().downcast_mut::() else { + vortex_bail!("append_to_builder for VarBin requires a variable-binary builder") + }; + append_to_varbinview(array, builder, ctx) + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done( + varbin_to_canonical(array.as_view(), ctx)?.into_array(), )) } +} + +impl ArrayPlugin for VarBin { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_::(); + 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( + ArrayPlugin::id(self), + array, + metadata, + ))) + } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + let metadata = VarBinMetadata::decode(metadata)?; let offsets_ptype = metadata.get_offsets_ptype()?; let validity = if children.len() == 1 { @@ -207,50 +263,14 @@ impl VTable for VarBin { &mut session.create_execution_ctx(), )?; let slots = VarBinData::make_slots(offsets, &validity, len); - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots, - )) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - VarBinSlots::NAMES[idx].to_string() - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - PARENT_RULES.evaluate(array, parent, child_idx) - } - - fn append_to_builder( - array: ArrayView<'_, Self>, - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - if let Some(result) = - match_each_varbin_builder!(builder, |builder| builder.append_varbin(array, ctx)) - { - return result; - } - - // The two arms here are every builder a `Utf8`/`Binary` dtype has: all four - // `VarBinBuilder` widths above, and `VarBinViewBuilder` below. - let Some(builder) = builder.as_any_mut().downcast_mut::() else { - vortex_bail!("append_to_builder for VarBin requires a variable-binary builder") - }; - append_to_varbinview(array, builder, ctx) - } - - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - Ok(ExecutionResult::done( - varbin_to_canonical(array.as_view(), ctx)?.into_array(), - )) + ))? + .into_array()) } } @@ -296,6 +316,8 @@ mod tests { use super::VarBin; use super::VarBinMetadata; use crate::ArrayRef; + use crate::array::ArrayDeserialization; + use crate::array::ArrayPlugin; use crate::array::VTable; use crate::dtype::DType; use crate::dtype::Nullability; @@ -308,11 +330,14 @@ mod tests { let metadata = VarBinMetadata { offsets_ptype: 11 }.encode_to_vec(); let children: Vec = Vec::new(); let result = VarBin.deserialize( - &DType::Utf8(Nullability::NonNullable), - 0, - &metadata, - &[], - &children, + ArrayDeserialization::new( + VTable::id(&VarBin), + &DType::Utf8(Nullability::NonNullable), + 0, + &metadata, + &[], + &children, + ), &SESSION, ); assert!(result.is_err_and(|err| { diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index f3d8d4147ed..f88a6538904 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -15,11 +15,15 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::ArrayDeserialization; use crate::ArrayParts; +use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::EqMode; use crate::ExecutionCtx; use crate::ExecutionResult; +use crate::IntoArray; use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayId; @@ -36,7 +40,6 @@ 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; @@ -152,23 +155,70 @@ impl VTable for VarBinView { )) } + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + VarBinViewSlots::NAMES[idx].to_string() + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if let Some(builder) = builder.as_any_mut().downcast_mut::() { + return builder.append_varbinview_array(&array.into_owned(), ctx); + } + if let Some(result) = + match_each_varbin_builder!(builder, |builder| builder.append_varbinview(array, ctx)) + { + return result; + } + vortex_bail!("append_to_builder for VarBinView requires a variable-binary builder") + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(array)) + } +} + +impl ArrayPlugin for VarBinView { + fn id(&self) -> ArrayId { + VTable::id(self) + } + fn serialize( - _array: ArrayView<'_, Self>, + &self, + array: &ArrayRef, _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) + ) -> VortexResult> { + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + vec![], + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + if !metadata.is_empty() { vortex_bail!( "VarBinViewArray expects empty metadata, got {} bytes", @@ -209,13 +259,14 @@ impl VTable for VarBinView { validity.clone(), )?; let slots = VarBinViewData::make_slots(&validity, len); - return Ok(ArrayParts::new( + return Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots, - )); + ))? + .into_array()); } let data_buffers = data_handles @@ -232,45 +283,14 @@ impl VTable for VarBinView { &mut session.create_execution_ctx(), )?; let slots = VarBinViewData::make_slots(&validity, len); - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, data, slots, - )) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - VarBinViewSlots::NAMES[idx].to_string() - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - PARENT_RULES.evaluate(array, parent, child_idx) - } - - fn append_to_builder( - array: ArrayView<'_, Self>, - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - if let Some(builder) = builder.as_any_mut().downcast_mut::() { - return builder.append_varbinview_array(&array.into_owned(), ctx); - } - if let Some(result) = - match_each_varbin_builder!(builder, |builder| builder.append_varbinview(array, ctx)) - { - return result; - } - vortex_bail!("append_to_builder for VarBinView requires a variable-binary builder") - } - - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { - Ok(ExecutionResult::done(array)) + ))? + .into_array()) } } diff --git a/vortex-array/src/arrays/variant/vtable/mod.rs b/vortex-array/src/arrays/variant/vtable/mod.rs index 47cb8f0f18e..fa44ef91776 100644 --- a/vortex-array/src/arrays/variant/vtable/mod.rs +++ b/vortex-array/src/arrays/variant/vtable/mod.rs @@ -15,9 +15,13 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use vortex_utils::aliases::hash_set::HashSet; +use crate::ArrayDeserialization; +use crate::ArrayPlugin; use crate::ArrayRef; +use crate::ArraySerialization; use crate::ExecutionCtx; use crate::ExecutionResult; +use crate::IntoArray; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayParts; @@ -36,7 +40,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; @@ -125,28 +128,64 @@ impl VTable for Variant { with_empty_buffers(self, array, buffers) } - fn serialize( + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + match VariantSlots::NAMES.get(idx) { + Some(name) => (*name).to_string(), + None => vortex_panic!("VariantArray slot_name index {idx} out of bounds"), + } + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(array)) + } + + fn reduce_parent( array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + RULES.evaluate(array, parent, child_idx) + } +} + +impl ArrayPlugin for Variant { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn serialize( + &self, + array: &ArrayRef, _session: &VortexSession, - ) -> VortexResult>> { - let shredded_dtype = array.slots()[VariantSlots::SHREDDED] + ) -> VortexResult> { + let view = array.as_::(); + let shredded_dtype = view.slots()[VariantSlots::SHREDDED] .as_ref() .map(|shredded| shredded.dtype().try_into()) .transpose()?; - Ok(Some( - VariantMetadataProto { shredded_dtype }.encode_to_vec(), - )) + let metadata = VariantMetadataProto { shredded_dtype }.encode_to_vec(); + + Ok(Some(ArraySerialization::from_array( + ArrayPlugin::id(self), + array, + metadata, + ))) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + parts: ArrayDeserialization<'_>, session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + vortex_ensure!( buffers.is_empty(), "VariantArray expects 0 buffers, got {}", @@ -165,7 +204,7 @@ impl VTable for Variant { let shredded = shredded_dtype .map(|dtype| children.get(1, &dtype, len)) .transpose()?; - Ok(ArrayParts::new( + Ok(Array::::try_from_parts(ArrayParts::new( self.clone(), dtype.clone(), len, @@ -175,26 +214,8 @@ impl VTable for Variant { shredded, } .into_slots(), - )) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - match VariantSlots::NAMES.get(idx) { - Some(name) => (*name).to_string(), - None => vortex_panic!("VariantArray slot_name index {idx} out of bounds"), - } - } - - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { - Ok(ExecutionResult::done(array)) - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - RULES.evaluate(array, parent, child_idx) + ))? + .into_array()) } } diff --git a/vortex-array/src/arrays/widen/vtable/mod.rs b/vortex-array/src/arrays/widen/vtable/mod.rs index 6d8dec5ea98..739618c8faf 100644 --- a/vortex-array/src/arrays/widen/vtable/mod.rs +++ b/vortex-array/src/arrays/widen/vtable/mod.rs @@ -14,11 +14,9 @@ pub use plugin::WidenPlugin; mod validity; use vortex_error::VortexResult; -use vortex_error::vortex_bail; 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; @@ -47,7 +45,6 @@ use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::match_each_decimal_value_type; use crate::require_child; -use crate::serde::ArrayChildren; /// The encoding of a [`WidenArray`]. #[derive(Clone, Debug)] @@ -103,25 +100,6 @@ impl VTable for Widen { with_empty_buffers(self, array, buffers) } - 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!("Widen deserialization is not supported") - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { WidenSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/serde.rs b/vortex-array/src/serde.rs index fb1a430f822..3b94498027f 100644 --- a/vortex-array/src/serde.rs +++ b/vortex-array/src/serde.rs @@ -789,7 +789,6 @@ mod tests { use vortex_session::registry::CachedId; use super::*; - use crate::Array; use crate::ArrayPlugin; use crate::ArraySerialization; use crate::ArrayVTable; @@ -868,16 +867,13 @@ mod tests { "old primitive wire ID cannot represent length {}", parts.len, ); - Ok(Array::::try_from_parts(ArrayVTable::deserialize( - &Primitive, - parts.dtype, - parts.len, - parts.metadata, - parts.buffers, - parts.children, + Primitive.deserialize( + ArrayDeserialization { + serialized_id: ArrayPlugin::id(&Primitive), + ..parts + }, session, - )?)? - .into_array()) + ) } } diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index 62c500f459d..85c61191de7 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -35,7 +35,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::RunEnd; use crate::arrays::Struct; @@ -113,7 +112,6 @@ impl Default for ArraySession { this.register(Dict); this.register(List); this.register(Masked); - this.register(PiecewiseSequence); this.register(RunEnd); this.register(VarBin); this.register(WidenPlugin); diff --git a/vortex-array/src/test_harness/trace/tests.rs b/vortex-array/src/test_harness/trace/tests.rs index ea452860b84..e8657e3df31 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; @@ -178,25 +177,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>, @@ -287,25 +267,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-layout/src/arrays/shared/vtable.rs b/vortex-layout/src/arrays/shared/vtable.rs index 435de3b4271..4ffbecffdc8 100644 --- a/vortex-layout/src/arrays/shared/vtable.rs +++ b/vortex-layout/src/arrays/shared/vtable.rs @@ -27,7 +27,6 @@ use vortex_array::with_empty_buffers; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::arrays::shared::SharedArrayExt; @@ -100,26 +99,6 @@ impl VTable for Shared { SharedSlots::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 vortex_array::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-python/src/arrays/py/vtable.rs b/vortex-python/src/arrays/py/vtable.rs index 333a4b0f1a9..17e599fa2ab 100644 --- a/vortex-python/src/arrays/py/vtable.rs +++ b/vortex-python/src/arrays/py/vtable.rs @@ -5,11 +5,14 @@ use std::hash::Hash; use std::sync::Arc; use vortex::array::Array; +use vortex::array::ArrayDeserialization; use vortex::array::ArrayEq; use vortex::array::ArrayHash; use vortex::array::ArrayId; use vortex::array::ArrayParts; +use vortex::array::ArrayPlugin; use vortex::array::ArrayRef; +use vortex::array::ArraySerialization; use vortex::array::ArrayView; use vortex::array::EqMode; use vortex::array::ExecutionCtx; @@ -19,7 +22,6 @@ use vortex::array::ProbeState; 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; @@ -93,33 +95,36 @@ impl VTable for PythonVTable { with_empty_buffers(self, array, buffers) } + fn slot_name(_array: ArrayView<'_, Self>, _idx: usize) -> String { + vortex_panic!("PythonArray has no slots") + } + + fn execute(_array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + vortex_bail!("PythonArray execution is not supported"); + } +} + +/// Python encodings register their IDs with the session, but have no wire form. +impl ArrayPlugin for PythonVTable { + fn id(&self) -> ArrayId { + self.id + } + fn serialize( - _array: ArrayView<'_, Self>, + &self, + _array: &ArrayRef, _session: &VortexSession, - ) -> VortexResult>> { + ) -> VortexResult> { Ok(None) } fn deserialize( &self, - _dtype: &DType, - _len: usize, - bytes: &[u8], - _buffers: &[BufferHandle], - _children: &dyn ArrayChildren, + _parts: ArrayDeserialization<'_>, _session: &VortexSession, - ) -> VortexResult> { - _ = bytes; + ) -> VortexResult { vortex_bail!("PythonArray deserialization is not supported"); } - - fn slot_name(_array: ArrayView<'_, Self>, _idx: usize) -> String { - vortex_panic!("PythonArray has no slots") - } - - fn execute(_array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { - vortex_bail!("PythonArray execution is not supported"); - } } impl OperationsVTable for PythonVTable {