From a18ed1b795a63cad9645f645a684afddac396665 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 5 Nov 2026 14:19:12 -0500 Subject: [PATCH 1/8] move array serde from VTable to explicit ArrayPlugin impls Signed-off-by: Matt Katz --- encodings/alp/src/alp/array.rs | 79 +------ encodings/alp/src/alp/array/plugin.rs | 135 ++++++++++++ encodings/alp/src/alp/mod.rs | 24 ++- encodings/alp/src/alp/plugin.rs | 14 +- encodings/alp/src/alp_rd/array.rs | 118 +---------- encodings/alp/src/alp_rd/array/plugin.rs | 177 ++++++++++++++++ encodings/bytebool/src/array.rs | 51 +---- encodings/bytebool/src/array/plugin.rs | 103 +++++++++ encodings/datetime-parts/src/array.rs | 64 +----- encodings/datetime-parts/src/array/plugin.rs | 119 +++++++++++ .../src/decimal_byte_parts/array.rs | 21 -- .../src/decimal_byte_parts/plugin/tests.rs | 32 --- encodings/elias-fano/src/unary/tests.rs | 37 ++-- encodings/elias-fano/src/unary/vtable/mod.rs | 44 +--- .../elias-fano/src/unary/vtable/plugin.rs | 98 +++++++++ .../fastlanes/src/bitpacking/plugin/tests.rs | 10 - .../fastlanes/src/bitpacking/vtable/mod.rs | 20 -- encodings/fastlanes/src/delta/vtable/mod.rs | 75 +------ .../fastlanes/src/delta/vtable/plugin.rs | 133 ++++++++++++ encodings/fastlanes/src/for_/plugin/tests.rs | 12 -- encodings/fastlanes/src/for_/vtable/mod.rs | 20 -- encodings/fastlanes/src/rle/vtable/mod.rs | 75 +------ encodings/fastlanes/src/rle/vtable/plugin.rs | 131 ++++++++++++ encodings/fsst/src/array.rs | 132 +----------- encodings/fsst/src/array/plugin.rs | 197 ++++++++++++++++++ encodings/onpair/src/array.rs | 104 +-------- encodings/onpair/src/array/plugin.rs | 163 +++++++++++++++ encodings/parquet-variant/src/vtable.rs | 104 +-------- .../parquet-variant/src/vtable/plugin.rs | 162 ++++++++++++++ encodings/pco/src/array.rs | 70 +------ encodings/pco/src/array/plugin.rs | 127 +++++++++++ encodings/sequence/src/array.rs | 86 +------- encodings/sequence/src/array/plugin.rs | 145 +++++++++++++ encodings/sparse/src/lib.rs | 64 +----- encodings/sparse/src/plugin.rs | 123 +++++++++++ encodings/zigzag/src/array.rs | 45 +--- encodings/zigzag/src/array/plugin.rs | 98 +++++++++ encodings/zstd/src/array.rs | 62 +----- encodings/zstd/src/array/plugin.rs | 118 +++++++++++ encodings/zstd/src/zstd_buffers.rs | 127 ++--------- encodings/zstd/src/zstd_buffers/plugin.rs | 154 ++++++++++++++ vortex-array/src/array/foreign.rs | 31 --- vortex-array/src/array/plugin.rs | 54 ----- vortex-array/src/array/vtable/mod.rs | 39 +--- vortex-array/src/arrays/bool/vtable/mod.rs | 54 +---- vortex-array/src/arrays/bool/vtable/plugin.rs | 108 ++++++++++ vortex-array/src/arrays/chunked/vtable/mod.rs | 77 +------ .../src/arrays/chunked/vtable/plugin.rs | 133 ++++++++++++ .../src/arrays/constant/vtable/mod.rs | 39 +--- .../src/arrays/constant/vtable/plugin.rs | 90 ++++++++ vortex-array/src/arrays/decimal/vtable/mod.rs | 20 -- vortex-array/src/arrays/dict/array.rs | 18 +- vortex-array/src/arrays/dict/vtable/mod.rs | 67 +----- vortex-array/src/arrays/dict/vtable/plugin.rs | 120 +++++++++++ .../src/arrays/extension/vtable/mod.rs | 41 +--- .../src/arrays/extension/vtable/plugin.rs | 94 +++++++++ vortex-array/src/arrays/filter/vtable.rs | 24 --- .../src/arrays/fixed_size_list/vtable/mod.rs | 64 +----- .../arrays/fixed_size_list/vtable/plugin.rs | 118 +++++++++++ vortex-array/src/arrays/interleave/mod.rs | 20 -- vortex-array/src/arrays/list/vtable/mod.rs | 70 +------ vortex-array/src/arrays/list/vtable/plugin.rs | 126 +++++++++++ .../src/arrays/listview/vtable/mod.rs | 101 +-------- .../src/arrays/listview/vtable/plugin.rs | 150 +++++++++++++ vortex-array/src/arrays/map/vtable/mod.rs | 52 +---- vortex-array/src/arrays/map/vtable/plugin.rs | 107 ++++++++++ vortex-array/src/arrays/masked/vtable/mod.rs | 63 +----- .../src/arrays/masked/vtable/plugin.rs | 116 +++++++++++ vortex-array/src/arrays/null/mod.rs | 34 +-- vortex-array/src/arrays/null/plugin.rs | 84 ++++++++ vortex-array/src/arrays/patched/vtable/mod.rs | 57 +---- .../src/arrays/patched/vtable/plugin.rs | 113 ++++++++++ .../src/arrays/piecewise_sequence/tests.rs | 2 +- .../src/arrays/piecewise_sequence/vtable.rs | 21 -- .../src/arrays/primitive/vtable/mod.rs | 61 +----- .../src/arrays/primitive/vtable/plugin.rs | 114 ++++++++++ vortex-array/src/arrays/runend/array.rs | 53 +---- .../src/arrays/runend/array/plugin.rs | 113 ++++++++++ .../src/arrays/scalar_fn/vtable/mod.rs | 23 -- .../src/arrays/scalar_fn/vtable/validity.rs | 22 -- vortex-array/src/arrays/slice/vtable.rs | 23 -- vortex-array/src/arrays/struct_/vtable/mod.rs | 67 +----- .../src/arrays/struct_/vtable/plugin.rs | 119 +++++++++++ vortex-array/src/arrays/union/vtable/mod.rs | 48 +---- .../src/arrays/union/vtable/plugin.rs | 99 +++++++++ vortex-array/src/arrays/varbin/vtable/mod.rs | 66 +----- .../src/arrays/varbin/vtable/plugin.rs | 125 +++++++++++ .../src/arrays/varbinview/vtable/mod.rs | 97 +-------- .../src/arrays/varbinview/vtable/plugin.rs | 152 ++++++++++++++ vortex-array/src/arrays/variant/vtable/mod.rs | 56 +---- .../src/arrays/variant/vtable/plugin.rs | 112 ++++++++++ vortex-array/src/arrays/widen/vtable/mod.rs | 22 -- vortex-array/src/serde.rs | 19 +- vortex-array/src/session/mod.rs | 2 - vortex-array/src/test_harness/trace/tests.rs | 39 ---- vortex-file/src/tests.rs | 4 +- vortex-layout/src/arrays/shared/vtable.rs | 21 -- vortex-python/src/arrays/py/vtable.rs | 24 +-- vortex-python/src/arrays/py/vtable/plugin.rs | 46 ++++ 99 files changed, 4555 insertions(+), 2952 deletions(-) create mode 100644 encodings/alp/src/alp/array/plugin.rs create mode 100644 encodings/alp/src/alp_rd/array/plugin.rs create mode 100644 encodings/bytebool/src/array/plugin.rs create mode 100644 encodings/datetime-parts/src/array/plugin.rs create mode 100644 encodings/elias-fano/src/unary/vtable/plugin.rs create mode 100644 encodings/fastlanes/src/delta/vtable/plugin.rs create mode 100644 encodings/fastlanes/src/rle/vtable/plugin.rs create mode 100644 encodings/fsst/src/array/plugin.rs create mode 100644 encodings/onpair/src/array/plugin.rs create mode 100644 encodings/parquet-variant/src/vtable/plugin.rs create mode 100644 encodings/pco/src/array/plugin.rs create mode 100644 encodings/sequence/src/array/plugin.rs create mode 100644 encodings/sparse/src/plugin.rs create mode 100644 encodings/zigzag/src/array/plugin.rs create mode 100644 encodings/zstd/src/array/plugin.rs create mode 100644 encodings/zstd/src/zstd_buffers/plugin.rs create mode 100644 vortex-array/src/arrays/bool/vtable/plugin.rs create mode 100644 vortex-array/src/arrays/chunked/vtable/plugin.rs create mode 100644 vortex-array/src/arrays/constant/vtable/plugin.rs create mode 100644 vortex-array/src/arrays/dict/vtable/plugin.rs create mode 100644 vortex-array/src/arrays/extension/vtable/plugin.rs create mode 100644 vortex-array/src/arrays/fixed_size_list/vtable/plugin.rs create mode 100644 vortex-array/src/arrays/list/vtable/plugin.rs create mode 100644 vortex-array/src/arrays/listview/vtable/plugin.rs create mode 100644 vortex-array/src/arrays/map/vtable/plugin.rs create mode 100644 vortex-array/src/arrays/masked/vtable/plugin.rs create mode 100644 vortex-array/src/arrays/null/plugin.rs create mode 100644 vortex-array/src/arrays/patched/vtable/plugin.rs create mode 100644 vortex-array/src/arrays/primitive/vtable/plugin.rs create mode 100644 vortex-array/src/arrays/runend/array/plugin.rs create mode 100644 vortex-array/src/arrays/struct_/vtable/plugin.rs create mode 100644 vortex-array/src/arrays/union/vtable/plugin.rs create mode 100644 vortex-array/src/arrays/varbin/vtable/plugin.rs create mode 100644 vortex-array/src/arrays/varbinview/vtable/plugin.rs create mode 100644 vortex-array/src/arrays/variant/vtable/plugin.rs create mode 100644 vortex-python/src/arrays/py/vtable/plugin.rs diff --git a/encodings/alp/src/alp/array.rs b/encodings/alp/src/alp/array.rs index 3eae9b59580..48967c719f4 100644 --- a/encodings/alp/src/alp/array.rs +++ b/encodings/alp/src/alp/array.rs @@ -7,7 +7,6 @@ use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; -use prost::Message; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -21,7 +20,6 @@ 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; use vortex_array::arrays::Primitive; use vortex_array::buffer::BufferHandle; @@ -34,7 +32,6 @@ use vortex_array::patches::PatchesData; use vortex_array::patches::PatchesMetadata; use vortex_array::require_child; use vortex_array::require_patches; -use vortex_array::serde::ArrayChildren; use vortex_array::smallvec::smallvec; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityChild; @@ -45,7 +42,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; use crate::ALPFloat; @@ -53,6 +49,8 @@ use crate::alp::Exponents; use crate::alp::decompress::execute_decompress; use crate::alp::rules::RULES; +mod plugin; + /// A [`ALP`]-encoded Vortex array. pub type ALPArray = Array; @@ -113,79 +111,6 @@ impl VTable for ALP { vortex_array::vtable::with_empty_buffers(self, array, buffers) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - let exponents = array.exponents(); - Ok(Some( - ALPMetadata { - exp_e: exponents.e as u32, - exp_f: exponents.f as u32, - patches: array - .patches() - .map(|p| p.to_metadata(array.len(), array.dtype())) - .transpose()?, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - let metadata = ALPMetadata::decode(metadata)?; - let encoded_ptype = match &dtype { - DType::Primitive(PType::F32, n) => DType::Primitive(PType::I32, *n), - DType::Primitive(PType::F64, n) => DType::Primitive(PType::I64, *n), - d => vortex_bail!("expected type: f32 or f64 but instead got {}", d), - }; - let encoded = children.get(0, &encoded_ptype, len)?; - - let patches = metadata - .patches - .map(|p| { - let indices = children.get(1, &p.indices_dtype()?, p.len()?)?; - let values = children.get(2, dtype, p.len()?)?; - let chunk_offsets = p - .chunk_offsets_dtype()? - .map(|dtype| children.get(3, &dtype, usize::try_from(p.chunk_offsets_len())?)) - .transpose()?; - - Patches::from_metadata( - len, - &p, - indices, - values, - chunk_offsets, - &mut session.create_execution_ctx(), - ) - }) - .transpose()?; - - let slots = ALPData::make_slots(&encoded, patches.as_ref()); - let data = ALPData::new( - Exponents { - e: u8::try_from(metadata.exp_e)?, - f: u8::try_from(metadata.exp_f)?, - }, - patches, - ); - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { ALPSlots::NAMES[idx].to_string() } diff --git a/encodings/alp/src/alp/array/plugin.rs b/encodings/alp/src/alp/array/plugin.rs new file mode 100644 index 00000000000..58462cf7475 --- /dev/null +++ b/encodings/alp/src/alp/array/plugin.rs @@ -0,0 +1,135 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`ALP`]. + +use prost::Message; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::VortexSessionExecute; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::patches::Patches; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use super::ALP; +use super::ALPArrayExt; +use super::ALPData; +use super::ALPMetadata; +use crate::alp::Exponents; + +impl ArrayPlugin for ALP { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "ALP plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + let exponents = array.exponents(); + let metadata = 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(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "ALP plugin does not recognize the serialized ID", + ); + 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), + DType::Primitive(PType::F64, n) => DType::Primitive(PType::I64, *n), + d => vortex_bail!("expected type: f32 or f64 but instead got {}", d), + }; + let encoded = children.get(0, &encoded_ptype, len)?; + + let patches = metadata + .patches + .map(|p| { + let indices = children.get(1, &p.indices_dtype()?, p.len()?)?; + let values = children.get(2, dtype, p.len()?)?; + let chunk_offsets = p + .chunk_offsets_dtype()? + .map(|dtype| children.get(3, &dtype, usize::try_from(p.chunk_offsets_len())?)) + .transpose()?; + + Patches::from_metadata( + len, + &p, + indices, + values, + chunk_offsets, + &mut session.create_execution_ctx(), + ) + }) + .transpose()?; + + let slots = ALPData::make_slots(&encoded, patches.as_ref()); + let data = ALPData::new( + Exponents { + e: u8::try_from(metadata.exp_e)?, + f: u8::try_from(metadata.exp_f)?, + }, + patches, + ); + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots, + ))? + .into_array()) + } +} diff --git a/encodings/alp/src/alp/mod.rs b/encodings/alp/src/alp/mod.rs index 66c30bba019..ea6f03ad038 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,20 @@ 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 = + ArrayPlugin::serialize(&ALP, &sliced, &session)?.expect("ALP serializes"); + let read = ArrayPlugin::deserialize( + &ALP, + ArrayDeserialization::new( + serialization.serialized_id, + sliced.dtype(), + sliced.len(), + &serialization.metadata, + &[], + &serialization.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..dca212568e9 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; @@ -65,16 +64,9 @@ 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 = ArrayPlugin::deserialize(&ALP, parts, session)? + .try_downcast::() + .map_err(|_| vortex_err!("ALP plugin should only deserialize vortex.alp"))?; // Check if there are interior patches to externalize. let Some(patches) = alp_array.patches() else { diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array.rs index e6eeff1d1fa..47df00cce0d 100644 --- a/encodings/alp/src/alp_rd/array.rs +++ b/encodings/alp/src/alp_rd/array.rs @@ -8,8 +8,6 @@ use std::hash::Hash; use std::hash::Hasher; use ::alp::MAX_DICT_SIZE; -use itertools::Itertools; -use prost::Message; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -23,7 +21,6 @@ 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; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; @@ -37,7 +34,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; @@ -51,13 +47,14 @@ 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::alp_rd::check_left_parts_codes; use crate::alp_rd::rules::RULES; use crate::alp_rd_decode; +mod plugin; + /// A [`ALPRD`]-encoded Vortex array. pub type ALPRDArray = Array; @@ -152,117 +149,6 @@ impl VTable for ALPRD { vortex_array::vtable::with_empty_buffers(self, array, buffers) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - let dict = array - .left_parts_dictionary() - .iter() - .map(|&i| i as u32) - .collect::>(); - - Ok(Some( - ALPRDMetadata { - right_bit_width: array.right_bit_width() as u32, - dict_len: array.left_parts_dictionary().len() as u32, - dict, - left_parts_ptype: array.left_parts().dtype().as_ptype() as i32, - patches: array - .left_parts_patches() - .map(|p| p.to_metadata(array.len(), p.dtype())) - .transpose()?, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - let metadata = ALPRDMetadata::decode(metadata)?; - let left_parts_ptype = metadata.get_left_parts_ptype()?; - if children.len() < 2 { - vortex_bail!( - "Expected at least 2 children for ALPRD encoding, found {}", - children.len() - ); - } - - // Stored metadata is untrusted: dict_len must stay within the stored entries. - if metadata.dict_len as usize > metadata.dict.len() { - vortex_bail!( - "ALP-RD dictionary length {} exceeds the {} stored entries", - metadata.dict_len, - metadata.dict.len() - ); - } - - let left_parts_dtype = DType::Primitive(left_parts_ptype, dtype.nullability()); - let left_parts = children.get(0, &left_parts_dtype, len)?; - let left_parts_dictionary: Buffer = metadata.dict.as_slice() - [0..metadata.dict_len as usize] - .iter() - .map(|&i| { - u16::try_from(i) - .map_err(|_| vortex_err!("left_parts_dictionary code {i} does not fit in u16")) - }) - .try_collect()?; - - let right_parts_dtype = match &dtype { - DType::Primitive(PType::F32, _) => { - DType::Primitive(PType::U32, Nullability::NonNullable) - } - DType::Primitive(PType::F64, _) => { - DType::Primitive(PType::U64, Nullability::NonNullable) - } - _ => vortex_bail!("Expected f32 or f64 dtype, got {:?}", dtype), - }; - let right_parts = children.get(1, &right_parts_dtype, len)?; - - let left_parts_patches = metadata - .patches - .map(|p| { - let indices = children.get(2, &p.indices_dtype()?, p.len()?)?; - let values = children.get(3, &left_parts_dtype.as_nonnullable(), p.len()?)?; - - Patches::new( - len, - p.offset()?, - indices, - values, - // TODO(0ax1): handle chunk offsets - None, - &mut session.create_execution_ctx(), - ) - }) - .transpose()?; - let slots = ALPRDData::make_slots(&left_parts, &right_parts, left_parts_patches.as_ref()); - let data = ALPRDData::new( - left_parts_dictionary, - u8::try_from(metadata.right_bit_width).map_err(|_| { - vortex_err!( - "right_bit_width {} out of u8 range", - metadata.right_bit_width - ) - })?, - left_parts_patches, - ); - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { ALPRDSlots::NAMES[idx].to_string() } diff --git a/encodings/alp/src/alp_rd/array/plugin.rs b/encodings/alp/src/alp_rd/array/plugin.rs new file mode 100644 index 00000000000..00e83ed39ac --- /dev/null +++ b/encodings/alp/src/alp_rd/array/plugin.rs @@ -0,0 +1,177 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`ALPRD`]. + +use itertools::Itertools; +use prost::Message; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::VortexSessionExecute; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::patches::Patches; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use super::ALPRD; +use super::ALPRDArrayExt; +use super::ALPRDArraySlotsExt; +use super::ALPRDData; +use super::ALPRDMetadata; + +impl ArrayPlugin for ALPRD { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "ALPRD plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + let dict = array + .left_parts_dictionary() + .iter() + .map(|&i| i as u32) + .collect::>(); + + let metadata = 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(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "ALPRD plugin does not recognize the serialized ID", + ); + 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 { + vortex_bail!( + "Expected at least 2 children for ALPRD encoding, found {}", + children.len() + ); + } + + // Stored metadata is untrusted: dict_len must stay within the stored entries. + if metadata.dict_len as usize > metadata.dict.len() { + vortex_bail!( + "ALP-RD dictionary length {} exceeds the {} stored entries", + metadata.dict_len, + metadata.dict.len() + ); + } + + let left_parts_dtype = DType::Primitive(left_parts_ptype, dtype.nullability()); + let left_parts = children.get(0, &left_parts_dtype, len)?; + let left_parts_dictionary: Buffer = metadata.dict.as_slice() + [0..metadata.dict_len as usize] + .iter() + .map(|&i| { + u16::try_from(i) + .map_err(|_| vortex_err!("left_parts_dictionary code {i} does not fit in u16")) + }) + .try_collect()?; + + let right_parts_dtype = match &dtype { + DType::Primitive(PType::F32, _) => { + DType::Primitive(PType::U32, Nullability::NonNullable) + } + DType::Primitive(PType::F64, _) => { + DType::Primitive(PType::U64, Nullability::NonNullable) + } + _ => vortex_bail!("Expected f32 or f64 dtype, got {:?}", dtype), + }; + let right_parts = children.get(1, &right_parts_dtype, len)?; + + let left_parts_patches = metadata + .patches + .map(|p| { + let indices = children.get(2, &p.indices_dtype()?, p.len()?)?; + let values = children.get(3, &left_parts_dtype.as_nonnullable(), p.len()?)?; + + Patches::new( + len, + p.offset()?, + indices, + values, + // TODO(0ax1): handle chunk offsets + None, + &mut session.create_execution_ctx(), + ) + }) + .transpose()?; + let slots = ALPRDData::make_slots(&left_parts, &right_parts, left_parts_patches.as_ref()); + let data = ALPRDData::new( + left_parts_dictionary, + u8::try_from(metadata.right_bit_width).map_err(|_| { + vortex_err!( + "right_bit_width {} out of u8 range", + metadata.right_bit_width + ) + })?, + left_parts_patches, + ); + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots, + ))? + .into_array()) + } +} diff --git a/encodings/bytebool/src/array.rs b/encodings/bytebool/src/array.rs index 88007ecfd01..4356712d944 100644 --- a/encodings/bytebool/src/array.rs +++ b/encodings/bytebool/src/array.rs @@ -25,7 +25,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; @@ -35,12 +34,13 @@ use vortex_array::vtable::validity_to_child; use vortex_buffer::BitBufferMut; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +mod plugin; + /// A [`ByteBool`]-encoded Vortex array. pub type ByteBoolArray = Array; @@ -113,53 +113,6 @@ impl VTable for ByteBool { )) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - if !metadata.is_empty() { - vortex_bail!( - "ByteBoolArray expects empty metadata, got {} bytes", - metadata.len() - ); - } - let validity = if children.is_empty() { - Validity::from(dtype.nullability()) - } else if children.len() == 1 { - let validity = children.get(0, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("Expected 0 or 1 child, got {}", children.len()); - }; - - if buffers.len() != 1 { - vortex_bail!("Expected 1 buffer, got {}", buffers.len()); - } - let buffer = buffers[0].clone(); - - let data = ByteBoolData::new(buffer); - let slots = ByteBoolData::make_slots(&validity, len); - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { ByteBoolSlots::NAMES[idx].to_string() } diff --git a/encodings/bytebool/src/array/plugin.rs b/encodings/bytebool/src/array/plugin.rs new file mode 100644 index 00000000000..85da7f8632e --- /dev/null +++ b/encodings/bytebool/src/array/plugin.rs @@ -0,0 +1,103 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`ByteBool`]. + +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use super::ByteBool; +use super::ByteBoolData; + +impl ArrayPlugin for ByteBool { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "ByteBool plugin cannot serialize an array of another encoding", + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "ByteBool plugin does not recognize the serialized ID", + ); + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + if !metadata.is_empty() { + vortex_bail!( + "ByteBoolArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + let validity = if children.is_empty() { + Validity::from(dtype.nullability()) + } else if children.len() == 1 { + let validity = children.get(0, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + vortex_bail!("Expected 0 or 1 child, got {}", children.len()); + }; + + if buffers.len() != 1 { + vortex_bail!("Expected 1 buffer, got {}", buffers.len()); + } + let buffer = buffers[0].clone(); + + let data = ByteBoolData::new(buffer); + let slots = ByteBoolData::make_slots(&validity, len); + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots, + ))? + .into_array()) + } +} diff --git a/encodings/datetime-parts/src/array.rs b/encodings/datetime-parts/src/array.rs index 4416327acaf..706a856088e 100644 --- a/encodings/datetime-parts/src/array.rs +++ b/encodings/datetime-parts/src/array.rs @@ -6,7 +6,6 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hasher; -use prost::Message; use vortex_array::AnyCanonical; use vortex_array::Array; use vortex_array::ArrayEq; @@ -24,10 +23,8 @@ use vortex_array::arrays::Primitive; use vortex_array::arrays::TemporalArray; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::require_child; -use vortex_array::serde::ArrayChildren; use vortex_array::smallvec::smallvec; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityChild; @@ -37,7 +34,6 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::TemporalParts; @@ -45,6 +41,8 @@ use crate::canonical::decode_to_temporal; use crate::compute::rules::PARENT_RULES; use crate::split_temporal; +mod plugin; + /// A [`DateTimeParts`]-encoded Vortex array. pub type DateTimePartsArray = Array; @@ -130,64 +128,6 @@ impl VTable for DateTimeParts { vortex_array::vtable::with_empty_buffers(self, array, buffers) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - DateTimePartsMetadata { - days_ptype: PType::try_from(array.days().dtype())? as i32, - seconds_ptype: PType::try_from(array.seconds().dtype())? as i32, - subseconds_ptype: PType::try_from(array.subseconds().dtype())? as i32, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - let metadata = DateTimePartsMetadata::decode(metadata)?; - if children.len() != 3 { - vortex_bail!( - "Expected 3 children for datetime-parts encoding, found {}", - children.len() - ) - } - - let days = children.get( - 0, - &DType::Primitive(metadata.get_days_ptype()?, dtype.nullability()), - len, - )?; - let seconds = children.get( - 1, - &DType::Primitive(metadata.get_seconds_ptype()?, Nullability::NonNullable), - len, - )?; - let subseconds = children.get( - 2, - &DType::Primitive(metadata.get_subseconds_ptype()?, Nullability::NonNullable), - len, - )?; - - let slots = smallvec![Some(days), Some(seconds), Some(subseconds)]; - let data = DateTimePartsData {}; - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { DateTimePartsSlots::NAMES[idx].to_string() } diff --git a/encodings/datetime-parts/src/array/plugin.rs b/encodings/datetime-parts/src/array/plugin.rs new file mode 100644 index 00000000000..47d6cab44e0 --- /dev/null +++ b/encodings/datetime-parts/src/array/plugin.rs @@ -0,0 +1,119 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`DateTimeParts`]. + +use prost::Message; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::smallvec::smallvec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use super::DateTimeParts; +use super::DateTimePartsArraySlotsExt; +use super::DateTimePartsData; +use super::DateTimePartsMetadata; + +impl ArrayPlugin for DateTimeParts { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "DateTimeParts plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + let metadata = 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(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "DateTimeParts plugin does not recognize the serialized ID", + ); + let ArrayDeserialization { + dtype, + len, + metadata, + children, + .. + } = parts; + let metadata = DateTimePartsMetadata::decode(metadata)?; + if children.len() != 3 { + vortex_bail!( + "Expected 3 children for datetime-parts encoding, found {}", + children.len() + ) + } + + let days = children.get( + 0, + &DType::Primitive(metadata.get_days_ptype()?, dtype.nullability()), + len, + )?; + let seconds = children.get( + 1, + &DType::Primitive(metadata.get_seconds_ptype()?, Nullability::NonNullable), + len, + )?; + let subseconds = children.get( + 2, + &DType::Primitive(metadata.get_subseconds_ptype()?, Nullability::NonNullable), + len, + )?; + + let slots = smallvec![Some(days), Some(seconds), Some(subseconds)]; + let data = DateTimePartsData {}; + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots, + ))? + .into_array()) + } +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs index 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/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..993687eec07 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,19 @@ 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. 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) - } + ArrayPlugin::deserialize( + &Unary, + 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..edf773f3b7e 100644 --- a/encodings/elias-fano/src/unary/vtable/mod.rs +++ b/encodings/elias-fano/src/unary/vtable/mod.rs @@ -5,7 +5,6 @@ use std::hash::Hash; use std::hash::Hasher; use num_traits::AsPrimitive; -use prost::Message; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -24,7 +23,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; @@ -32,12 +30,12 @@ use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use super::array::UnaryData; mod operations; +mod plugin; mod validity; /// A [`Unary`]-encoded Vortex array. @@ -134,46 +132,6 @@ impl VTable for Unary { )) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - UnaryMetadata { - num_ones: num_ones_of(array)?, - offset: u32::try_from(array.meta.offset())?, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_ensure_eq!(buffers.len(), 1); - vortex_ensure_eq!(children.len(), 0); - let metadata = UnaryMetadata::decode(metadata)?; - let data = UnaryData::from_buffer( - buffers[0].clone(), - metadata.num_ones, - usize::try_from(metadata.offset)?, - len, - )?; - Ok(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}") diff --git a/encodings/elias-fano/src/unary/vtable/plugin.rs b/encodings/elias-fano/src/unary/vtable/plugin.rs new file mode 100644 index 00000000000..222cc8b78c0 --- /dev/null +++ b/encodings/elias-fano/src/unary/vtable/plugin.rs @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Unary`]. + +use prost::Message; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::ArraySlots; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use super::Unary; +use super::UnaryMetadata; +use super::num_ones_of; +use crate::UnaryData; + +impl ArrayPlugin for Unary { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Unary plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + let metadata = UnaryMetadata { + num_ones: num_ones_of(array)?, + offset: u32::try_from(array.meta.offset())?, + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Unary plugin does not recognize the serialized ID", + ); + 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)?; + let data = UnaryData::from_buffer( + buffers[0].clone(), + metadata.num_ones, + usize::try_from(metadata.offset)?, + len, + )?; + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + ArraySlots::new(), + ))? + .into_array()) + } +} 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..6fd2e6de2c6 100644 --- a/encodings/fastlanes/src/delta/vtable/mod.rs +++ b/encodings/fastlanes/src/delta/vtable/mod.rs @@ -4,7 +4,6 @@ use std::hash::Hash; use std::hash::Hasher; -use prost::Message; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -20,23 +19,18 @@ use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; 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; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; -use vortex_error::vortex_err; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::DeltaData; use crate::FL_CHUNK_SIZE; -use crate::delta::array::DeltaArrayExt; use crate::delta::array::DeltaArraySlotsExt; use crate::delta::array::DeltaSlots; use crate::delta::array::DeltaSlotsView; @@ -45,6 +39,7 @@ use crate::delta::array::lane_count; use crate::delta_compress; mod operations; +mod plugin; mod rules; mod slice; mod validity; @@ -134,74 +129,6 @@ impl VTable for Delta { DeltaSlots::NAMES[idx].to_string() } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - DeltaMetadata { - deltas_len: array.deltas().len() as u64, - offset: array.offset() as u32, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_ensure!( - buffers.is_empty(), - "DeltaArray expects 0 buffers, got {}", - buffers.len() - ); - vortex_ensure!( - children.len() == 2 || children.len() == 3, - "DeltaArray expects 2 or 3 children, got {}", - children.len() - ); - let metadata = DeltaMetadata::decode(metadata)?; - let ptype = PType::try_from(dtype)?; - let lanes = lane_count(ptype); - - let deltas_len = usize::try_from(metadata.deltas_len) - .map_err(|_| vortex_err!("deltas_len {} overflowed usize", metadata.deltas_len))?; - vortex_ensure!( - deltas_len.is_multiple_of(FL_CHUNK_SIZE), - "deltas length must be a multiple of {FL_CHUNK_SIZE}" - ); - let bases_len = deltas_len / FL_CHUNK_SIZE * lanes; - - let bases = children.get(0, &dtype.as_nonnullable(), bases_len)?; - let deltas = children.get(1, &dtype.as_nonnullable(), deltas_len)?; - let validity_child = if children.len() == 3 { - Some(children.get(2, &Validity::DTYPE, len)?) - } else { - None - }; - - let data = DeltaData::try_new(metadata.offset as usize)?; - let slots = DeltaSlots { - bases, - deltas, - validity_child, - } - .into_slots(); - Ok(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); diff --git a/encodings/fastlanes/src/delta/vtable/plugin.rs b/encodings/fastlanes/src/delta/vtable/plugin.rs new file mode 100644 index 00000000000..fa55d39a23e --- /dev/null +++ b/encodings/fastlanes/src/delta/vtable/plugin.rs @@ -0,0 +1,133 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Delta`]. + +use prost::Message; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::dtype::PType; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use super::Delta; +use super::DeltaMetadata; +use crate::DeltaData; +use crate::FL_CHUNK_SIZE; +use crate::delta::array::DeltaArrayExt; +use crate::delta::array::DeltaArraySlotsExt; +use crate::delta::array::DeltaSlots; +use crate::delta::array::lane_count; + +impl ArrayPlugin for Delta { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Delta plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + let metadata = DeltaMetadata { + deltas_len: array.deltas().len() as u64, + offset: array.offset() as u32, + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Delta plugin does not recognize the serialized ID", + ); + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + vortex_ensure!( + buffers.is_empty(), + "DeltaArray expects 0 buffers, got {}", + buffers.len() + ); + vortex_ensure!( + children.len() == 2 || children.len() == 3, + "DeltaArray expects 2 or 3 children, got {}", + children.len() + ); + let metadata = DeltaMetadata::decode(metadata)?; + let ptype = PType::try_from(dtype)?; + let lanes = lane_count(ptype); + + let deltas_len = usize::try_from(metadata.deltas_len) + .map_err(|_| vortex_err!("deltas_len {} overflowed usize", metadata.deltas_len))?; + vortex_ensure!( + deltas_len.is_multiple_of(FL_CHUNK_SIZE), + "deltas length must be a multiple of {FL_CHUNK_SIZE}" + ); + let bases_len = deltas_len / FL_CHUNK_SIZE * lanes; + + let bases = children.get(0, &dtype.as_nonnullable(), bases_len)?; + let deltas = children.get(1, &dtype.as_nonnullable(), deltas_len)?; + let validity_child = if children.len() == 3 { + Some(children.get(2, &Validity::DTYPE, len)?) + } else { + None + }; + + let data = DeltaData::try_new(metadata.offset as usize)?; + let slots = DeltaSlots { + bases, + deltas, + validity_child, + } + .into_slots(); + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots, + ))? + .into_array()) + } +} 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..e1fec5b5476 100644 --- a/encodings/fastlanes/src/rle/vtable/mod.rs +++ b/encodings/fastlanes/src/rle/vtable/mod.rs @@ -4,7 +4,6 @@ use std::hash::Hash; use std::hash::Hasher; -use prost::Message; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -21,25 +20,22 @@ use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::serde::ArrayChildren; use vortex_array::vtable::VTable; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::RLEData; -use crate::rle::array::RLEArrayExt; -use crate::rle::array::RLEArraySlotsExt; use crate::rle::array::RLESlots; use crate::rle::array::RLESlotsView; use crate::rle::array::rle_decompress::rle_decompress; use crate::rle::vtable::rules::RULES; mod operations; +mod plugin; mod rules; mod validity; @@ -155,75 +151,6 @@ impl VTable for RLE { RLESlots::NAMES[idx].to_string() } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - RLEMetadata { - values_len: array.values().len() as u64, - indices_len: array.indices().len() as u64, - indices_ptype: PType::try_from(array.indices().dtype())? as i32, - values_idx_offsets_len: array.values_idx_offsets().len() as u64, - values_idx_offsets_ptype: PType::try_from(array.values_idx_offsets().dtype())? - as i32, - offset: array.offset() as u64, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - let metadata = RLEMetadata::decode(metadata)?; - let indices_ptype = metadata.get_indices_ptype()?; - let values_idx_offsets_ptype = metadata.get_values_idx_offsets_ptype()?; - vortex_ensure!( - buffers.is_empty(), - "RLEArray expects 0 buffers, got {}", - buffers.len() - ); - let values = children.get( - 0, - &DType::Primitive(dtype.as_ptype(), Nullability::NonNullable), - usize::try_from(metadata.values_len)?, - )?; - - let indices = children.get( - 1, - &DType::Primitive(indices_ptype, dtype.nullability()), - usize::try_from(metadata.indices_len)?, - )?; - - let values_idx_offsets = children.get( - 2, - &DType::Primitive(values_idx_offsets_ptype, Nullability::NonNullable), - usize::try_from(metadata.values_idx_offsets_len)?, - )?; - - let slots = RLESlots { - values, - indices, - values_idx_offsets, - } - .into_slots(); - let data = RLEData::try_new(metadata.offset as usize)?; - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )) - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done( rle_decompress(&array, ctx)?.into_array(), diff --git a/encodings/fastlanes/src/rle/vtable/plugin.rs b/encodings/fastlanes/src/rle/vtable/plugin.rs new file mode 100644 index 00000000000..c3e104d8373 --- /dev/null +++ b/encodings/fastlanes/src/rle/vtable/plugin.rs @@ -0,0 +1,131 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`RLE`]. + +use prost::Message; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use super::RLE; +use super::RLEMetadata; +use crate::RLEData; +use crate::rle::array::RLEArrayExt; +use crate::rle::array::RLEArraySlotsExt; +use crate::rle::array::RLESlots; + +impl ArrayPlugin for RLE { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "RLE plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + let metadata = 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(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "RLE plugin does not recognize the serialized ID", + ); + 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()?; + vortex_ensure!( + buffers.is_empty(), + "RLEArray expects 0 buffers, got {}", + buffers.len() + ); + let values = children.get( + 0, + &DType::Primitive(dtype.as_ptype(), Nullability::NonNullable), + usize::try_from(metadata.values_len)?, + )?; + + let indices = children.get( + 1, + &DType::Primitive(indices_ptype, dtype.nullability()), + usize::try_from(metadata.indices_len)?, + )?; + + let values_idx_offsets = children.get( + 2, + &DType::Primitive(values_idx_offsets_ptype, Nullability::NonNullable), + usize::try_from(metadata.values_idx_offsets_len)?, + )?; + + let slots = RLESlots { + values, + indices, + values_idx_offsets, + } + .into_slots(); + let data = RLEData::try_new(metadata.offset as usize)?; + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots, + ))? + .into_array()) + } +} diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index ae8d909b052..c30ac6a5b02 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -13,7 +13,6 @@ use fsst::Compressor; use fsst::Decompressor; use fsst::Symbol; use num_traits::AsPrimitive; -use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -26,7 +25,6 @@ use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::TypedArrayRef; -use vortex_array::VortexSessionExecute; use vortex_array::array_slots; use vortex_array::arrays::Primitive; use vortex_array::arrays::VarBin; @@ -60,7 +58,6 @@ 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::canonical::FSST_DECODE_SLACK; @@ -69,6 +66,8 @@ use crate::canonical::canonicalize_fsst; use crate::canonical::fsst_decode_bytes; use crate::rules::RULES; +mod plugin; + /// A [`FSST`]-encoded Vortex array. pub type FSSTArray = Array; @@ -187,133 +186,6 @@ impl VTable for FSST { )) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - let codes_offsets = array.codes_offsets(); - Ok(Some( - FSSTMetadata { - uncompressed_lengths_ptype: array.uncompressed_lengths().dtype().as_ptype().into(), - codes_offsets_ptype: codes_offsets.dtype().as_ptype().into(), - } - .encode_to_vec(), - )) - } - - /// Deserializes an FSST array from its serialized components. - /// - /// Supports two serialization formats: - /// - /// ## Legacy format (2 buffers, 2 children) - /// - /// The original FSST layout stored the compressed codes as a full `VarBinArray` child. - /// - **Buffers**: `[symbols, symbol_lengths]` - /// - **Children**: `[codes (VarBinArray), uncompressed_lengths (Primitive)]` - /// - /// The codes VarBinArray child is decomposed: its bytes become the `codes_bytes` buffer, - /// and its offsets/validity are extracted into slots. - /// See `FSST::deserialize_legacy`. - /// - /// ## Current format (3 buffers, 2-3 children) - /// - /// The current layout stores the compressed bytes as a raw buffer alongside the symbol - /// table, with offsets and validity as separate children. - /// - **Buffers**: `[symbols, symbol_lengths, compressed_codes_bytes]` - /// - **Children**: `[uncompressed_lengths, codes_offsets, (optional) codes_validity]` - /// - /// The `codes_bytes` buffer is stored directly in `FSSTData`. A `VarBinArray` for the - /// codes can be reconstructed on demand via [`FSSTArrayExt::codes()`] using the bytes - /// from `FSSTData` combined with offsets and validity from the array's slots. - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - let metadata = FSSTMetadata::decode(metadata)?; - let symbols = Buffer::::from_byte_buffer(buffers[0].clone().try_to_host_sync()?); - let symbol_lengths = Buffer::::from_byte_buffer(buffers[1].clone().try_to_host_sync()?); - - let mut ctx = session.create_execution_ctx(); - if buffers.len() == 2 { - return Self::deserialize_legacy( - self, - dtype, - len, - &metadata, - &symbols, - &symbol_lengths, - children, - &mut ctx, - ); - } - - if buffers.len() == 3 { - let uncompressed_lengths = children.get( - 0, - &DType::Primitive( - metadata.get_uncompressed_lengths_ptype()?, - Nullability::NonNullable, - ), - len, - )?; - - let codes_bytes = buffers[2].clone(); - let codes_offsets = children.get( - 1, - &DType::Primitive( - PType::try_from(metadata.codes_offsets_ptype)?, - Nullability::NonNullable, - ), - // VarBin offsets are len + 1 - len + 1, - )?; - - let codes_validity = if children.len() == 2 { - Validity::from(dtype.nullability()) - } else if children.len() == 3 { - let validity = children.get(2, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("Expected 2 or 3 children, got {}", children.len()); - }; - - FSSTData::validate_parts( - symbols.as_slice(), - symbol_lengths.as_slice(), - &codes_offsets, - dtype.nullability(), - &uncompressed_lengths, - dtype, - len, - )?; - FSSTData::validate_codes_bounds(&codes_bytes, &codes_offsets, &mut ctx)?; - let slots = FSSTSlots { - uncompressed_lengths, - codes_offsets, - codes_validity: validity_to_child(&codes_validity, len), - } - .into_slots(); - let data = FSSTData::try_new(symbols, symbol_lengths, codes_bytes, len)?; - return Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )); - } - - vortex_bail!( - "InvalidArgument: Expected 2 or 3 buffers, got {}", - buffers.len() - ); - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { FSSTSlots::NAMES[idx].to_string() } diff --git a/encodings/fsst/src/array/plugin.rs b/encodings/fsst/src/array/plugin.rs new file mode 100644 index 00000000000..93d8a40df59 --- /dev/null +++ b/encodings/fsst/src/array/plugin.rs @@ -0,0 +1,197 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`FSST`]. + +use fsst::Symbol; +use prost::Message; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::VortexSessionExecute; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::validity::Validity; +use vortex_array::vtable::validity_to_child; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use crate::FSST; +use crate::array::FSSTArraySlotsExt; +use crate::array::FSSTData; +use crate::array::FSSTMetadata; +use crate::array::FSSTSlots; + +impl ArrayPlugin for FSST { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "FSST plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + let codes_offsets = array.codes_offsets(); + let metadata = FSSTMetadata { + uncompressed_lengths_ptype: array.uncompressed_lengths().dtype().as_ptype().into(), + codes_offsets_ptype: codes_offsets.dtype().as_ptype().into(), + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + /// Deserializes an FSST array from its serialized components. + /// + /// Supports two serialization formats: + /// + /// ## Legacy format (2 buffers, 2 children) + /// + /// The original FSST layout stored the compressed codes as a full `VarBinArray` child. + /// - **Buffers**: `[symbols, symbol_lengths]` + /// - **Children**: `[codes (VarBinArray), uncompressed_lengths (Primitive)]` + /// + /// The codes VarBinArray child is decomposed: its bytes become the `codes_bytes` buffer, + /// and its offsets/validity are extracted into slots. + /// See `FSST::deserialize_legacy`. + /// + /// ## Current format (3 buffers, 2-3 children) + /// + /// The current layout stores the compressed bytes as a raw buffer alongside the symbol + /// table, with offsets and validity as separate children. + /// - **Buffers**: `[symbols, symbol_lengths, compressed_codes_bytes]` + /// - **Children**: `[uncompressed_lengths, codes_offsets, (optional) codes_validity]` + /// + /// The `codes_bytes` buffer is stored directly in `FSSTData`. A `VarBinArray` for the + /// codes can be reconstructed on demand via [`FSSTArrayExt::codes()`] using the bytes + /// from `FSSTData` combined with offsets and validity from the array's slots. + /// + /// [`FSSTArrayExt::codes()`]: crate::array::FSSTArrayExt::codes + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "FSST plugin does not recognize the serialized ID", + ); + 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 Ok(Array::try_from_parts(Self::deserialize_legacy( + self, + dtype, + len, + &metadata, + &symbols, + &symbol_lengths, + children, + &mut ctx, + )?)? + .into_array()); + } + + if buffers.len() == 3 { + let uncompressed_lengths = children.get( + 0, + &DType::Primitive( + metadata.get_uncompressed_lengths_ptype()?, + Nullability::NonNullable, + ), + len, + )?; + + let codes_bytes = buffers[2].clone(); + let codes_offsets = children.get( + 1, + &DType::Primitive( + PType::try_from(metadata.codes_offsets_ptype)?, + Nullability::NonNullable, + ), + // VarBin offsets are len + 1 + len + 1, + )?; + + let codes_validity = if children.len() == 2 { + Validity::from(dtype.nullability()) + } else if children.len() == 3 { + let validity = children.get(2, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + vortex_bail!("Expected 2 or 3 children, got {}", children.len()); + }; + + FSSTData::validate_parts( + symbols.as_slice(), + symbol_lengths.as_slice(), + &codes_offsets, + dtype.nullability(), + &uncompressed_lengths, + dtype, + len, + )?; + FSSTData::validate_codes_bounds(&codes_bytes, &codes_offsets, &mut ctx)?; + let slots = FSSTSlots { + uncompressed_lengths, + codes_offsets, + codes_validity: validity_to_child(&codes_validity, len), + } + .into_slots(); + let data = FSSTData::try_new(symbols, symbol_lengths, codes_bytes, len)?; + return Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots, + ))? + .into_array()); + } + + vortex_bail!( + "InvalidArgument: Expected 2 or 3 buffers, got {}", + buffers.len() + ); + } +} diff --git a/encodings/onpair/src/array.rs b/encodings/onpair/src/array.rs index 8b45fc13438..344d76f6274 100644 --- a/encodings/onpair/src/array.rs +++ b/encodings/onpair/src/array.rs @@ -14,7 +14,6 @@ use onpair::CompactDictionary; use onpair::CompactDictionaryView; use onpair::Dictionary; use onpair::DictionaryStorage; -use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -32,14 +31,12 @@ use vortex_array::builders::ArrayBuilder; use vortex_array::builders::VarBinBuilder; use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; use vortex_array::dtype::OffsetBuilderPType; use vortex_array::dtype::PType; use vortex_array::match_each_integer_ptype; use vortex_array::match_each_varbin_builder; use vortex_array::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; @@ -53,7 +50,6 @@ 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::canonical::OnPairDecodePlan; @@ -62,6 +58,8 @@ use crate::canonical::onpair_decode_bytes; use crate::decode::collect_widened; use crate::rules::RULES; +mod plugin; + /// An [`OnPair`]-encoded Vortex array. pub type OnPairArray = Array; @@ -510,104 +508,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() } diff --git a/encodings/onpair/src/array/plugin.rs b/encodings/onpair/src/array/plugin.rs new file mode 100644 index 00000000000..79fcf92ac30 --- /dev/null +++ b/encodings/onpair/src/array/plugin.rs @@ -0,0 +1,163 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`OnPair`]. + +use prost::Message; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::validity::Validity; +use vortex_array::vtable::validity_to_child; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use crate::OnPair; +use crate::array::OnPairArraySlotsExt; +use crate::array::OnPairData; +use crate::array::OnPairMetadata; +use crate::array::OnPairSlots; + +impl ArrayPlugin for OnPair { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "OnPair plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + 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; + let metadata = 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(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "OnPair plugin does not recognize the serialized ID", + ); + 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()) + } +} diff --git a/encodings/parquet-variant/src/vtable.rs b/encodings/parquet-variant/src/vtable.rs index b7b8105a2ab..96b89a1ab33 100644 --- a/encodings/parquet-variant/src/vtable.rs +++ b/encodings/parquet-variant/src/vtable.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use prost::Message; use vortex_array::Array; use vortex_array::ArrayId; use vortex_array::ArrayParts; @@ -16,17 +15,14 @@ use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::proto::dtype as pb; -use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::child_to_validity; -use vortex_array::vtable::validity_to_child; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::array::ParquetVariantArraySlotsExt; @@ -34,6 +30,8 @@ use crate::array::ParquetVariantSlots; use crate::array::core_storage_without_typed_value; use crate::array::logical_shredded_from_parquet_typed_value; +mod plugin; + /// VTable for Arrow's canonical `arrow.parquet.variant` extension storage. /// /// `ParquetVariantArray` preserves semi-structured data stored as Parquet Variant values in a @@ -176,104 +174,6 @@ impl VTable for ParquetVariant { ParquetVariantSlots::NAMES[idx].to_string() } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - let typed_value_dtype = array - .typed_value() - .map(|tv| tv.dtype().try_into()) - .transpose()?; - Ok(Some( - ParquetVariantMetadataProto { - has_value: array.value().is_some(), - typed_value_dtype, - value_nullable: array.value().is_some_and(|v| v.dtype().is_nullable()), - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - vortex_ensure!( - buffers.is_empty(), - "ParquetVariantArray expects 0 buffers, got {}", - buffers.len() - ); - - let proto = ParquetVariantMetadataProto::decode(metadata)?; - let typed_value_dtype = match proto.typed_value_dtype.as_ref() { - Some(dtype) => Some(DType::from_proto(dtype, session)?), - None => None, - }; - - vortex_ensure!(matches!(dtype, DType::Variant(_)), "Expected Variant DType"); - let has_typed_value = typed_value_dtype.is_some(); - vortex_ensure!( - proto.has_value || has_typed_value, - "At least one of value or typed_value must be present" - ); - - let expected_children = 1 + proto.has_value as usize + has_typed_value as usize; - vortex_ensure!( - children.len() == expected_children || children.len() == expected_children + 1, - "Expected {} or {} children, got {}", - expected_children, - expected_children + 1, - children.len() - ); - - let (validity, mut child_idx) = if children.len() == expected_children { - (Validity::from(dtype.nullability()), 0) - } else { - (Validity::Array(children.get(0, &Validity::DTYPE, len)?), 1) - }; - let variant_metadata = - children.get(child_idx, &DType::Binary(Nullability::NonNullable), len)?; - child_idx += 1; - - let value = if proto.has_value { - let v = children.get(child_idx, &DType::Binary(proto.value_nullable.into()), len)?; - child_idx += 1; - Some(v) - } else { - None - }; - - let typed_value = if has_typed_value { - // typed_value can be any type — primitive, list, struct, etc. - let dtype = typed_value_dtype - .ok_or_else(|| vortex_err!("typed_value_dtype missing for typed_value child"))?; - let tv = children.get(child_idx, &dtype, len)?; - Some(tv) - } else { - None - }; - - let slots = ParquetVariantSlots { - validity: validity_to_child(&validity, len), - metadata: variant_metadata, - value, - typed_value, - } - .into_slots(); - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - EmptyArrayData, - slots, - )) - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { let shredded = array .typed_value() diff --git a/encodings/parquet-variant/src/vtable/plugin.rs b/encodings/parquet-variant/src/vtable/plugin.rs new file mode 100644 index 00000000000..83c6059ad1d --- /dev/null +++ b/encodings/parquet-variant/src/vtable/plugin.rs @@ -0,0 +1,162 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`ParquetVariant`]. + +use prost::Message; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::EmptyArrayData; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::validity::Validity; +use vortex_array::vtable::validity_to_child; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use super::ParquetVariantMetadataProto; +use crate::ParquetVariant; +use crate::array::ParquetVariantArraySlotsExt; +use crate::array::ParquetVariantSlots; + +impl ArrayPlugin for ParquetVariant { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "ParquetVariant plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + let typed_value_dtype = array + .typed_value() + .map(|tv| tv.dtype().try_into()) + .transpose()?; + let metadata = 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(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "ParquetVariant plugin does not recognize the serialized ID", + ); + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + vortex_ensure!( + buffers.is_empty(), + "ParquetVariantArray expects 0 buffers, got {}", + buffers.len() + ); + + let proto = ParquetVariantMetadataProto::decode(metadata)?; + let typed_value_dtype = match proto.typed_value_dtype.as_ref() { + Some(dtype) => Some(DType::from_proto(dtype, session)?), + None => None, + }; + + vortex_ensure!(matches!(dtype, DType::Variant(_)), "Expected Variant DType"); + let has_typed_value = typed_value_dtype.is_some(); + vortex_ensure!( + proto.has_value || has_typed_value, + "At least one of value or typed_value must be present" + ); + + let expected_children = 1 + proto.has_value as usize + has_typed_value as usize; + vortex_ensure!( + children.len() == expected_children || children.len() == expected_children + 1, + "Expected {} or {} children, got {}", + expected_children, + expected_children + 1, + children.len() + ); + + let (validity, mut child_idx) = if children.len() == expected_children { + (Validity::from(dtype.nullability()), 0) + } else { + (Validity::Array(children.get(0, &Validity::DTYPE, len)?), 1) + }; + let variant_metadata = + children.get(child_idx, &DType::Binary(Nullability::NonNullable), len)?; + child_idx += 1; + + let value = if proto.has_value { + let v = children.get(child_idx, &DType::Binary(proto.value_nullable.into()), len)?; + child_idx += 1; + Some(v) + } else { + None + }; + + let typed_value = if has_typed_value { + // typed_value can be any type — primitive, list, struct, etc. + let dtype = typed_value_dtype + .ok_or_else(|| vortex_err!("typed_value_dtype missing for typed_value child"))?; + let tv = children.get(child_idx, &dtype, len)?; + Some(tv) + } else { + None + }; + + let slots = ParquetVariantSlots { + validity: validity_to_child(&validity, len), + metadata: variant_metadata, + value, + typed_value, + } + .into_slots(); + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + EmptyArrayData, + slots, + ))? + .into_array()) + } +} diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index f7cc8c0a261..05c2c2810ba 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -17,7 +17,6 @@ use pco::match_number_enum; use pco::wrapped::ChunkDecompressor; use pco::wrapped::FileCompressor; use pco::wrapped::FileDecompressor; -use prost::Message; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -39,7 +38,6 @@ use vortex_array::dtype::DType; use vortex_array::dtype::PType; use vortex_array::dtype::half; use vortex_array::scalar::Scalar; -use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::OperationsVTable; use vortex_array::vtable::VTable; @@ -50,17 +48,17 @@ use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_error::VortexError; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::PcoChunkInfo; use crate::PcoMetadata; use crate::PcoPageInfo; +mod plugin; + // Overall approach here: // Chunk the array into Pco chunks (currently using the default recommended size // for good compression), and into finer-grained Pco pages. As we go, write each @@ -194,70 +192,6 @@ impl VTable for Pco { )) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - vortex_ensure!( - array.slice_start() == 0 && array.slice_stop() == array.unsliced_n_rows(), - "cannot serialize sliced PCO arrays; materialize the slice first" - ); - Ok(Some(array.metadata.clone().encode_to_vec())) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - let metadata = PcoMetadata::decode(metadata)?; - let validity = if children.is_empty() { - Validity::from(dtype.nullability()) - } else if children.len() == 1 { - let validity = children.get(0, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("PcoArray expected 0 or 1 child, got {}", children.len()); - }; - - vortex_ensure!(buffers.len() >= metadata.chunks.len()); - let chunk_metas = buffers[..metadata.chunks.len()] - .iter() - .map(|b| b.clone().try_to_host_sync()) - .collect::>>()?; - let pages = buffers[metadata.chunks.len()..] - .iter() - .map(|b| b.clone().try_to_host_sync()) - .collect::>>()?; - - let expected_n_pages = metadata - .chunks - .iter() - .map(|info| info.pages.len()) - .sum::(); - vortex_ensure_eq!(pages.len(), expected_n_pages); - - let slots = PcoSlots { - validity: validity_to_child(&validity, len), - } - .into_slots(); - // SAFETY: `Array::try_from_parts`, which consumes these parts, validates the data before - // publishing the array. - let data = - unsafe { PcoData::new_unchecked(chunk_metas, pages, dtype.as_ptype(), metadata, len) }; - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { PcoSlots::NAMES[idx].to_string() } diff --git a/encodings/pco/src/array/plugin.rs b/encodings/pco/src/array/plugin.rs new file mode 100644 index 00000000000..b68187ec500 --- /dev/null +++ b/encodings/pco/src/array/plugin.rs @@ -0,0 +1,127 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Pco`]. + +use prost::Message; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::validity::Validity; +use vortex_array::vtable::validity_to_child; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use super::Pco; +use super::PcoData; +use super::PcoSlots; +use crate::PcoMetadata; + +impl ArrayPlugin for Pco { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Pco plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + vortex_ensure!( + array.slice_start() == 0 && array.slice_stop() == array.unsliced_n_rows(), + "cannot serialize sliced PCO arrays; materialize the slice first" + ); + let metadata = array.metadata.clone().encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Pco plugin does not recognize the serialized ID", + ); + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + let metadata = PcoMetadata::decode(metadata)?; + let validity = if children.is_empty() { + Validity::from(dtype.nullability()) + } else if children.len() == 1 { + let validity = children.get(0, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + vortex_bail!("PcoArray expected 0 or 1 child, got {}", children.len()); + }; + + vortex_ensure!(buffers.len() >= metadata.chunks.len()); + let chunk_metas = buffers[..metadata.chunks.len()] + .iter() + .map(|b| b.clone().try_to_host_sync()) + .collect::>>()?; + let pages = buffers[metadata.chunks.len()..] + .iter() + .map(|b| b.clone().try_to_host_sync()) + .collect::>>()?; + + let expected_n_pages = metadata + .chunks + .iter() + .map(|info| info.pages.len()) + .sum::(); + vortex_ensure_eq!(pages.len(), expected_n_pages); + + let slots = PcoSlots { + validity: validity_to_child(&validity, len), + } + .into_slots(); + // SAFETY: `Array::try_from_parts`, which consumes these parts, validates the data before + // publishing the array. + let data = + unsafe { PcoData::new_unchecked(chunk_metas, pages, dtype.as_ptype(), metadata, len) }; + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots, + ))? + .into_array()) + } +} diff --git a/encodings/sequence/src/array.rs b/encodings/sequence/src/array.rs index 22438e35494..9d42611cc36 100644 --- a/encodings/sequence/src/array.rs +++ b/encodings/sequence/src/array.rs @@ -7,7 +7,6 @@ use std::hash::Hash; use std::hash::Hasher; use num_traits::AsPrimitive; -use prost::Message; use smallvec::smallvec; use vortex_array::Array; use vortex_array::ArrayEq; @@ -25,7 +24,6 @@ use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; -use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::PType; use vortex_array::expr::stats::Precision as StatPrecision; use vortex_array::expr::stats::Stat; @@ -35,7 +33,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; @@ -47,7 +44,6 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::compress::sequence_decompress; @@ -55,6 +51,8 @@ use crate::eval; use crate::eval::SequenceValue; use crate::rules::RULES; +mod plugin; + /// A [`Sequence`]-encoded Vortex array. pub type SequenceArray = Array; @@ -337,86 +335,6 @@ impl VTable for Sequence { vortex_array::vtable::with_empty_buffers(self, array, buffers) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - let metadata = SequenceMetadata { - base: Some((&array.base()).into()), - multiplier: Some((&array.multiplier()).into()), - }; - - Ok(Some(metadata.encode_to_vec())) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - vortex_ensure!( - buffers.is_empty(), - "SequenceArray expects 0 buffers, got {}", - buffers.len() - ); - vortex_ensure!( - children.is_empty(), - "SequenceArray expects 0 children, got {}", - children.len() - ); - let DType::Primitive(output_ptype, _) = dtype else { - vortex_bail!( - "only primitive dtypes are supported in SequenceArray currently, got {dtype}" - ); - }; - let metadata = SequenceMetadata::decode(metadata)?; - - let base_metadata = metadata - .base - .as_ref() - .ok_or_else(|| vortex_err!("base required"))?; - - let multiplier_metadata = metadata - .multiplier - .as_ref() - .ok_or_else(|| vortex_err!("multiplier required"))?; - - // We go via Scalar to validate that the value is valid for the ptype. - let base = Scalar::from_proto_value( - base_metadata, - &DType::Primitive(*output_ptype, NonNullable), - session, - )? - .as_primitive() - .pvalue() - .vortex_expect("sequence array base should be a non-nullable primitive"); - - // The serialized step preserves signedness independently of the output ptype. - let multiplier_ptype = SequenceData::multiplier_ptype_from_proto(multiplier_metadata)?; - let multiplier = Scalar::from_proto_value( - multiplier_metadata, - &DType::Primitive(multiplier_ptype, NonNullable), - session, - )? - .as_primitive() - .pvalue() - .vortex_expect("sequence array multiplier should be a non-nullable primitive"); - - let data = - SequenceData::try_new(base, multiplier, *output_ptype, dtype.nullability(), len)?; - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - ArraySlots::new(), - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { SLOT_NAMES[idx].to_string() } diff --git a/encodings/sequence/src/array/plugin.rs b/encodings/sequence/src/array/plugin.rs new file mode 100644 index 00000000000..76aa9c0bb6f --- /dev/null +++ b/encodings/sequence/src/array/plugin.rs @@ -0,0 +1,145 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Sequence`]. + +use prost::Message; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::ArraySlots; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::scalar::Scalar; +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_session::VortexSession; + +use crate::Sequence; +use crate::SequenceData; +use crate::array::SequenceMetadata; + +impl ArrayPlugin for Sequence { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Sequence plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + let metadata = SequenceMetadata { + base: Some((&array.base()).into()), + multiplier: Some((&array.multiplier()).into()), + }; + + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata.encode_to_vec(), + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Sequence plugin does not recognize the serialized ID", + ); + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + vortex_ensure!( + buffers.is_empty(), + "SequenceArray expects 0 buffers, got {}", + buffers.len() + ); + vortex_ensure!( + children.is_empty(), + "SequenceArray expects 0 children, got {}", + children.len() + ); + let DType::Primitive(output_ptype, _) = dtype else { + vortex_bail!( + "only primitive dtypes are supported in SequenceArray currently, got {dtype}" + ); + }; + let metadata = SequenceMetadata::decode(metadata)?; + + let base_metadata = metadata + .base + .as_ref() + .ok_or_else(|| vortex_err!("base required"))?; + + let multiplier_metadata = metadata + .multiplier + .as_ref() + .ok_or_else(|| vortex_err!("multiplier required"))?; + + // We go via Scalar to validate that the value is valid for the ptype. + let base = Scalar::from_proto_value( + base_metadata, + &DType::Primitive(*output_ptype, NonNullable), + session, + )? + .as_primitive() + .pvalue() + .vortex_expect("sequence array base should be a non-nullable primitive"); + + // The serialized step preserves signedness independently of the output ptype. + let multiplier_ptype = SequenceData::multiplier_ptype_from_proto(multiplier_metadata)?; + let multiplier = Scalar::from_proto_value( + multiplier_metadata, + &DType::Primitive(multiplier_ptype, NonNullable), + session, + )? + .as_primitive() + .pvalue() + .vortex_expect("sequence array multiplier should be a non-nullable primitive"); + + let data = + SequenceData::try_new(base, multiplier, *output_ptype, dtype.nullability(), len)?; + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + ArraySlots::new(), + ))? + .into_array()) + } +} diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index ecf00243f15..8aecb179b40 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -7,7 +7,6 @@ use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; -use prost::Message as _; use vortex_array::AnyCanonical; use vortex_array::Array; use vortex_array::ArrayEq; @@ -22,7 +21,6 @@ use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; -use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Primitive; @@ -44,7 +42,6 @@ use vortex_array::require_opt_child; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_array::scalar_fn::fns::operators::Operator; -use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; @@ -68,6 +65,7 @@ mod canonical; mod compute; mod kernel; mod ops; +mod plugin; mod rules; mod slice; @@ -250,66 +248,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() } diff --git a/encodings/sparse/src/plugin.rs b/encodings/sparse/src/plugin.rs new file mode 100644 index 00000000000..12e649b55c2 --- /dev/null +++ b/encodings/sparse/src/plugin.rs @@ -0,0 +1,123 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Sparse`]. + +use prost::Message; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::VortexSessionExecute; +use vortex_array::patches::Patches; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use crate::Sparse; +use crate::SparseData; +use crate::SparseExt; +use crate::SparseMetadata; + +impl ArrayPlugin for Sparse { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Sparse plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + 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(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata.encode_to_vec(), + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Sparse plugin does not recognize the serialized ID", + ); + 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()) + } +} diff --git a/encodings/zigzag/src/array.rs b/encodings/zigzag/src/array.rs index 101da38b521..760ebf53501 100644 --- a/encodings/zigzag/src/array.rs +++ b/encodings/zigzag/src/array.rs @@ -24,7 +24,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; @@ -36,7 +35,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 zigzag::ZigZag as ExternalZigZag; @@ -44,6 +42,8 @@ use crate::compute::ZigZagEncoded; use crate::rules::RULES; use crate::zigzag_decode; +mod plugin; + /// A [`ZigZag`]-encoded Vortex array. pub type ZigZagArray = Array; @@ -92,47 +92,6 @@ impl VTable for ZigZag { vortex_array::vtable::with_empty_buffers(self, array, buffers) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - if !metadata.is_empty() { - vortex_bail!( - "ZigZagArray expects empty metadata, got {} bytes", - metadata.len() - ); - } - if children.len() != 1 { - vortex_bail!("Expected 1 child, got {}", children.len()); - } - - let ptype = PType::try_from(dtype)?; - let encoded_type = DType::Primitive(ptype.to_unsigned(), dtype.nullability()); - - let encoded = children.get(0, &encoded_type, len)?; - let slots = smallvec![Some(encoded.clone())]; - let data = ZigZagData::try_new(encoded.dtype())?; - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { ZigZagSlots::NAMES[idx].to_string() } diff --git a/encodings/zigzag/src/array/plugin.rs b/encodings/zigzag/src/array/plugin.rs new file mode 100644 index 00000000000..924e08f0ffd --- /dev/null +++ b/encodings/zigzag/src/array/plugin.rs @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`ZigZag`]. + +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::smallvec::smallvec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use super::ZigZag; +use super::ZigZagData; + +impl ArrayPlugin for ZigZag { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "ZigZag plugin cannot serialize an array of another encoding", + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "ZigZag plugin does not recognize the serialized ID", + ); + let ArrayDeserialization { + dtype, + len, + metadata, + children, + .. + } = parts; + if !metadata.is_empty() { + vortex_bail!( + "ZigZagArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + if children.len() != 1 { + vortex_bail!("Expected 1 child, got {}", children.len()); + } + + let ptype = PType::try_from(dtype)?; + let encoded_type = DType::Primitive(ptype.to_unsigned(), dtype.nullability()); + + let encoded = children.get(0, &encoded_type, len)?; + let slots = smallvec![Some(encoded.clone())]; + let data = ZigZagData::try_new(encoded.dtype())?; + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots, + ))? + .into_array()) + } +} diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index 78034a5fd83..fae256f94b9 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -12,7 +12,6 @@ use std::sync::Arc; use itertools::Itertools as _; use num_traits::AsPrimitive; -use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -39,7 +38,6 @@ use vortex_array::dtype::DType; use vortex_array::dtype::OffsetBuilderPType; use vortex_array::match_each_varbin_builder; use vortex_array::scalar::Scalar; -use vortex_array::serde::ArrayChildren; use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; use vortex_array::vtable::OperationsVTable; @@ -61,7 +59,6 @@ use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_mask::AllOr; use vortex_mask::Mask; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use zstd::zstd_safe::WriteBuf; @@ -70,6 +67,8 @@ use crate::ZstdMetadata; use crate::decompressor::with_decompressor; use crate::validate_frame_content_size; +mod plugin; + // Zstd doesn't support training dictionaries on very few samples. const MIN_SAMPLES_FOR_DICTIONARY: usize = 8; type ViewLen = u32; @@ -217,63 +216,6 @@ impl VTable for Zstd { )) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(array.metadata.clone().encode_to_vec())) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - let metadata = ZstdMetadata::decode(metadata)?; - let validity = if children.is_empty() { - Validity::from(dtype.nullability()) - } else if children.len() == 1 { - let validity = children.get(0, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("ZstdArray expected 0 or 1 child, got {}", children.len()); - }; - - let (dictionary_buffer, compressed_buffers) = if metadata.dictionary_size == 0 { - // no dictionary - ( - None, - buffers - .iter() - .map(|b| b.clone().try_to_host_sync()) - .collect::>>()?, - ) - } else { - // with dictionary - ( - Some(buffers[0].clone().try_to_host_sync()?), - buffers[1..] - .iter() - .map(|b| b.clone().try_to_host_sync()) - .collect::>>()?, - ) - }; - - let slots = smallvec![validity_to_child(&validity, len)]; - let data = ZstdData::new(dictionary_buffer, compressed_buffers, metadata, len); - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { ZstdSlots::NAMES[idx].to_string() } diff --git a/encodings/zstd/src/array/plugin.rs b/encodings/zstd/src/array/plugin.rs new file mode 100644 index 00000000000..58995f4d08e --- /dev/null +++ b/encodings/zstd/src/array/plugin.rs @@ -0,0 +1,118 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Zstd`]. + +use prost::Message; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::smallvec::smallvec; +use vortex_array::validity::Validity; +use vortex_array::vtable::validity_to_child; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use crate::Zstd; +use crate::ZstdData; +use crate::ZstdMetadata; + +impl ArrayPlugin for Zstd { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Zstd plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + array.metadata.clone().encode_to_vec(), + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Zstd plugin does not recognize the serialized ID", + ); + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + let metadata = ZstdMetadata::decode(metadata)?; + let validity = if children.is_empty() { + Validity::from(dtype.nullability()) + } else if children.len() == 1 { + let validity = children.get(0, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + vortex_bail!("ZstdArray expected 0 or 1 child, got {}", children.len()); + }; + + let (dictionary_buffer, compressed_buffers) = if metadata.dictionary_size == 0 { + // no dictionary + ( + None, + buffers + .iter() + .map(|b| b.clone().try_to_host_sync()) + .collect::>>()?, + ) + } else { + // with dictionary + ( + Some(buffers[0].clone().try_to_host_sync()?), + buffers[1..] + .iter() + .map(|b| b.clone().try_to_host_sync()) + .collect::>>()?, + ) + }; + + let slots = smallvec![validity_to_child(&validity, len)]; + let data = ZstdData::new(dictionary_buffer, compressed_buffers, metadata, len); + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots, + ))? + .into_array()) + } +} diff --git a/encodings/zstd/src/zstd_buffers.rs b/encodings/zstd/src/zstd_buffers.rs index 7619ce9ec6e..0c916ac3bf1 100644 --- a/encodings/zstd/src/zstd_buffers.rs +++ b/encodings/zstd/src/zstd_buffers.rs @@ -8,7 +8,6 @@ use std::hash::Hash; use std::hash::Hasher; use std::sync::Arc; -use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayDeserialization; use vortex_array::ArrayEq; @@ -24,7 +23,6 @@ use vortex_array::ExecutionResult; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::scalar::Scalar; -use vortex_array::serde::ArrayChildren; use vortex_array::session::ArraySessionExt; use vortex_array::validity::Validity; use vortex_array::vtable::OperationsVTable; @@ -46,6 +44,8 @@ use crate::ZstdBuffersMetadata; use crate::decompressor::with_decompressor; use crate::validate_frame_content_size; +mod plugin; + /// A [`ZstdBuffers`]-encoded Vortex array. pub type ZstdBuffersArray = Array; @@ -531,100 +531,6 @@ impl VTable for ZstdBuffers { format!("child_{idx}") } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - let children: Vec<&ArrayRef> = array.slots().iter().flatten().collect(); - let child_dtypes = children - .iter() - .map(|child| child.dtype().try_into()) - .collect::>>()?; - let child_lens = children.iter().map(|child| child.len() as u64).collect(); - - Ok(Some( - ZstdBuffersMetadata { - inner_encoding_id: array.inner_encoding_id.to_string(), - inner_metadata: array.inner_metadata.clone(), - uncompressed_sizes: array.uncompressed_sizes.clone(), - buffer_alignments: array.buffer_alignments.clone(), - child_dtypes, - child_lens, - inner_child_count: Some(u32::try_from(array.inner_child_count)?), - validity_child: array.validity_slot.map(u32::try_from).transpose()?, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - let metadata = ZstdBuffersMetadata::decode(metadata)?; - let compressed_buffers: Vec = buffers.to_vec(); - - // Children belong to inner encodings, and serialization doesn't - // preserve their dtypes and values. Check dtypes are recovered from - // metadata. - vortex_ensure_eq!(metadata.child_dtypes.len(), children.len()); - vortex_ensure_eq!(metadata.child_lens.len(), children.len()); - - let mut slots = (0..children.len()) - .map(|i| { - let child_dtype = DType::from_proto(&metadata.child_dtypes[i], session)?; - let child_len = usize::try_from(metadata.child_lens[i])?; - children.get(i, &child_dtype, child_len).map(Some) - }) - .collect::>>()?; - - let mut data = ZstdBuffersData { - inner_encoding_id: array_id_from_string(&metadata.inner_encoding_id), - inner_metadata: metadata.inner_metadata.clone(), - compressed_buffers, - uncompressed_sizes: metadata.uncompressed_sizes.clone(), - buffer_alignments: metadata.buffer_alignments.clone(), - inner_child_count: children.len(), - validity_slot: None, - }; - data.validate()?; - - match metadata.inner_child_count { - Some(inner_child_count) => { - data.inner_child_count = usize::try_from(inner_child_count)?; - data.validity_slot = metadata.validity_child.map(usize::try_from).transpose()?; - } - // Arrays written before the validity was lifted out of the wrapped array only hold - // its children, so rebuild it once here to find its validity. - None if dtype.is_nullable() => { - let legacy = Array::try_from_parts(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data.clone(), - slots.clone().into(), - ))?; - let validity = - ZstdBuffers::decompress_and_build_inner(&legacy, session)?.validity()?; - data.validity_slot = lift_validity(&validity, len, &mut slots); - } - None => {} - } - - Ok(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 { @@ -667,8 +573,10 @@ impl ValidityVTable for ZstdBuffers { mod tests { use std::sync::LazyLock; + use prost::Message; use rstest::rstest; use vortex_array::ArrayContext; + use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -861,31 +769,32 @@ mod tests { let compressed = ZstdBuffers::compress(&input, 3, &session)?; let mut metadata = ZstdBuffersMetadata::decode( - ZstdBuffers::serialize(compressed.as_view(), &session)? + ArrayPlugin::serialize(&ZstdBuffers, compressed.as_array(), &session)? .vortex_expect("metadata") + .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(), + let decoded = ArrayPlugin::deserialize( + &ZstdBuffers, + ArrayDeserialization::new( + VTable::id(&ZstdBuffers), + input.dtype(), + input.len(), + &metadata.encode_to_vec(), + &compressed.data().compressed_buffers, + &children.as_slice(), + ), &session, )?; - let decoded = Array::try_from_parts(parts)?; - assert!(decoded.data().validity_slot.is_some()); + assert!(decoded.as_::().data().validity_slot.is_some()); let mut ctx = session.create_execution_ctx(); for i in 0..input.len() { - assert_eq!( - decoded.clone().into_array().is_valid(i, &mut ctx)?, - input.is_valid(i, &mut ctx)? - ); + assert_eq!(decoded.is_valid(i, &mut ctx)?, input.is_valid(i, &mut ctx)?); } Ok(()) } diff --git a/encodings/zstd/src/zstd_buffers/plugin.rs b/encodings/zstd/src/zstd_buffers/plugin.rs new file mode 100644 index 00000000000..5fb0cbb4437 --- /dev/null +++ b/encodings/zstd/src/zstd_buffers/plugin.rs @@ -0,0 +1,154 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`ZstdBuffers`]. + +use prost::Message; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use super::ZstdBuffersData; +use super::array_id_from_string; +use super::lift_validity; +use crate::ZstdBuffers; +use crate::ZstdBuffersMetadata; + +impl ArrayPlugin for ZstdBuffers { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "ZstdBuffers plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + let children: Vec<&ArrayRef> = array.slots().iter().flatten().collect(); + let child_dtypes = children + .iter() + .map(|child| child.dtype().try_into()) + .collect::>>()?; + let child_lens = children.iter().map(|child| child.len() as u64).collect(); + + let metadata = 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(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "ZstdBuffers plugin does not recognize the serialized ID", + ); + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + let metadata = ZstdBuffersMetadata::decode(metadata)?; + let compressed_buffers: Vec = buffers.to_vec(); + + // Children belong to inner encodings, and serialization doesn't + // preserve their dtypes and values. Check dtypes are recovered from + // metadata. + vortex_ensure_eq!(metadata.child_dtypes.len(), children.len()); + vortex_ensure_eq!(metadata.child_lens.len(), children.len()); + + let mut slots = (0..children.len()) + .map(|i| { + let child_dtype = DType::from_proto(&metadata.child_dtypes[i], session)?; + let child_len = usize::try_from(metadata.child_lens[i])?; + children.get(i, &child_dtype, child_len).map(Some) + }) + .collect::>>()?; + + let mut data = ZstdBuffersData { + inner_encoding_id: array_id_from_string(&metadata.inner_encoding_id), + inner_metadata: metadata.inner_metadata.clone(), + compressed_buffers, + uncompressed_sizes: metadata.uncompressed_sizes.clone(), + buffer_alignments: metadata.buffer_alignments.clone(), + inner_child_count: children.len(), + validity_slot: None, + }; + data.validate()?; + + match metadata.inner_child_count { + Some(inner_child_count) => { + data.inner_child_count = usize::try_from(inner_child_count)?; + data.validity_slot = metadata.validity_child.map(usize::try_from).transpose()?; + } + // Arrays written before the validity was lifted out of the wrapped array only hold + // its children, so rebuild it once here to find its validity. + None if dtype.is_nullable() => { + let legacy = Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data.clone(), + slots.clone().into(), + ))?; + let validity = + ZstdBuffers::decompress_and_build_inner(&legacy, session)?.validity()?; + data.validity_slot = lift_validity(&validity, len, &mut slots); + } + None => {} + } + + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots.into(), + ))? + .into_array()) + } +} 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..7b7219297bd 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; @@ -164,53 +160,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..20bbb34a972 100644 --- a/vortex-array/src/array/vtable/mod.rs +++ b/vortex-array/src/array/vtable/mod.rs @@ -5,7 +5,7 @@ //! //! A Vortex array encoding is implemented by a small static vtable type plus an associated //! `TypedArrayData` value stored in each array instance. The vtable owns behavior such as -//! validation, serialization, execution, child traversal, scalar access, and validity access. +//! validation, execution, child traversal, scalar access, and validity access. //! //! The public [`ArrayRef`] API performs common precondition checks before calling //! into these traits. Implementations should focus on encoding-specific work and uphold the @@ -26,7 +26,7 @@ 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,9 @@ use crate::validity::Validity; /// The logic is split across several "VTable" traits to enable easier code organization than /// simply lumping everything into a single trait. /// -/// From this [`VTable`] trait, we derive implementations for the sealed `DynArrayData` trait and the -/// public [`ArrayPlugin`] registry trait. +/// From this [`VTable`] trait, we derive implementations for the sealed `DynArrayData` trait. +/// Serialization is not part of the vtable: encodings that support serde implement +/// [`ArrayPlugin`] separately and register it with the session. /// /// The functions defined in these vtable traits will typically document their pre- and /// post-conditions. The pre-conditions are validated inside the `DynArrayData` and [`ArrayRef`] @@ -115,35 +115,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..e4f6181a542 100644 --- a/vortex-array/src/arrays/bool/vtable/mod.rs +++ b/vortex-array/src/arrays/bool/vtable/mod.rs @@ -4,8 +4,6 @@ use std::hash::Hash; use std::hash::Hasher; -use prost::Message; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -27,10 +25,9 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::BoolBuilder; use crate::dtype::DType; -use crate::serde::ArrayChildren; -use crate::validity::Validity; mod kernel; mod operations; +mod plugin; mod validity; use vortex_session::registry::CachedId; @@ -114,20 +111,6 @@ impl VTable for Bool { )) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - let offset = array.meta.offset(); - assert!(offset < 8, "Offset must be <8, got {offset}"); - Ok(Some( - BoolMetadata { - offset: u32::try_from(offset).vortex_expect("checked"), - } - .encode_to_vec(), - )) - } - fn validate( &self, data: &BoolData, @@ -158,41 +141,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() } diff --git a/vortex-array/src/arrays/bool/vtable/plugin.rs b/vortex-array/src/arrays/bool/vtable/plugin.rs new file mode 100644 index 00000000000..73407213aed --- /dev/null +++ b/vortex-array/src/arrays/bool/vtable/plugin.rs @@ -0,0 +1,108 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Bool`]. + +use prost::Message; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use super::BoolMetadata; +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::IntoArray; +use crate::VTable; +use crate::arrays::Bool; +use crate::arrays::bool::BoolData; +use crate::validity::Validity; + +impl ArrayPlugin for Bool { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Bool plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + let offset = array.meta.offset(); + assert!(offset < 8, "Offset must be <8, got {offset}"); + let metadata = BoolMetadata { + offset: u32::try_from(offset).vortex_expect("checked"), + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Bool plugin does not recognize the serialized ID", + ); + 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()) + } +} diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index af6709b3f2b..f7c046f36e4 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -4,15 +4,12 @@ use std::hash::Hasher; use itertools::Itertools; -use smallvec::SmallVec; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; -use vortex_error::vortex_err; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::ArrayEq; @@ -23,14 +20,12 @@ use crate::EqMode; use crate::ExecutionCtx; use crate::ExecutionResult; use crate::IntoArray; -use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayParts; use crate::array::ArrayView; use crate::array::VTable; use crate::array::with_empty_buffers; -use crate::arrays::PrimitiveArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::chunked::ChunkedData; use crate::arrays::chunked::array::ChunkedSlots; @@ -41,9 +36,9 @@ use crate::builders::ArrayBuilder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; -use crate::serde::ArrayChildren; mod canonical; mod operations; +mod plugin; mod validity; /// A [`Chunked`]-encoded Vortex array. @@ -158,76 +153,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, diff --git a/vortex-array/src/arrays/chunked/vtable/plugin.rs b/vortex-array/src/arrays/chunked/vtable/plugin.rs new file mode 100644 index 00000000000..e484881c6b8 --- /dev/null +++ b/vortex-array/src/arrays/chunked/vtable/plugin.rs @@ -0,0 +1,133 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Chunked`]. + +use itertools::Itertools; +use smallvec::SmallVec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::IntoArray; +use crate::VTable; +use crate::VortexSessionExecute; +use crate::arrays::Chunked; +use crate::arrays::PrimitiveArray; +use crate::arrays::chunked::ChunkedData; +use crate::arrays::chunked::array::ChunkedSlots; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; + +impl ArrayPlugin for Chunked { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Chunked plugin cannot serialize an array of another encoding", + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Chunked plugin does not recognize the serialized ID", + ); + 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..fcf4fcce79a 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -17,7 +17,6 @@ use crate::ArrayEq; use crate::ArrayHash; use crate::ArrayParts; use crate::ArrayRef; -use crate::ArraySlots; use crate::EqMode; use crate::ExecutionCtx; use crate::ExecutionResult; @@ -52,9 +51,9 @@ use crate::scalar::DecimalValue; use crate::scalar::ListScalar; use crate::scalar::Scalar; use crate::scalar::ScalarValue; -use crate::serde::ArrayChildren; pub(crate) mod canonical; mod operations; +mod plugin; mod validity; /// A [`Constant`]-encoded Vortex array. @@ -133,42 +132,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, diff --git a/vortex-array/src/arrays/constant/vtable/plugin.rs b/vortex-array/src/arrays/constant/vtable/plugin.rs new file mode 100644 index 00000000000..40740f57c3d --- /dev/null +++ b/vortex-array/src/arrays/constant/vtable/plugin.rs @@ -0,0 +1,90 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Constant`]. + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::ArraySlots; +use crate::IntoArray; +use crate::VTable; +use crate::arrays::Constant; +use crate::arrays::constant::ConstantData; +use crate::scalar::Scalar; +use crate::scalar::ScalarValue; + +impl ArrayPlugin for Constant { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Constant plugin cannot serialize an array of another encoding", + ); + // HACK: Because the scalar is stored in the buffers, we do not need to serialize the + // metadata at all. + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Constant plugin does not recognize the serialized ID", + ); + 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()) + } +} diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index 9153448d4c0..e21adb51186 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -42,7 +42,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 +102,6 @@ impl VTable for Decimal { )) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(None) - } - fn validate( &self, data: &DecimalData, @@ -138,18 +130,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() } diff --git a/vortex-array/src/arrays/dict/array.rs b/vortex-array/src/arrays/dict/array.rs index 6c8a6f22a65..16ea3b1fe5b 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; @@ -606,12 +608,16 @@ mod test { } .encode_to_vec(); let children: Vec = Vec::new(); - let result = Dict.deserialize( - &DType::Primitive(PType::I32, NonNullable), - 0, - &metadata, - &[], - &children, + let result = ArrayPlugin::deserialize( + &Dict, + 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..54c18fe9470 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/vortex-array/src/arrays/dict/vtable/mod.rs @@ -4,12 +4,8 @@ use std::hash::Hasher; use num_traits::AsPrimitive; -use prost::Message; -use smallvec::smallvec; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; -use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_mask::AllOr; use vortex_mask::Mask; @@ -17,7 +13,6 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use super::DictData; -use super::DictMetadata; use super::DictOwnedExt; use super::DictParts; use super::array::DictSlots; @@ -40,7 +35,6 @@ use crate::arrays::ConstantArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinView; -use crate::arrays::dict::DictArrayExt; use crate::arrays::dict::DictArraySlotsExt; use crate::arrays::dict::compute::rules::PARENT_RULES; use crate::arrays::dict::execute::take_canonical; @@ -49,19 +43,17 @@ use crate::builders::ArrayBuilder; use crate::builders::VarBinBuilder; use crate::builders::VarBinViewBuilder; use crate::dtype::DType; -use crate::dtype::Nullability; use crate::dtype::OffsetBuilderPType; -use crate::dtype::PType; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; use crate::match_each_integer_ptype; use crate::match_each_varbin_builder; use crate::require_child; use crate::scalar::Scalar; -use crate::serde::ArrayChildren; mod kernel; mod operations; +mod plugin; mod validity; /// A [`Dict`]-encoded Vortex array. @@ -136,63 +128,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() } diff --git a/vortex-array/src/arrays/dict/vtable/plugin.rs b/vortex-array/src/arrays/dict/vtable/plugin.rs new file mode 100644 index 00000000000..80af8ecf3ce --- /dev/null +++ b/vortex-array/src/arrays/dict/vtable/plugin.rs @@ -0,0 +1,120 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Dict`]. + +use prost::Message; +use smallvec::smallvec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::IntoArray; +use crate::VTable; +use crate::arrays::Dict; +use crate::arrays::dict::DictArrayExt; +use crate::arrays::dict::DictArraySlotsExt; +use crate::arrays::dict::DictData; +use crate::arrays::dict::DictMetadata; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; + +impl ArrayPlugin for Dict { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Dict plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + let metadata = 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(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Dict plugin does not recognize the serialized ID", + ); + 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()) + } +} diff --git a/vortex-array/src/arrays/extension/vtable/mod.rs b/vortex-array/src/arrays/extension/vtable/mod.rs index e8efe3a44f8..ad095e8c223 100644 --- a/vortex-array/src/arrays/extension/vtable/mod.rs +++ b/vortex-array/src/arrays/extension/vtable/mod.rs @@ -28,10 +28,10 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::ExtensionBuilder; use crate::dtype::DType; -use crate::serde::ArrayChildren; mod kernel; mod operations; +mod plugin; mod validity; /// An extension array that wraps another array with additional type information. @@ -138,45 +138,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() } diff --git a/vortex-array/src/arrays/extension/vtable/plugin.rs b/vortex-array/src/arrays/extension/vtable/plugin.rs new file mode 100644 index 00000000000..3688feb0d97 --- /dev/null +++ b/vortex-array/src/arrays/extension/vtable/plugin.rs @@ -0,0 +1,94 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Extension`]. + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::IntoArray; +use crate::VTable; +use crate::array::EmptyArrayData; +use crate::arrays::Extension; +use crate::arrays::extension::array::ExtensionSlots; +use crate::dtype::DType; + +impl ArrayPlugin for Extension { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Extension plugin cannot serialize an array of another encoding", + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Extension plugin does not recognize the serialized ID", + ); + 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..8aa24dd8455 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs @@ -33,10 +33,9 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::FixedSizeListBuilder; use crate::dtype::DType; -use crate::serde::ArrayChildren; -use crate::validity::Validity; mod kernel; mod operations; +mod plugin; mod validity; /// A [`FixedSizeList`]-encoded Vortex array. @@ -99,13 +98,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,60 +134,6 @@ impl VTable for FixedSizeList { Ok(()) } - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - if !metadata.is_empty() { - vortex_bail!( - "FixedSizeListArray expects empty metadata, got {} bytes", - metadata.len() - ); - } - vortex_ensure!( - buffers.is_empty(), - "`FixedSizeList::build` expects no buffers" - ); - - let DType::FixedSizeList(element_dtype, list_size, _) = &dtype else { - vortex_bail!("Expected `DType::FixedSizeList`, got {:?}", dtype); - }; - - let validity = { - if children.len() > 2 { - vortex_bail!("`FixedSizeList::build` method expected 1 or 2 children") - } - - if children.len() == 2 { - let validity = children.get(1, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - debug_assert_eq!(children.len(), 1); - Validity::from(dtype.nullability()) - } - }; - - let num_elements = len * (*list_size as usize); - let elements = children.get(0, element_dtype.as_ref(), num_elements)?; - - let data = - FixedSizeListData::try_build(elements.clone(), *list_size, validity.clone(), len)?; - let slots = FixedSizeListData::make_slots(&elements, &validity, len); - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { FixedSizeListSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/plugin.rs b/vortex-array/src/arrays/fixed_size_list/vtable/plugin.rs new file mode 100644 index 00000000000..046def3b94a --- /dev/null +++ b/vortex-array/src/arrays/fixed_size_list/vtable/plugin.rs @@ -0,0 +1,118 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`FixedSizeList`]. + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::IntoArray; +use crate::VTable; +use crate::arrays::FixedSizeList; +use crate::arrays::fixed_size_list::FixedSizeListData; +use crate::dtype::DType; +use crate::validity::Validity; + +impl ArrayPlugin for FixedSizeList { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "FixedSizeList plugin cannot serialize an array of another encoding", + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "FixedSizeList plugin does not recognize the serialized ID", + ); + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + if !metadata.is_empty() { + vortex_bail!( + "FixedSizeListArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + vortex_ensure!( + buffers.is_empty(), + "`FixedSizeList::build` expects no buffers" + ); + + let DType::FixedSizeList(element_dtype, list_size, _) = &dtype else { + vortex_bail!("Expected `DType::FixedSizeList`, got {:?}", dtype); + }; + + let validity = { + if children.len() > 2 { + vortex_bail!("`FixedSizeList::build` method expected 1 or 2 children") + } + + if children.len() == 2 { + let validity = children.get(1, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + debug_assert_eq!(children.len(), 1); + Validity::from(dtype.nullability()) + } + }; + + let num_elements = len * (*list_size as usize); + let elements = children.get(0, element_dtype.as_ref(), num_elements)?; + + let data = + FixedSizeListData::try_build(elements.clone(), *list_size, validity.clone(), len)?; + let slots = FixedSizeListData::make_slots(&elements, &validity, len); + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots, + ))? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/interleave/mod.rs b/vortex-array/src/arrays/interleave/mod.rs index b9c5c7374ee..cdaf8edffb6 100644 --- a/vortex-array/src/arrays/interleave/mod.rs +++ b/vortex-array/src/arrays/interleave/mod.rs @@ -73,7 +73,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 +216,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..9e41261cba8 100644 --- a/vortex-array/src/arrays/list/vtable/mod.rs +++ b/vortex-array/src/arrays/list/vtable/mod.rs @@ -21,14 +21,12 @@ use crate::EqMode; use crate::ExecutionCtx; use crate::ExecutionResult; use crate::IntoArray; -use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayParts; use crate::array::ArrayView; use crate::array::VTable; use crate::array::with_empty_buffers; -use crate::arrays::list::ListArraySlotsExt; use crate::arrays::list::ListData; use crate::arrays::list::ListSlots; use crate::arrays::list::compute::rules::PARENT_RULES; @@ -39,9 +37,8 @@ use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::match_each_list_builder; -use crate::serde::ArrayChildren; -use crate::validity::Validity; mod operations; +mod plugin; mod validity; /// A [`List`]-encoded Vortex array. pub type ListArray = Array; @@ -113,19 +110,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,58 +140,6 @@ impl VTable for List { Ok(()) } - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - let metadata = ListMetadata::decode(metadata)?; - let offset_ptype = metadata.get_offset_ptype()?; - let validity = if children.len() == 2 { - Validity::from(dtype.nullability()) - } else if children.len() == 3 { - let validity = children.get(2, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("Expected 2 or 3 children, got {}", children.len()); - }; - - let DType::List(element_dtype, _) = &dtype else { - vortex_bail!("Expected List dtype, got {:?}", dtype); - }; - let elements = children.get( - 0, - element_dtype.as_ref(), - usize::try_from(metadata.elements_len)?, - )?; - - let offsets = children.get( - 1, - &DType::Primitive(offset_ptype, Nullability::NonNullable), - len + 1, - )?; - - let data = ListData::try_build( - elements.clone(), - offsets.clone(), - validity.clone(), - &mut session.create_execution_ctx(), - )?; - let slots = ListData::make_slots(&elements, &offsets, &validity, len); - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { ListSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/list/vtable/plugin.rs b/vortex-array/src/arrays/list/vtable/plugin.rs new file mode 100644 index 00000000000..0d128c9b9c5 --- /dev/null +++ b/vortex-array/src/arrays/list/vtable/plugin.rs @@ -0,0 +1,126 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`List`]. + +use prost::Message; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use super::ListMetadata; +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::IntoArray; +use crate::VTable; +use crate::VortexSessionExecute; +use crate::arrays::List; +use crate::arrays::list::ListArraySlotsExt; +use crate::arrays::list::ListData; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::validity::Validity; + +impl ArrayPlugin for List { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "List plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + let metadata = ListMetadata { + elements_len: array.elements().len() as u64, + offset_ptype: PType::try_from(array.offsets().dtype())? as i32, + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "List plugin does not recognize the serialized ID", + ); + 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 { + Validity::from(dtype.nullability()) + } else if children.len() == 3 { + let validity = children.get(2, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + vortex_bail!("Expected 2 or 3 children, got {}", children.len()); + }; + + let DType::List(element_dtype, _) = &dtype else { + vortex_bail!("Expected List dtype, got {:?}", dtype); + }; + let elements = children.get( + 0, + element_dtype.as_ref(), + usize::try_from(metadata.elements_len)?, + )?; + + let offsets = children.get( + 1, + &DType::Primitive(offset_ptype, Nullability::NonNullable), + len + 1, + )?; + + let data = ListData::try_build( + elements.clone(), + offsets.clone(), + validity.clone(), + &mut session.create_execution_ctx(), + )?; + let slots = ListData::make_slots(&elements, &offsets, &validity, len); + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots, + ))? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/listview/vtable/mod.rs b/vortex-array/src/arrays/listview/vtable/mod.rs index c2cac004a31..a3060f2289a 100644 --- a/vortex-array/src/arrays/listview/vtable/mod.rs +++ b/vortex-array/src/arrays/listview/vtable/mod.rs @@ -5,7 +5,6 @@ use std::hash::Hash; use std::hash::Hasher; use std::sync::Arc; -use prost::Message; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -23,26 +22,22 @@ use crate::ArrayRef; use crate::EqMode; use crate::ExecutionCtx; use crate::ExecutionResult; -use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; use crate::array::with_empty_buffers; -use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::ListViewData; use crate::arrays::listview::ListViewSlots; use crate::arrays::listview::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::dtype::DType; -use crate::dtype::Nullability; use crate::dtype::PType; use crate::match_each_list_builder; -use crate::serde::ArrayChildren; -use crate::validity::Validity; mod kernel; mod operations; +mod plugin; mod validity; /// A [`ListView`]-encoded Vortex array. pub type ListViewArray = Array; @@ -126,20 +121,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,79 +155,6 @@ impl VTable for ListView { Ok(()) } - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - let metadata = ListViewMetadata::decode(metadata)?; - let offset_ptype = metadata.get_offset_ptype()?; - let size_ptype = metadata.get_size_ptype()?; - vortex_ensure!( - buffers.is_empty(), - "`ListViewArray::build` expects no buffers" - ); - - let DType::List(element_dtype, _) = dtype else { - vortex_bail!("Expected List dtype, got {:?}", dtype); - }; - - let validity = if children.len() == 3 { - Validity::from(dtype.nullability()) - } else if children.len() == 4 { - let validity = children.get(3, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!( - "`ListViewArray::build` expects 3 or 4 children, got {}", - children.len() - ); - }; - - // Get elements with the correct length from metadata. - let elements = children.get( - 0, - element_dtype.as_ref(), - usize::try_from(metadata.elements_len)?, - )?; - - // Get offsets with proper type from metadata. - let offsets = children.get( - 1, - &DType::Primitive(offset_ptype, Nullability::NonNullable), - len, - )?; - - // Get sizes with proper type from metadata. - let sizes = children.get( - 2, - &DType::Primitive(size_ptype, Nullability::NonNullable), - len, - )?; - - ListViewData::validate( - &elements, - &offsets, - &sizes, - &validity, - &mut session.create_execution_ctx(), - )?; - let data = ListViewData::try_new()?; - let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len); - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { ListViewSlots::NAMES[idx].to_string() } @@ -293,6 +201,8 @@ mod tests { use super::ListView; use super::ListViewMetadata; + use crate::ArrayDeserialization; + use crate::ArrayPlugin; use crate::ArrayRef; use crate::array::VTable; use crate::dtype::DType; @@ -323,7 +233,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/listview/vtable/plugin.rs b/vortex-array/src/arrays/listview/vtable/plugin.rs new file mode 100644 index 00000000000..c29482d8ab5 --- /dev/null +++ b/vortex-array/src/arrays/listview/vtable/plugin.rs @@ -0,0 +1,150 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`ListView`]. + +use prost::Message; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use super::ListViewMetadata; +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::IntoArray; +use crate::VTable; +use crate::VortexSessionExecute; +use crate::arrays::ListView; +use crate::arrays::listview::ListViewArraySlotsExt; +use crate::arrays::listview::ListViewData; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::validity::Validity; + +impl ArrayPlugin for ListView { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "ListView plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + let metadata = 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(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "ListView plugin does not recognize the serialized ID", + ); + 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()?; + vortex_ensure!( + buffers.is_empty(), + "`ListViewArray::build` expects no buffers" + ); + + let DType::List(element_dtype, _) = dtype else { + vortex_bail!("Expected List dtype, got {:?}", dtype); + }; + + let validity = if children.len() == 3 { + Validity::from(dtype.nullability()) + } else if children.len() == 4 { + let validity = children.get(3, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + vortex_bail!( + "`ListViewArray::build` expects 3 or 4 children, got {}", + children.len() + ); + }; + + // Get elements with the correct length from metadata. + let elements = children.get( + 0, + element_dtype.as_ref(), + usize::try_from(metadata.elements_len)?, + )?; + + // Get offsets with proper type from metadata. + let offsets = children.get( + 1, + &DType::Primitive(offset_ptype, Nullability::NonNullable), + len, + )?; + + // Get sizes with proper type from metadata. + let sizes = children.get( + 2, + &DType::Primitive(size_ptype, Nullability::NonNullable), + len, + )?; + + ListViewData::validate( + &elements, + &offsets, + &sizes, + &validity, + &mut session.create_execution_ctx(), + )?; + let data = ListViewData::try_new()?; + let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len); + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots, + ))? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/map/vtable/mod.rs b/vortex-array/src/arrays/map/vtable/mod.rs index 17980991d3a..ee5220cca1d 100644 --- a/vortex-array/src/arrays/map/vtable/mod.rs +++ b/vortex-array/src/arrays/map/vtable/mod.rs @@ -3,7 +3,6 @@ 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; @@ -19,7 +18,6 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::array::ValidityVTableFromChild; use crate::array::with_empty_buffers; -use crate::arrays::ListView; use crate::arrays::map::MapData; use crate::arrays::map::MapSlots; use crate::arrays::map::MapSlotsView; @@ -29,10 +27,10 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::dtype::DType; use crate::match_each_map_builder; -use crate::serde::ArrayChildren; mod kernel; mod operations; +mod plugin; mod validity; /// A [`Map`]-encoded Vortex array. @@ -96,54 +94,6 @@ impl VTable for Map { with_empty_buffers(self, array, buffers) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - if !metadata.is_empty() { - vortex_bail!( - "MapArray expects empty metadata, got {} bytes", - metadata.len() - ); - } - vortex_ensure!(buffers.is_empty(), "MapArray expects no buffers"); - - let DType::Map(map_dtype, nullability) = dtype else { - vortex_bail!("Expected map dtype, got {dtype}"); - }; - vortex_ensure_eq!(children.len(), MapSlots::COUNT); - - let expected_entries_dtype = - DType::List(std::sync::Arc::new(map_dtype.entries_dtype()), *nullability); - let entries = children.get(MapSlots::ENTRIES, &expected_entries_dtype, len)?; - vortex_ensure!( - entries.is::(), - "MapArray entries must use vortex.listview encoding, got {}", - entries.encoding_id() - ); - - let slots = MapData::make_slots(entries); - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - MapData, - slots, - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { MapSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/map/vtable/plugin.rs b/vortex-array/src/arrays/map/vtable/plugin.rs new file mode 100644 index 00000000000..c536f72985a --- /dev/null +++ b/vortex-array/src/arrays/map/vtable/plugin.rs @@ -0,0 +1,107 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Map`]. + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::IntoArray; +use crate::VTable; +use crate::arrays::ListView; +use crate::arrays::Map; +use crate::arrays::map::MapData; +use crate::arrays::map::MapSlots; +use crate::dtype::DType; + +impl ArrayPlugin for Map { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Map plugin cannot serialize an array of another encoding", + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Map plugin does not recognize the serialized ID", + ); + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + if !metadata.is_empty() { + vortex_bail!( + "MapArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + vortex_ensure!(buffers.is_empty(), "MapArray expects no buffers"); + + let DType::Map(map_dtype, nullability) = dtype else { + vortex_bail!("Expected map dtype, got {dtype}"); + }; + vortex_ensure_eq!(children.len(), MapSlots::COUNT); + + let expected_entries_dtype = + DType::List(std::sync::Arc::new(map_dtype.entries_dtype()), *nullability); + let entries = children.get(MapSlots::ENTRIES, &expected_entries_dtype, len)?; + vortex_ensure!( + entries.is::(), + "MapArray entries must use vortex.listview encoding, got {}", + entries.encoding_id() + ); + + let slots = MapData::make_slots(entries); + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + MapData, + slots, + ))? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/masked/vtable/mod.rs b/vortex-array/src/arrays/masked/vtable/mod.rs index 507223a94a0..383e08979fd 100644 --- a/vortex-array/src/arrays/masked/vtable/mod.rs +++ b/vortex-array/src/arrays/masked/vtable/mod.rs @@ -1,14 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors mod operations; +mod plugin; mod validity; use std::hash::Hasher; -use smallvec::smallvec; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_panic; @@ -23,12 +22,10 @@ use crate::ArrayRef; use crate::Canonical; use crate::EqMode; use crate::IntoArray; -use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; -use crate::array::validity_to_child; use crate::array::with_empty_buffers; use crate::arrays::ConstantArray; use crate::arrays::masked::MaskedArrayExt; @@ -43,7 +40,6 @@ use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; use crate::require_child; use crate::scalar::Scalar; -use crate::serde::ArrayChildren; use crate::validity::Validity; /// A [`Masked`]-encoded Vortex array. pub type MaskedArray = Array; @@ -115,63 +111,6 @@ impl VTable for Masked { 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!( - "MaskedArray expects empty metadata, got {} bytes", - metadata.len() - ); - } - if !buffers.is_empty() { - vortex_bail!("Expected 0 buffer, got {}", buffers.len()); - } - - vortex_ensure!( - children.len() == 1 || children.len() == 2, - "`MaskedArray::build` expects 1 or 2 children, got {}", - children.len() - ); - - let child = children.get(0, &dtype.as_nonnullable(), len)?; - - let validity = if children.len() == 2 { - let validity = children.get(1, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - Validity::from(dtype.nullability()) - }; - - let validity_slot = validity_to_child(&validity, len); - let data = MaskedData::try_new( - len, - child.all_valid(&mut session.create_execution_ctx())?, - validity, - )?; - Ok(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); diff --git a/vortex-array/src/arrays/masked/vtable/plugin.rs b/vortex-array/src/arrays/masked/vtable/plugin.rs new file mode 100644 index 00000000000..0c62c450cd0 --- /dev/null +++ b/vortex-array/src/arrays/masked/vtable/plugin.rs @@ -0,0 +1,116 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Masked`]. + +use smallvec::smallvec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::IntoArray; +use crate::VTable; +use crate::VortexSessionExecute; +use crate::array::validity_to_child; +use crate::arrays::Masked; +use crate::arrays::masked::MaskedData; +use crate::validity::Validity; + +impl ArrayPlugin for Masked { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Masked plugin cannot serialize an array of another encoding", + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Masked plugin does not recognize the serialized ID", + ); + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + if !metadata.is_empty() { + vortex_bail!( + "MaskedArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + if !buffers.is_empty() { + vortex_bail!("Expected 0 buffer, got {}", buffers.len()); + } + + vortex_ensure!( + children.len() == 1 || children.len() == 2, + "`MaskedArray::build` expects 1 or 2 children, got {}", + children.len() + ); + + let child = children.get(0, &dtype.as_nonnullable(), len)?; + + let validity = if children.len() == 2 { + let validity = children.get(1, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + Validity::from(dtype.nullability()) + }; + + let validity_slot = validity_to_child(&validity, len); + let data = MaskedData::try_new( + len, + child.all_valid(&mut session.create_execution_ctx())?, + validity, + )?; + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + smallvec![Some(child), validity_slot], + ))? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/null/mod.rs b/vortex-array/src/arrays/null/mod.rs index e55212db65a..c7170e2c3eb 100644 --- a/vortex-array/src/arrays/null/mod.rs +++ b/vortex-array/src/arrays/null/mod.rs @@ -3,7 +3,6 @@ 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; @@ -33,10 +32,10 @@ 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; +mod plugin; /// A [`Null`]-encoded Vortex array. pub type NullArray = Array; @@ -93,37 +92,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, diff --git a/vortex-array/src/arrays/null/plugin.rs b/vortex-array/src/arrays/null/plugin.rs new file mode 100644 index 00000000000..b7e93b6f37e --- /dev/null +++ b/vortex-array/src/arrays/null/plugin.rs @@ -0,0 +1,84 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Null`]. + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::ArraySlots; +use crate::IntoArray; +use crate::VTable; +use crate::array::EmptyArrayData; +use crate::arrays::Null; + +impl ArrayPlugin for Null { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Null plugin cannot serialize an array of another encoding", + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Null plugin does not recognize the serialized ID", + ); + 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()) + } +} diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs index 01c71d52997..e528b363205 100644 --- a/vortex-array/src/arrays/patched/vtable/mod.rs +++ b/vortex-array/src/arrays/patched/vtable/mod.rs @@ -7,6 +7,7 @@ use crate::ArrayEq; use crate::ArrayHash; mod kernels; mod operations; +mod plugin; mod slice; use std::hash::Hash; @@ -47,10 +48,8 @@ use crate::builders::ArrayBuilder; use crate::builders::PrimitiveBuilder; use crate::dtype::DType; use crate::dtype::NativePType; -use crate::dtype::PType; use crate::match_each_native_ptype; use crate::require_child; -use crate::serde::ArrayChildren; /// A [`Patched`]-encoded Vortex array. pub type PatchedArray = Array; @@ -138,60 +137,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, diff --git a/vortex-array/src/arrays/patched/vtable/plugin.rs b/vortex-array/src/arrays/patched/vtable/plugin.rs new file mode 100644 index 00000000000..4ae9ea592ba --- /dev/null +++ b/vortex-array/src/arrays/patched/vtable/plugin.rs @@ -0,0 +1,113 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Patched`]. + +use prost::Message; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use super::PatchedMetadata; +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::IntoArray; +use crate::VTable; +use crate::arrays::Patched; +use crate::arrays::patched::PatchedArrayExt; +use crate::arrays::patched::PatchedArraySlotsExt; +use crate::arrays::patched::PatchedData; +use crate::arrays::patched::PatchedSlots; +use crate::dtype::PType; + +impl ArrayPlugin for Patched { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Patched plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + let metadata = 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(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Patched plugin does not recognize the serialized ID", + ); + 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()) + } +} 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..c80bc1d0d47 100644 --- a/vortex-array/src/arrays/primitive/vtable/mod.rs +++ b/vortex-array/src/arrays/primitive/vtable/mod.rs @@ -3,7 +3,6 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use crate::ArrayParts; @@ -19,16 +18,14 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::PrimitiveBuilder; use crate::dtype::DType; -use crate::dtype::PType; use crate::match_each_native_ptype; -use crate::serde::ArrayChildren; mod kernel; mod operations; +mod plugin; mod validity; use std::hash::Hasher; -use vortex_buffer::Alignment; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -97,13 +94,6 @@ impl VTable for Primitive { )) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - fn validate( &self, data: &PrimitiveData, @@ -132,55 +122,6 @@ impl VTable for Primitive { Ok(()) } - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - if !metadata.is_empty() { - vortex_bail!( - "PrimitiveArray expects empty metadata, got {} bytes", - metadata.len() - ); - } - let buffer = fixed_width::single_buffer(buffers)?; - - let validity = fixed_width::deserialize_validity(dtype.nullability(), len, children)?; - - let ptype = PType::try_from(dtype)?; - - vortex_ensure!( - buffer.is_aligned_to(Alignment::new(ptype.byte_width())), - "Misaligned buffer cannot be used to build PrimitiveArray of {ptype}" - ); - - if buffer.len() != ptype.byte_width() * len { - vortex_bail!( - "Buffer length {} does not match expected length {} for {}, {}", - buffer.len(), - ptype.byte_width() * len, - ptype.byte_width(), - len, - ); - } - - // SAFETY: the buffer length and alignment are checked above. - let slots = PrimitiveData::make_slots(&validity, len); - let data = unsafe { PrimitiveData::new_unchecked_from_handle(buffer, ptype, validity) }; - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { PrimitiveSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/primitive/vtable/plugin.rs b/vortex-array/src/arrays/primitive/vtable/plugin.rs new file mode 100644 index 00000000000..e4aab246037 --- /dev/null +++ b/vortex-array/src/arrays/primitive/vtable/plugin.rs @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Primitive`]. + +use vortex_buffer::Alignment; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::IntoArray; +use crate::VTable; +use crate::arrays::Primitive; +use crate::arrays::fixed_width::vtable as fixed_width; +use crate::arrays::primitive::PrimitiveData; +use crate::dtype::PType; + +impl ArrayPlugin for Primitive { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Primitive plugin cannot serialize an array of another encoding", + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Primitive plugin does not recognize the serialized ID", + ); + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + if !metadata.is_empty() { + vortex_bail!( + "PrimitiveArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + let buffer = fixed_width::single_buffer(buffers)?; + + let validity = fixed_width::deserialize_validity(dtype.nullability(), len, children)?; + + let ptype = PType::try_from(dtype)?; + + vortex_ensure!( + buffer.is_aligned_to(Alignment::new(ptype.byte_width())), + "Misaligned buffer cannot be used to build PrimitiveArray of {ptype}" + ); + + if buffer.len() != ptype.byte_width() * len { + vortex_bail!( + "Buffer length {} does not match expected length {} for {}, {}", + buffer.len(), + ptype.byte_width() * len, + ptype.byte_width(), + len, + ); + } + + // SAFETY: the buffer length and alignment are checked above. + let slots = PrimitiveData::make_slots(&validity, len); + let data = unsafe { PrimitiveData::new_unchecked_from_handle(buffer, ptype, validity) }; + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots, + ))? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/runend/array.rs b/vortex-array/src/arrays/runend/array.rs index 3a4e5493387..5ce5210d5d7 100644 --- a/vortex-array/src/arrays/runend/array.rs +++ b/vortex-array/src/arrays/runend/array.rs @@ -52,11 +52,12 @@ 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; +mod plugin; + /// A [`RunEnd`]-encoded Vortex array. pub type RunEndArray = Array; @@ -139,56 +140,6 @@ impl VTable for RunEnd { crate::vtable::with_empty_buffers(self, array, buffers) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - RunEndMetadata { - ends_ptype: PType::try_from(array.ends().dtype()) - .vortex_expect("Must be a valid PType") as i32, - num_runs: array.ends().len() as u64, - offset: array.offset() as u64, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - let metadata = RunEndMetadata::decode(metadata)?; - let ends_ptype = metadata.get_ends_ptype()?; - let ends_dtype = DType::Primitive(ends_ptype, Nullability::NonNullable); - let runs = usize::try_from(metadata.num_runs).vortex_expect("Must be a valid usize"); - let ends = children.get(0, &ends_dtype, runs)?; - - let values = children.get(1, dtype, runs)?; - let offset = usize::try_from(metadata.offset).vortex_expect("Offset must be a valid usize"); - RunEndData::validate_parts( - &ends, - &values, - offset, - len, - &mut session.create_execution_ctx(), - )?; - let slots = RunEndSlots { ends, values }.into_slots(); - let data = RunEndData::new(offset); - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { RunEndSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/runend/array/plugin.rs b/vortex-array/src/arrays/runend/array/plugin.rs new file mode 100644 index 00000000000..039e3bceaba --- /dev/null +++ b/vortex-array/src/arrays/runend/array/plugin.rs @@ -0,0 +1,113 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`RunEnd`]. + +use prost::Message; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use super::RunEnd; +use super::RunEndArrayExt; +use super::RunEndArraySlotsExt; +use super::RunEndData; +use super::RunEndMetadata; +use super::RunEndSlots; +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::IntoArray; +use crate::VTable; +use crate::VortexSessionExecute; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; + +impl ArrayPlugin for RunEnd { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "RunEnd plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + let metadata = 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(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "RunEnd plugin does not recognize the serialized ID", + ); + 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); + let runs = usize::try_from(metadata.num_runs).vortex_expect("Must be a valid usize"); + let ends = children.get(0, &ends_dtype, runs)?; + + let values = children.get(1, dtype, runs)?; + let offset = usize::try_from(metadata.offset).vortex_expect("Offset must be a valid usize"); + RunEndData::validate_parts( + &ends, + &values, + offset, + len, + &mut session.create_execution_ctx(), + )?; + let slots = RunEndSlots { ends, values }.into_slots(); + let data = RunEndData::new(offset); + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots, + ))? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs index 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..7f65b5c5f3a 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -9,7 +9,6 @@ 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; @@ -41,7 +40,6 @@ use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; use crate::require_child; use crate::scalar::Scalar; -use crate::serde::ArrayChildren; use crate::validity::Validity; /// A [`Slice`]-encoded Vortex array. @@ -129,27 +127,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..36dc69683de 100644 --- a/vortex-array/src/arrays/struct_/vtable/mod.rs +++ b/vortex-array/src/arrays/struct_/vtable/mod.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_panic; @@ -18,16 +17,14 @@ use crate::array::VTable; use crate::array::child_to_validity; use crate::array::with_empty_buffers; use crate::arrays::struct_::array::StructSlots; -use crate::arrays::struct_::array::struct_slots_with_capacity; use crate::arrays::struct_::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::StructBuilder; use crate::dtype::DType; -use crate::serde::ArrayChildren; -use crate::validity::Validity; mod kernel; mod operations; +mod plugin; mod validity; use vortex_session::registry::CachedId; @@ -131,68 +128,6 @@ impl VTable for Struct { with_empty_buffers(self, array, buffers) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - if !metadata.is_empty() { - vortex_bail!( - "StructArray expects empty metadata, got {} bytes", - metadata.len() - ); - } - let DType::Struct(struct_dtype, nullability) = dtype else { - vortex_bail!("Expected struct dtype, found {:?}", dtype) - }; - - let (validity, non_data_children) = if children.len() == struct_dtype.nfields() { - (Validity::from(*nullability), 0_usize) - } else if children.len() == struct_dtype.nfields() + 1 { - let validity = children.get(0, &Validity::DTYPE, len)?; - (Validity::Array(validity), 1_usize) - } else { - vortex_bail!( - "Expected {} or {} children, found {}", - struct_dtype.nfields(), - struct_dtype.nfields() + 1, - children.len() - ); - }; - - let mut slots = struct_slots_with_capacity(&validity, len, struct_dtype.nfields()); - for i in 0..struct_dtype.nfields() { - let child_dtype = struct_dtype - .field_by_index(i) - .vortex_expect("no out of bounds"); - slots.push(Some(children.get( - non_data_children + i, - &child_dtype, - len, - )?)); - } - - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - EmptyArrayData, - slots, - )) - } - fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { if idx == StructSlots::VALIDITY { "validity".to_string() diff --git a/vortex-array/src/arrays/struct_/vtable/plugin.rs b/vortex-array/src/arrays/struct_/vtable/plugin.rs new file mode 100644 index 00000000000..663db412e71 --- /dev/null +++ b/vortex-array/src/arrays/struct_/vtable/plugin.rs @@ -0,0 +1,119 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Struct`]. + +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::IntoArray; +use crate::VTable; +use crate::array::EmptyArrayData; +use crate::arrays::Struct; +use crate::arrays::struct_::array::struct_slots_with_capacity; +use crate::dtype::DType; +use crate::validity::Validity; + +impl ArrayPlugin for Struct { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Struct plugin cannot serialize an array of another encoding", + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Struct plugin does not recognize the serialized ID", + ); + let ArrayDeserialization { + dtype, + len, + metadata, + children, + .. + } = parts; + if !metadata.is_empty() { + vortex_bail!( + "StructArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + let DType::Struct(struct_dtype, nullability) = dtype else { + vortex_bail!("Expected struct dtype, found {:?}", dtype) + }; + + let (validity, non_data_children) = if children.len() == struct_dtype.nfields() { + (Validity::from(*nullability), 0_usize) + } else if children.len() == struct_dtype.nfields() + 1 { + let validity = children.get(0, &Validity::DTYPE, len)?; + (Validity::Array(validity), 1_usize) + } else { + vortex_bail!( + "Expected {} or {} children, found {}", + struct_dtype.nfields(), + struct_dtype.nfields() + 1, + children.len() + ); + }; + + let mut slots = struct_slots_with_capacity(&validity, len, struct_dtype.nfields()); + for i in 0..struct_dtype.nfields() { + let child_dtype = struct_dtype + .field_by_index(i) + .vortex_expect("no out of bounds"); + slots.push(Some(children.get( + non_data_children + i, + &child_dtype, + len, + )?)); + } + + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + EmptyArrayData, + slots, + ))? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/union/vtable/mod.rs b/vortex-array/src/arrays/union/vtable/mod.rs index 7800f9430d2..5365fd44b5e 100644 --- a/vortex-array/src/arrays/union/vtable/mod.rs +++ b/vortex-array/src/arrays/union/vtable/mod.rs @@ -3,11 +3,8 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_error::vortex_ensure; -use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::ArrayRef; @@ -22,15 +19,13 @@ use crate::array::VTable; use crate::array::with_empty_buffers; use crate::arrays::union::UnionArrayExt; use crate::arrays::union::UnionSlots; -use crate::arrays::union::array::make_union_parts; use crate::arrays::union::compute::rules::PARENT_RULES; -use crate::arrays::union::union_type_ids_dtype; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::dtype::DType; -use crate::serde::ArrayChildren; mod operations; +mod plugin; mod validate; mod validity; @@ -100,47 +95,6 @@ impl VTable for Union { with_empty_buffers(self, array, buffers) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_ensure!(metadata.is_empty(), "UnionArray expects empty metadata"); - vortex_ensure!(buffers.is_empty(), "UnionArray expects no buffers"); - let DType::Union(variants, nullability) = dtype else { - vortex_bail!("Expected union dtype, found {dtype}") - }; - vortex_ensure_eq!(children.len(), UnionSlots::CHILDREN_OFFSET + variants.len()); - - let type_ids = children.get( - UnionSlots::TYPE_IDS, - &union_type_ids_dtype(*nullability), - len, - )?; - let sparse_children = variants - .variants() - .enumerate() - .map(|(index, dtype)| children.get(UnionSlots::CHILDREN_OFFSET + index, &dtype, len)) - .collect::>>()?; - - Ok(make_union_parts( - type_ids, - variants.clone(), - sparse_children, - )) - } - fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { if idx == UnionSlots::TYPE_IDS { "type_ids".to_string() diff --git a/vortex-array/src/arrays/union/vtable/plugin.rs b/vortex-array/src/arrays/union/vtable/plugin.rs new file mode 100644 index 00000000000..3983e148dd0 --- /dev/null +++ b/vortex-array/src/arrays/union/vtable/plugin.rs @@ -0,0 +1,99 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Union`]. + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::IntoArray; +use crate::VTable; +use crate::arrays::Union; +use crate::arrays::union::UnionSlots; +use crate::arrays::union::array::make_union_parts; +use crate::arrays::union::union_type_ids_dtype; +use crate::dtype::DType; + +impl ArrayPlugin for Union { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Union plugin cannot serialize an array of another encoding", + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Union plugin does not recognize the serialized ID", + ); + 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 { + vortex_bail!("Expected union dtype, found {dtype}") + }; + vortex_ensure_eq!(children.len(), UnionSlots::CHILDREN_OFFSET + variants.len()); + + let type_ids = children.get( + UnionSlots::TYPE_IDS, + &union_type_ids_dtype(*nullability), + len, + )?; + let sparse_children = variants + .variants() + .enumerate() + .map(|(index, dtype)| children.get(UnionSlots::CHILDREN_OFFSET + index, &dtype, len)) + .collect::>>()?; + + Ok(Array::try_from_parts(make_union_parts( + type_ids, + variants.clone(), + sparse_children, + ))? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/varbin/vtable/mod.rs b/vortex-array/src/arrays/varbin/vtable/mod.rs index e20611c4671..d9e3254cfd9 100644 --- a/vortex-array/src/arrays/varbin/vtable/mod.rs +++ b/vortex-array/src/arrays/varbin/vtable/mod.rs @@ -18,14 +18,12 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::ExecutionResult; use crate::IntoArray; -use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; use crate::arrays::PrimitiveArray; use crate::arrays::varbin::VarBinArrayExt; -use crate::arrays::varbin::VarBinArraySlotsExt; use crate::arrays::varbin::VarBinData; use crate::arrays::varbin::VarBinSlots; use crate::buffer::BufferHandle; @@ -36,11 +34,10 @@ use crate::dtype::Nullability; use crate::dtype::PType; use crate::match_each_integer_ptype; use crate::match_each_varbin_builder; -use crate::serde::ArrayChildren; -use crate::validity::Validity; pub(crate) mod canonical; mod kernel; mod operations; +mod plugin; mod validity; use canonical::varbin_to_canonical; @@ -155,67 +152,6 @@ impl VTable for VarBin { )) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - VarBinMetadata { - offsets_ptype: PType::try_from(array.offsets().dtype()) - .vortex_expect("Must be a valid PType") as i32, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - let metadata = VarBinMetadata::decode(metadata)?; - let offsets_ptype = metadata.get_offsets_ptype()?; - let validity = if children.len() == 1 { - Validity::from(dtype.nullability()) - } else if children.len() == 2 { - let validity = children.get(1, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("Expected 1 or 2 children, got {}", children.len()); - }; - - let offsets = children.get( - 0, - &DType::Primitive(offsets_ptype, Nullability::NonNullable), - len + 1, - )?; - - if buffers.len() != 1 { - vortex_bail!("Expected 1 buffer, got {}", buffers.len()); - } - let bytes = buffers[0].clone().try_to_host_sync()?; - - let data = VarBinData::try_build( - offsets.clone(), - bytes, - dtype.clone(), - validity.clone(), - &mut session.create_execution_ctx(), - )?; - let slots = VarBinData::make_slots(offsets, &validity, len); - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { VarBinSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/varbin/vtable/plugin.rs b/vortex-array/src/arrays/varbin/vtable/plugin.rs new file mode 100644 index 00000000000..9c48280d40c --- /dev/null +++ b/vortex-array/src/arrays/varbin/vtable/plugin.rs @@ -0,0 +1,125 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`VarBin`]. + +use prost::Message; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use super::VarBinMetadata; +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::IntoArray; +use crate::VTable; +use crate::VortexSessionExecute; +use crate::arrays::VarBin; +use crate::arrays::varbin::VarBinArraySlotsExt; +use crate::arrays::varbin::VarBinData; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::validity::Validity; + +impl ArrayPlugin for VarBin { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "VarBin plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + let metadata = VarBinMetadata { + offsets_ptype: PType::try_from(array.offsets().dtype()) + .vortex_expect("Must be a valid PType") as i32, + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "VarBin plugin does not recognize the serialized ID", + ); + 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 { + Validity::from(dtype.nullability()) + } else if children.len() == 2 { + let validity = children.get(1, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + vortex_bail!("Expected 1 or 2 children, got {}", children.len()); + }; + + let offsets = children.get( + 0, + &DType::Primitive(offsets_ptype, Nullability::NonNullable), + len + 1, + )?; + + if buffers.len() != 1 { + vortex_bail!("Expected 1 buffer, got {}", buffers.len()); + } + let bytes = buffers[0].clone().try_to_host_sync()?; + + let data = VarBinData::try_build( + offsets.clone(), + bytes, + dtype.clone(), + validity.clone(), + &mut session.create_execution_ctx(), + )?; + let slots = VarBinData::make_slots(offsets, &validity, len); + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots, + ))? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index f3d8d4147ed..8fad7406aed 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -2,15 +2,12 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::hash::Hasher; -use std::mem::size_of; use std::sync::Arc; -use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; -use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -20,12 +17,10 @@ use crate::ArrayRef; use crate::EqMode; use crate::ExecutionCtx; use crate::ExecutionResult; -use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; -use crate::arrays::varbinview::BinaryView; use crate::arrays::varbinview::VarBinViewData; use crate::arrays::varbinview::array::VarBinViewSlots; use crate::arrays::varbinview::compute::rules::PARENT_RULES; @@ -36,10 +31,9 @@ use crate::dtype::DType; use crate::hash::ArrayEq; use crate::hash::ArrayHash; use crate::match_each_varbin_builder; -use crate::serde::ArrayChildren; -use crate::validity::Validity; mod kernel; mod operations; +mod plugin; mod validity; /// A [`VarBinView`]-encoded Vortex array. pub type VarBinViewArray = Array; @@ -152,95 +146,6 @@ impl VTable for VarBinView { )) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some(vec![])) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - if !metadata.is_empty() { - vortex_bail!( - "VarBinViewArray expects empty metadata, got {} bytes", - metadata.len() - ); - } - let Some((views_handle, data_handles)) = buffers.split_last() else { - vortex_bail!("Expected at least 1 buffer, got 0"); - }; - - let validity = if children.is_empty() { - Validity::from(dtype.nullability()) - } else if children.len() == 1 { - let validity = children.get(0, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("Expected 0 or 1 children, got {}", children.len()); - }; - - let views_nbytes = views_handle.len(); - let expected_views_nbytes = len - .checked_mul(size_of::()) - .ok_or_else(|| vortex_err!("views byte length overflow for len={len}"))?; - if views_nbytes != expected_views_nbytes { - vortex_bail!( - "Expected views buffer length {} bytes, got {} bytes", - expected_views_nbytes, - views_nbytes - ); - } - - // If any buffer is on device, skip host validation and use try_new_handle. - if buffers.iter().any(|b| b.is_on_device()) { - let data = VarBinViewData::try_new_handle( - views_handle.clone(), - Arc::from(data_handles.to_vec()), - dtype.clone(), - validity.clone(), - )?; - let slots = VarBinViewData::make_slots(&validity, len); - return Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )); - } - - let data_buffers = data_handles - .iter() - .map(|b| b.as_host().clone()) - .collect::>(); - let views = Buffer::::from_byte_buffer(views_handle.clone().as_host().clone()); - - let data = VarBinViewData::try_new( - views, - Arc::from(data_buffers), - dtype.clone(), - validity.clone(), - &mut session.create_execution_ctx(), - )?; - let slots = VarBinViewData::make_slots(&validity, len); - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { VarBinViewSlots::NAMES[idx].to_string() } diff --git a/vortex-array/src/arrays/varbinview/vtable/plugin.rs b/vortex-array/src/arrays/varbinview/vtable/plugin.rs new file mode 100644 index 00000000000..ad9749a9383 --- /dev/null +++ b/vortex-array/src/arrays/varbinview/vtable/plugin.rs @@ -0,0 +1,152 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`VarBinView`]. + +use std::mem::size_of; +use std::sync::Arc; + +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::IntoArray; +use crate::VTable; +use crate::VortexSessionExecute; +use crate::arrays::VarBinView; +use crate::arrays::varbinview::BinaryView; +use crate::arrays::varbinview::VarBinViewData; +use crate::validity::Validity; + +impl ArrayPlugin for VarBinView { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "VarBinView plugin cannot serialize an array of another encoding", + ); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array, + vec![], + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "VarBinView plugin does not recognize the serialized ID", + ); + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + if !metadata.is_empty() { + vortex_bail!( + "VarBinViewArray expects empty metadata, got {} bytes", + metadata.len() + ); + } + let Some((views_handle, data_handles)) = buffers.split_last() else { + vortex_bail!("Expected at least 1 buffer, got 0"); + }; + + let validity = if children.is_empty() { + Validity::from(dtype.nullability()) + } else if children.len() == 1 { + let validity = children.get(0, &Validity::DTYPE, len)?; + Validity::Array(validity) + } else { + vortex_bail!("Expected 0 or 1 children, got {}", children.len()); + }; + + let views_nbytes = views_handle.len(); + let expected_views_nbytes = len + .checked_mul(size_of::()) + .ok_or_else(|| vortex_err!("views byte length overflow for len={len}"))?; + if views_nbytes != expected_views_nbytes { + vortex_bail!( + "Expected views buffer length {} bytes, got {} bytes", + expected_views_nbytes, + views_nbytes + ); + } + + // If any buffer is on device, skip host validation and use try_new_handle. + if buffers.iter().any(|b| b.is_on_device()) { + let data = VarBinViewData::try_new_handle( + views_handle.clone(), + Arc::from(data_handles.to_vec()), + dtype.clone(), + validity.clone(), + )?; + let slots = VarBinViewData::make_slots(&validity, len); + return Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots, + ))? + .into_array()); + } + + let data_buffers = data_handles + .iter() + .map(|b| b.as_host().clone()) + .collect::>(); + let views = Buffer::::from_byte_buffer(views_handle.clone().as_host().clone()); + + let data = VarBinViewData::try_new( + views, + Arc::from(data_buffers), + dtype.clone(), + validity.clone(), + &mut session.create_execution_ctx(), + )?; + let slots = VarBinViewData::make_slots(&validity, len); + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + data, + slots, + ))? + .into_array()) + } +} diff --git a/vortex-array/src/arrays/variant/vtable/mod.rs b/vortex-array/src/arrays/variant/vtable/mod.rs index 47cb8f0f18e..9105b4fba72 100644 --- a/vortex-array/src/arrays/variant/vtable/mod.rs +++ b/vortex-array/src/arrays/variant/vtable/mod.rs @@ -3,9 +3,9 @@ mod kernel; mod operations; +mod plugin; mod validity; -use prost::Message; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -36,7 +36,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,59 +124,6 @@ impl VTable for Variant { with_empty_buffers(self, array, buffers) } - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - let shredded_dtype = array.slots()[VariantSlots::SHREDDED] - .as_ref() - .map(|shredded| shredded.dtype().try_into()) - .transpose()?; - Ok(Some( - VariantMetadataProto { shredded_dtype }.encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - vortex_ensure!( - buffers.is_empty(), - "VariantArray expects 0 buffers, got {}", - buffers.len() - ); - let proto = VariantMetadataProto::decode(metadata)?; - let shredded_dtype = proto - .shredded_dtype - .as_ref() - .map(|dtype| DType::from_proto(dtype, session)) - .transpose()?; - vortex_ensure!(matches!(dtype, DType::Variant(_)), "Expected Variant DType"); - let expected_children = 1 + usize::from(shredded_dtype.is_some()); - vortex_ensure_eq!(children.len(), expected_children); - let core_storage = children.get(0, dtype, len)?; - let shredded = shredded_dtype - .map(|dtype| children.get(1, &dtype, len)) - .transpose()?; - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - EmptyArrayData, - VariantSlots { - core_storage, - shredded, - } - .into_slots(), - )) - } - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { match VariantSlots::NAMES.get(idx) { Some(name) => (*name).to_string(), diff --git a/vortex-array/src/arrays/variant/vtable/plugin.rs b/vortex-array/src/arrays/variant/vtable/plugin.rs new file mode 100644 index 00000000000..0ec40192759 --- /dev/null +++ b/vortex-array/src/arrays/variant/vtable/plugin.rs @@ -0,0 +1,112 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`Variant`]. + +use prost::Message; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use super::VariantMetadataProto; +use crate::Array; +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::IntoArray; +use crate::VTable; +use crate::array::EmptyArrayData; +use crate::arrays::Variant; +use crate::arrays::variant::VariantSlots; +use crate::dtype::DType; + +impl ArrayPlugin for Variant { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!( + array.encoding_id(), + VTable::id(self), + "Variant plugin cannot serialize an array of another encoding", + ); + let array = array.as_::(); + let shredded_dtype = array.slots()[VariantSlots::SHREDDED] + .as_ref() + .map(|shredded| shredded.dtype().try_into()) + .transpose()?; + let metadata = VariantMetadataProto { shredded_dtype }.encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + VTable::id(self), + array.array(), + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + VTable::id(self), + "Variant plugin does not recognize the serialized ID", + ); + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + vortex_ensure!( + buffers.is_empty(), + "VariantArray expects 0 buffers, got {}", + buffers.len() + ); + let proto = VariantMetadataProto::decode(metadata)?; + let shredded_dtype = proto + .shredded_dtype + .as_ref() + .map(|dtype| DType::from_proto(dtype, session)) + .transpose()?; + vortex_ensure!(matches!(dtype, DType::Variant(_)), "Expected Variant DType"); + let expected_children = 1 + usize::from(shredded_dtype.is_some()); + vortex_ensure_eq!(children.len(), expected_children); + let core_storage = children.get(0, dtype, len)?; + let shredded = shredded_dtype + .map(|dtype| children.get(1, &dtype, len)) + .transpose()?; + Ok(Array::try_from_parts(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + EmptyArrayData, + VariantSlots { + core_storage, + shredded, + } + .into_slots(), + ))? + .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..5f08b58c898 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,18 @@ mod tests { "old primitive wire ID cannot represent length {}", parts.len, ); - Ok(Array::::try_from_parts(ArrayVTable::deserialize( + ArrayPlugin::deserialize( &Primitive, - parts.dtype, - parts.len, - parts.metadata, - parts.buffers, - parts.children, + ArrayDeserialization::new( + old_primitive_id(), + parts.dtype, + parts.len, + parts.metadata, + parts.buffers, + parts.children, + ), session, - )?)? - .into_array()) + ) } } diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index 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-file/src/tests.rs b/vortex-file/src/tests.rs index 1dcf89a916e..1144b84b54b 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -174,13 +174,13 @@ async fn test_eof_values() { #[rstest] #[case::default( BtrBlocksCompressorBuilder::from_session(&SESSION), - if cfg!(feature = "zstd") { 70_060 } else { 69_996 } + if cfg!(feature = "zstd") { 70_012 } else { 69_948 } )] #[cfg_attr( feature = "zstd", case::compact( BtrBlocksCompressorBuilder::from_session(&SESSION).with_compact(), - 55_144 + 55_096 ) )] #[tokio::test] 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..3572d9e674d 100644 --- a/vortex-python/src/arrays/py/vtable.rs +++ b/vortex-python/src/arrays/py/vtable.rs @@ -19,7 +19,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; @@ -28,10 +27,11 @@ use vortex::error::vortex_bail; use vortex::error::vortex_ensure_eq; use vortex::error::vortex_panic; use vortex::scalar::Scalar; -use vortex::session::VortexSession; use crate::arrays::py::PythonArray; +mod plugin; + /// Wrapper struct encapsulating a Python encoding. #[derive(Debug, Clone)] pub struct PythonVTable { @@ -93,26 +93,6 @@ impl VTable for PythonVTable { with_empty_buffers(self, array, buffers) } - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(None) - } - - fn deserialize( - &self, - _dtype: &DType, - _len: usize, - bytes: &[u8], - _buffers: &[BufferHandle], - _children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - _ = bytes; - vortex_bail!("PythonArray deserialization is not supported"); - } - fn slot_name(_array: ArrayView<'_, Self>, _idx: usize) -> String { vortex_panic!("PythonArray has no slots") } diff --git a/vortex-python/src/arrays/py/vtable/plugin.rs b/vortex-python/src/arrays/py/vtable/plugin.rs new file mode 100644 index 00000000000..3b2f0ddb43d --- /dev/null +++ b/vortex-python/src/arrays/py/vtable/plugin.rs @@ -0,0 +1,46 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`ArrayPlugin`] implementation for [`PythonVTable`]. + +use vortex::array::ArrayDeserialization; +use vortex::array::ArrayId; +use vortex::array::ArrayPlugin; +use vortex::array::ArrayRef; +use vortex::array::ArraySerialization; +use vortex::array::VTable; +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::session::VortexSession; + +use super::PythonVTable; + +impl ArrayPlugin for PythonVTable { + fn id(&self) -> ArrayId { + VTable::id(self) + } + + fn read_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn write_ids(&self) -> Vec { + vec![VTable::id(self)] + } + + fn serialize( + &self, + _array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + Ok(None) + } + + fn deserialize( + &self, + _parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_bail!("PythonArray deserialization is not supported"); + } +} From d3d569e54d2d7fb23e271022739937dda7249c0d Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 6 Nov 2026 10:50:56 -0500 Subject: [PATCH 2/8] default read_ids and write_ids to the plugin id Signed-off-by: Matt Katz --- encodings/alp/src/alp/array/plugin.rs | 14 +++----------- encodings/alp/src/alp/plugin.rs | 8 -------- encodings/alp/src/alp_rd/array/plugin.rs | 14 +++----------- encodings/bytebool/src/array/plugin.rs | 14 +++----------- encodings/datetime-parts/src/array/plugin.rs | 14 +++----------- encodings/elias-fano/src/unary/vtable/plugin.rs | 14 +++----------- encodings/fastlanes/src/delta/vtable/plugin.rs | 14 +++----------- encodings/fastlanes/src/rle/vtable/plugin.rs | 14 +++----------- encodings/fsst/src/array/plugin.rs | 14 +++----------- encodings/onpair/src/array/plugin.rs | 14 +++----------- encodings/parquet-variant/src/vtable/plugin.rs | 14 +++----------- encodings/pco/src/array/plugin.rs | 14 +++----------- encodings/sequence/src/array/plugin.rs | 14 +++----------- encodings/sparse/src/plugin.rs | 14 +++----------- encodings/zigzag/src/array/plugin.rs | 14 +++----------- encodings/zstd/src/array/plugin.rs | 14 +++----------- encodings/zstd/src/zstd_buffers/plugin.rs | 14 +++----------- vortex-array/src/array/plugin.rs | 12 ++++++++++-- vortex-array/src/arrays/bool/vtable/plugin.rs | 14 +++----------- vortex-array/src/arrays/chunked/vtable/plugin.rs | 14 +++----------- vortex-array/src/arrays/constant/vtable/plugin.rs | 14 +++----------- .../src/arrays/decimal/vtable/plugin/mod.rs | 8 -------- vortex-array/src/arrays/dict/vtable/plugin.rs | 14 +++----------- vortex-array/src/arrays/extension/vtable/plugin.rs | 14 +++----------- .../src/arrays/fixed_size_list/vtable/plugin.rs | 14 +++----------- vortex-array/src/arrays/list/vtable/plugin.rs | 14 +++----------- vortex-array/src/arrays/listview/vtable/plugin.rs | 14 +++----------- vortex-array/src/arrays/map/vtable/plugin.rs | 14 +++----------- vortex-array/src/arrays/masked/vtable/plugin.rs | 14 +++----------- vortex-array/src/arrays/null/plugin.rs | 14 +++----------- vortex-array/src/arrays/patched/vtable/plugin.rs | 14 +++----------- vortex-array/src/arrays/primitive/vtable/plugin.rs | 14 +++----------- vortex-array/src/arrays/runend/array/plugin.rs | 14 +++----------- vortex-array/src/arrays/scalar_fn/plugin.rs | 8 -------- vortex-array/src/arrays/struct_/vtable/plugin.rs | 14 +++----------- vortex-array/src/arrays/union/vtable/plugin.rs | 14 +++----------- vortex-array/src/arrays/varbin/vtable/plugin.rs | 14 +++----------- .../src/arrays/varbinview/vtable/plugin.rs | 14 +++----------- vortex-array/src/arrays/variant/vtable/plugin.rs | 14 +++----------- vortex-python/src/arrays/py/vtable/plugin.rs | 8 -------- 40 files changed, 115 insertions(+), 419 deletions(-) diff --git a/encodings/alp/src/alp/array/plugin.rs b/encodings/alp/src/alp/array/plugin.rs index 58462cf7475..b4af77efd67 100644 --- a/encodings/alp/src/alp/array/plugin.rs +++ b/encodings/alp/src/alp/array/plugin.rs @@ -33,14 +33,6 @@ impl ArrayPlugin for ALP { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -48,7 +40,7 @@ impl ArrayPlugin for ALP { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "ALP plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -63,7 +55,7 @@ impl ArrayPlugin for ALP { } .encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -76,7 +68,7 @@ impl ArrayPlugin for ALP { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "ALP plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/encodings/alp/src/alp/plugin.rs b/encodings/alp/src/alp/plugin.rs index dca212568e9..4ae02ea455b 100644 --- a/encodings/alp/src/alp/plugin.rs +++ b/encodings/alp/src/alp/plugin.rs @@ -37,14 +37,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, diff --git a/encodings/alp/src/alp_rd/array/plugin.rs b/encodings/alp/src/alp_rd/array/plugin.rs index 00e83ed39ac..ecd4b3970a8 100644 --- a/encodings/alp/src/alp_rd/array/plugin.rs +++ b/encodings/alp/src/alp_rd/array/plugin.rs @@ -37,14 +37,6 @@ impl ArrayPlugin for ALPRD { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -52,7 +44,7 @@ impl ArrayPlugin for ALPRD { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "ALPRD plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -74,7 +66,7 @@ impl ArrayPlugin for ALPRD { } .encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -87,7 +79,7 @@ impl ArrayPlugin for ALPRD { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "ALPRD plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/encodings/bytebool/src/array/plugin.rs b/encodings/bytebool/src/array/plugin.rs index 85da7f8632e..bb506595202 100644 --- a/encodings/bytebool/src/array/plugin.rs +++ b/encodings/bytebool/src/array/plugin.rs @@ -26,14 +26,6 @@ impl ArrayPlugin for ByteBool { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -41,11 +33,11 @@ impl ArrayPlugin for ByteBool { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "ByteBool plugin cannot serialize an array of another encoding", ); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array, vec![], ))) @@ -58,7 +50,7 @@ impl ArrayPlugin for ByteBool { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "ByteBool plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/encodings/datetime-parts/src/array/plugin.rs b/encodings/datetime-parts/src/array/plugin.rs index 47d6cab44e0..3b30937f0e3 100644 --- a/encodings/datetime-parts/src/array/plugin.rs +++ b/encodings/datetime-parts/src/array/plugin.rs @@ -32,14 +32,6 @@ impl ArrayPlugin for DateTimeParts { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -47,7 +39,7 @@ impl ArrayPlugin for DateTimeParts { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "DateTimeParts plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -58,7 +50,7 @@ impl ArrayPlugin for DateTimeParts { } .encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -71,7 +63,7 @@ impl ArrayPlugin for DateTimeParts { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "DateTimeParts plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/encodings/elias-fano/src/unary/vtable/plugin.rs b/encodings/elias-fano/src/unary/vtable/plugin.rs index 222cc8b78c0..0c6f6e0c51b 100644 --- a/encodings/elias-fano/src/unary/vtable/plugin.rs +++ b/encodings/elias-fano/src/unary/vtable/plugin.rs @@ -28,14 +28,6 @@ impl ArrayPlugin for Unary { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -43,7 +35,7 @@ impl ArrayPlugin for Unary { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Unary plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -53,7 +45,7 @@ impl ArrayPlugin for Unary { } .encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -66,7 +58,7 @@ impl ArrayPlugin for Unary { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Unary plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/encodings/fastlanes/src/delta/vtable/plugin.rs b/encodings/fastlanes/src/delta/vtable/plugin.rs index fa55d39a23e..8f1cbabc988 100644 --- a/encodings/fastlanes/src/delta/vtable/plugin.rs +++ b/encodings/fastlanes/src/delta/vtable/plugin.rs @@ -35,14 +35,6 @@ impl ArrayPlugin for Delta { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -50,7 +42,7 @@ impl ArrayPlugin for Delta { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Delta plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -60,7 +52,7 @@ impl ArrayPlugin for Delta { } .encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -73,7 +65,7 @@ impl ArrayPlugin for Delta { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Delta plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/encodings/fastlanes/src/rle/vtable/plugin.rs b/encodings/fastlanes/src/rle/vtable/plugin.rs index c3e104d8373..6f78cb6bbb0 100644 --- a/encodings/fastlanes/src/rle/vtable/plugin.rs +++ b/encodings/fastlanes/src/rle/vtable/plugin.rs @@ -33,14 +33,6 @@ impl ArrayPlugin for RLE { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -48,7 +40,7 @@ impl ArrayPlugin for RLE { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "RLE plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -62,7 +54,7 @@ impl ArrayPlugin for RLE { } .encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -75,7 +67,7 @@ impl ArrayPlugin for RLE { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "RLE plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/encodings/fsst/src/array/plugin.rs b/encodings/fsst/src/array/plugin.rs index 93d8a40df59..14169bc1f6c 100644 --- a/encodings/fsst/src/array/plugin.rs +++ b/encodings/fsst/src/array/plugin.rs @@ -37,14 +37,6 @@ impl ArrayPlugin for FSST { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -52,7 +44,7 @@ impl ArrayPlugin for FSST { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "FSST plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -63,7 +55,7 @@ impl ArrayPlugin for FSST { } .encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -102,7 +94,7 @@ impl ArrayPlugin for FSST { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "FSST plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/encodings/onpair/src/array/plugin.rs b/encodings/onpair/src/array/plugin.rs index 79fcf92ac30..cd254ac4916 100644 --- a/encodings/onpair/src/array/plugin.rs +++ b/encodings/onpair/src/array/plugin.rs @@ -35,14 +35,6 @@ impl ArrayPlugin for OnPair { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -50,7 +42,7 @@ impl ArrayPlugin for OnPair { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "OnPair plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -67,7 +59,7 @@ impl ArrayPlugin for OnPair { } .encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -80,7 +72,7 @@ impl ArrayPlugin for OnPair { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "OnPair plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/encodings/parquet-variant/src/vtable/plugin.rs b/encodings/parquet-variant/src/vtable/plugin.rs index 83c6059ad1d..b0c981ac2f5 100644 --- a/encodings/parquet-variant/src/vtable/plugin.rs +++ b/encodings/parquet-variant/src/vtable/plugin.rs @@ -34,14 +34,6 @@ impl ArrayPlugin for ParquetVariant { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -49,7 +41,7 @@ impl ArrayPlugin for ParquetVariant { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "ParquetVariant plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -64,7 +56,7 @@ impl ArrayPlugin for ParquetVariant { } .encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -77,7 +69,7 @@ impl ArrayPlugin for ParquetVariant { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "ParquetVariant plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/encodings/pco/src/array/plugin.rs b/encodings/pco/src/array/plugin.rs index b68187ec500..bb67c854d89 100644 --- a/encodings/pco/src/array/plugin.rs +++ b/encodings/pco/src/array/plugin.rs @@ -31,14 +31,6 @@ impl ArrayPlugin for Pco { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -46,7 +38,7 @@ impl ArrayPlugin for Pco { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Pco plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -56,7 +48,7 @@ impl ArrayPlugin for Pco { ); let metadata = array.metadata.clone().encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -69,7 +61,7 @@ impl ArrayPlugin for Pco { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Pco plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/encodings/sequence/src/array/plugin.rs b/encodings/sequence/src/array/plugin.rs index 76aa9c0bb6f..46d864c65ac 100644 --- a/encodings/sequence/src/array/plugin.rs +++ b/encodings/sequence/src/array/plugin.rs @@ -34,14 +34,6 @@ impl ArrayPlugin for Sequence { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -49,7 +41,7 @@ impl ArrayPlugin for Sequence { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Sequence plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -59,7 +51,7 @@ impl ArrayPlugin for Sequence { }; Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata.encode_to_vec(), ))) @@ -72,7 +64,7 @@ impl ArrayPlugin for Sequence { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Sequence plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/encodings/sparse/src/plugin.rs b/encodings/sparse/src/plugin.rs index 12e649b55c2..3f3d55585d7 100644 --- a/encodings/sparse/src/plugin.rs +++ b/encodings/sparse/src/plugin.rs @@ -32,14 +32,6 @@ impl ArrayPlugin for Sparse { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -47,7 +39,7 @@ impl ArrayPlugin for Sparse { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Sparse plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -56,7 +48,7 @@ impl ArrayPlugin for Sparse { // Note that we DO NOT serialize the fill value since that is stored in the buffers. Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata.encode_to_vec(), ))) @@ -69,7 +61,7 @@ impl ArrayPlugin for Sparse { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Sparse plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/encodings/zigzag/src/array/plugin.rs b/encodings/zigzag/src/array/plugin.rs index 924e08f0ffd..60afb07e770 100644 --- a/encodings/zigzag/src/array/plugin.rs +++ b/encodings/zigzag/src/array/plugin.rs @@ -28,14 +28,6 @@ impl ArrayPlugin for ZigZag { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -43,11 +35,11 @@ impl ArrayPlugin for ZigZag { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "ZigZag plugin cannot serialize an array of another encoding", ); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array, vec![], ))) @@ -60,7 +52,7 @@ impl ArrayPlugin for ZigZag { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "ZigZag plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/encodings/zstd/src/array/plugin.rs b/encodings/zstd/src/array/plugin.rs index 58995f4d08e..a5c326a45bc 100644 --- a/encodings/zstd/src/array/plugin.rs +++ b/encodings/zstd/src/array/plugin.rs @@ -30,14 +30,6 @@ impl ArrayPlugin for Zstd { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -45,12 +37,12 @@ impl ArrayPlugin for Zstd { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Zstd plugin cannot serialize an array of another encoding", ); let array = array.as_::(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), array.metadata.clone().encode_to_vec(), ))) @@ -63,7 +55,7 @@ impl ArrayPlugin for Zstd { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Zstd plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/encodings/zstd/src/zstd_buffers/plugin.rs b/encodings/zstd/src/zstd_buffers/plugin.rs index 5fb0cbb4437..6517d64f54a 100644 --- a/encodings/zstd/src/zstd_buffers/plugin.rs +++ b/encodings/zstd/src/zstd_buffers/plugin.rs @@ -30,14 +30,6 @@ impl ArrayPlugin for ZstdBuffers { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -45,7 +37,7 @@ impl ArrayPlugin for ZstdBuffers { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "ZstdBuffers plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -68,7 +60,7 @@ impl ArrayPlugin for ZstdBuffers { } .encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -81,7 +73,7 @@ impl ArrayPlugin for ZstdBuffers { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "ZstdBuffers plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-array/src/array/plugin.rs b/vortex-array/src/array/plugin.rs index 7b7219297bd..d56943ecfc8 100644 --- a/vortex-array/src/array/plugin.rs +++ b/vortex-array/src/array/plugin.rs @@ -113,14 +113,22 @@ 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. /// diff --git a/vortex-array/src/arrays/bool/vtable/plugin.rs b/vortex-array/src/arrays/bool/vtable/plugin.rs index 73407213aed..07c089a0e27 100644 --- a/vortex-array/src/arrays/bool/vtable/plugin.rs +++ b/vortex-array/src/arrays/bool/vtable/plugin.rs @@ -29,14 +29,6 @@ impl ArrayPlugin for Bool { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -44,7 +36,7 @@ impl ArrayPlugin for Bool { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Bool plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -55,7 +47,7 @@ impl ArrayPlugin for Bool { } .encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -68,7 +60,7 @@ impl ArrayPlugin for Bool { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Bool plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-array/src/arrays/chunked/vtable/plugin.rs b/vortex-array/src/arrays/chunked/vtable/plugin.rs index e484881c6b8..c83114bf139 100644 --- a/vortex-array/src/arrays/chunked/vtable/plugin.rs +++ b/vortex-array/src/arrays/chunked/vtable/plugin.rs @@ -34,14 +34,6 @@ impl ArrayPlugin for Chunked { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -49,11 +41,11 @@ impl ArrayPlugin for Chunked { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Chunked plugin cannot serialize an array of another encoding", ); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array, vec![], ))) @@ -66,7 +58,7 @@ impl ArrayPlugin for Chunked { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Chunked plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-array/src/arrays/constant/vtable/plugin.rs b/vortex-array/src/arrays/constant/vtable/plugin.rs index 40740f57c3d..fbd965a6731 100644 --- a/vortex-array/src/arrays/constant/vtable/plugin.rs +++ b/vortex-array/src/arrays/constant/vtable/plugin.rs @@ -27,14 +27,6 @@ impl ArrayPlugin for Constant { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -42,13 +34,13 @@ impl ArrayPlugin for Constant { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Constant plugin cannot serialize an array of another encoding", ); // HACK: Because the scalar is stored in the buffers, we do not need to serialize the // metadata at all. Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array, vec![], ))) @@ -61,7 +53,7 @@ impl ArrayPlugin for Constant { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Constant plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs b/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs index 04925656780..324036e5b08 100644 --- a/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs @@ -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/vtable/plugin.rs b/vortex-array/src/arrays/dict/vtable/plugin.rs index 80af8ecf3ce..1487e256cd8 100644 --- a/vortex-array/src/arrays/dict/vtable/plugin.rs +++ b/vortex-array/src/arrays/dict/vtable/plugin.rs @@ -34,14 +34,6 @@ impl ArrayPlugin for Dict { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -49,7 +41,7 @@ impl ArrayPlugin for Dict { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Dict plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -66,7 +58,7 @@ impl ArrayPlugin for Dict { } .encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -79,7 +71,7 @@ impl ArrayPlugin for Dict { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Dict plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-array/src/arrays/extension/vtable/plugin.rs b/vortex-array/src/arrays/extension/vtable/plugin.rs index 3688feb0d97..f79a840264b 100644 --- a/vortex-array/src/arrays/extension/vtable/plugin.rs +++ b/vortex-array/src/arrays/extension/vtable/plugin.rs @@ -27,14 +27,6 @@ impl ArrayPlugin for Extension { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -42,11 +34,11 @@ impl ArrayPlugin for Extension { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Extension plugin cannot serialize an array of another encoding", ); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array, vec![], ))) @@ -59,7 +51,7 @@ impl ArrayPlugin for Extension { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Extension plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/plugin.rs b/vortex-array/src/arrays/fixed_size_list/vtable/plugin.rs index 046def3b94a..40ab024a769 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/plugin.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/plugin.rs @@ -28,14 +28,6 @@ impl ArrayPlugin for FixedSizeList { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -43,11 +35,11 @@ impl ArrayPlugin for FixedSizeList { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "FixedSizeList plugin cannot serialize an array of another encoding", ); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array, vec![], ))) @@ -60,7 +52,7 @@ impl ArrayPlugin for FixedSizeList { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "FixedSizeList plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-array/src/arrays/list/vtable/plugin.rs b/vortex-array/src/arrays/list/vtable/plugin.rs index 0d128c9b9c5..0d2b4a0770d 100644 --- a/vortex-array/src/arrays/list/vtable/plugin.rs +++ b/vortex-array/src/arrays/list/vtable/plugin.rs @@ -33,14 +33,6 @@ impl ArrayPlugin for List { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -48,7 +40,7 @@ impl ArrayPlugin for List { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "List plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -58,7 +50,7 @@ impl ArrayPlugin for List { } .encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -71,7 +63,7 @@ impl ArrayPlugin for List { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "List plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-array/src/arrays/listview/vtable/plugin.rs b/vortex-array/src/arrays/listview/vtable/plugin.rs index c29482d8ab5..784aed5c12e 100644 --- a/vortex-array/src/arrays/listview/vtable/plugin.rs +++ b/vortex-array/src/arrays/listview/vtable/plugin.rs @@ -34,14 +34,6 @@ impl ArrayPlugin for ListView { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -49,7 +41,7 @@ impl ArrayPlugin for ListView { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "ListView plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -60,7 +52,7 @@ impl ArrayPlugin for ListView { } .encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -73,7 +65,7 @@ impl ArrayPlugin for ListView { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "ListView plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-array/src/arrays/map/vtable/plugin.rs b/vortex-array/src/arrays/map/vtable/plugin.rs index c536f72985a..c6af458fdd0 100644 --- a/vortex-array/src/arrays/map/vtable/plugin.rs +++ b/vortex-array/src/arrays/map/vtable/plugin.rs @@ -29,14 +29,6 @@ impl ArrayPlugin for Map { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -44,11 +36,11 @@ impl ArrayPlugin for Map { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Map plugin cannot serialize an array of another encoding", ); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array, vec![], ))) @@ -61,7 +53,7 @@ impl ArrayPlugin for Map { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Map plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-array/src/arrays/masked/vtable/plugin.rs b/vortex-array/src/arrays/masked/vtable/plugin.rs index 0c62c450cd0..d0cf3560bbb 100644 --- a/vortex-array/src/arrays/masked/vtable/plugin.rs +++ b/vortex-array/src/arrays/masked/vtable/plugin.rs @@ -30,14 +30,6 @@ impl ArrayPlugin for Masked { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -45,11 +37,11 @@ impl ArrayPlugin for Masked { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Masked plugin cannot serialize an array of another encoding", ); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array, vec![], ))) @@ -62,7 +54,7 @@ impl ArrayPlugin for Masked { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Masked plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-array/src/arrays/null/plugin.rs b/vortex-array/src/arrays/null/plugin.rs index b7e93b6f37e..3b04ec22825 100644 --- a/vortex-array/src/arrays/null/plugin.rs +++ b/vortex-array/src/arrays/null/plugin.rs @@ -26,14 +26,6 @@ impl ArrayPlugin for Null { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -41,11 +33,11 @@ impl ArrayPlugin for Null { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Null plugin cannot serialize an array of another encoding", ); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array, vec![], ))) @@ -58,7 +50,7 @@ impl ArrayPlugin for Null { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Null plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-array/src/arrays/patched/vtable/plugin.rs b/vortex-array/src/arrays/patched/vtable/plugin.rs index 4ae9ea592ba..dd6a8280f49 100644 --- a/vortex-array/src/arrays/patched/vtable/plugin.rs +++ b/vortex-array/src/arrays/patched/vtable/plugin.rs @@ -30,14 +30,6 @@ impl ArrayPlugin for Patched { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -45,7 +37,7 @@ impl ArrayPlugin for Patched { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Patched plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -56,7 +48,7 @@ impl ArrayPlugin for Patched { } .encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -69,7 +61,7 @@ impl ArrayPlugin for Patched { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Patched plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-array/src/arrays/primitive/vtable/plugin.rs b/vortex-array/src/arrays/primitive/vtable/plugin.rs index e4aab246037..2757972b53f 100644 --- a/vortex-array/src/arrays/primitive/vtable/plugin.rs +++ b/vortex-array/src/arrays/primitive/vtable/plugin.rs @@ -29,14 +29,6 @@ impl ArrayPlugin for Primitive { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -44,11 +36,11 @@ impl ArrayPlugin for Primitive { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Primitive plugin cannot serialize an array of another encoding", ); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array, vec![], ))) @@ -61,7 +53,7 @@ impl ArrayPlugin for Primitive { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Primitive plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-array/src/arrays/runend/array/plugin.rs b/vortex-array/src/arrays/runend/array/plugin.rs index 039e3bceaba..3907726f882 100644 --- a/vortex-array/src/arrays/runend/array/plugin.rs +++ b/vortex-array/src/arrays/runend/array/plugin.rs @@ -34,14 +34,6 @@ impl ArrayPlugin for RunEnd { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -49,7 +41,7 @@ impl ArrayPlugin for RunEnd { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "RunEnd plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -61,7 +53,7 @@ impl ArrayPlugin for RunEnd { } .encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -74,7 +66,7 @@ impl ArrayPlugin for RunEnd { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "RunEnd plugin does not recognize the serialized ID", ); let ArrayDeserialization { 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/struct_/vtable/plugin.rs b/vortex-array/src/arrays/struct_/vtable/plugin.rs index 663db412e71..73125fbe18a 100644 --- a/vortex-array/src/arrays/struct_/vtable/plugin.rs +++ b/vortex-array/src/arrays/struct_/vtable/plugin.rs @@ -29,14 +29,6 @@ impl ArrayPlugin for Struct { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -44,11 +36,11 @@ impl ArrayPlugin for Struct { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Struct plugin cannot serialize an array of another encoding", ); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array, vec![], ))) @@ -61,7 +53,7 @@ impl ArrayPlugin for Struct { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Struct plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-array/src/arrays/union/vtable/plugin.rs b/vortex-array/src/arrays/union/vtable/plugin.rs index 3983e148dd0..6e15cdea14a 100644 --- a/vortex-array/src/arrays/union/vtable/plugin.rs +++ b/vortex-array/src/arrays/union/vtable/plugin.rs @@ -28,14 +28,6 @@ impl ArrayPlugin for Union { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -43,11 +35,11 @@ impl ArrayPlugin for Union { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Union plugin cannot serialize an array of another encoding", ); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array, vec![], ))) @@ -60,7 +52,7 @@ impl ArrayPlugin for Union { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Union plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-array/src/arrays/varbin/vtable/plugin.rs b/vortex-array/src/arrays/varbin/vtable/plugin.rs index 9c48280d40c..f163321f49a 100644 --- a/vortex-array/src/arrays/varbin/vtable/plugin.rs +++ b/vortex-array/src/arrays/varbin/vtable/plugin.rs @@ -34,14 +34,6 @@ impl ArrayPlugin for VarBin { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -49,7 +41,7 @@ impl ArrayPlugin for VarBin { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "VarBin plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -59,7 +51,7 @@ impl ArrayPlugin for VarBin { } .encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -72,7 +64,7 @@ impl ArrayPlugin for VarBin { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "VarBin plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-array/src/arrays/varbinview/vtable/plugin.rs b/vortex-array/src/arrays/varbinview/vtable/plugin.rs index ad9749a9383..34e1c9f6d39 100644 --- a/vortex-array/src/arrays/varbinview/vtable/plugin.rs +++ b/vortex-array/src/arrays/varbinview/vtable/plugin.rs @@ -33,14 +33,6 @@ impl ArrayPlugin for VarBinView { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -48,11 +40,11 @@ impl ArrayPlugin for VarBinView { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "VarBinView plugin cannot serialize an array of another encoding", ); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array, vec![], ))) @@ -65,7 +57,7 @@ impl ArrayPlugin for VarBinView { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "VarBinView plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-array/src/arrays/variant/vtable/plugin.rs b/vortex-array/src/arrays/variant/vtable/plugin.rs index 0ec40192759..3135c8eddc7 100644 --- a/vortex-array/src/arrays/variant/vtable/plugin.rs +++ b/vortex-array/src/arrays/variant/vtable/plugin.rs @@ -29,14 +29,6 @@ impl ArrayPlugin for Variant { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, array: &ArrayRef, @@ -44,7 +36,7 @@ impl ArrayPlugin for Variant { ) -> VortexResult> { vortex_ensure_eq!( array.encoding_id(), - VTable::id(self), + ArrayPlugin::id(self), "Variant plugin cannot serialize an array of another encoding", ); let array = array.as_::(); @@ -54,7 +46,7 @@ impl ArrayPlugin for Variant { .transpose()?; let metadata = VariantMetadataProto { shredded_dtype }.encode_to_vec(); Ok(Some(ArraySerialization::from_array( - VTable::id(self), + ArrayPlugin::id(self), array.array(), metadata, ))) @@ -67,7 +59,7 @@ impl ArrayPlugin for Variant { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - VTable::id(self), + ArrayPlugin::id(self), "Variant plugin does not recognize the serialized ID", ); let ArrayDeserialization { diff --git a/vortex-python/src/arrays/py/vtable/plugin.rs b/vortex-python/src/arrays/py/vtable/plugin.rs index 3b2f0ddb43d..e74401b0128 100644 --- a/vortex-python/src/arrays/py/vtable/plugin.rs +++ b/vortex-python/src/arrays/py/vtable/plugin.rs @@ -20,14 +20,6 @@ impl ArrayPlugin for PythonVTable { VTable::id(self) } - fn read_ids(&self) -> Vec { - vec![VTable::id(self)] - } - - fn write_ids(&self) -> Vec { - vec![VTable::id(self)] - } - fn serialize( &self, _array: &ArrayRef, From 157921a6ac73d27e38ab27ef962fbc94fb5f3bfb Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 6 Nov 2026 11:12:06 -0500 Subject: [PATCH 3/8] fix remaining vtable serde callers, unused imports and formatting Signed-off-by: Matt Katz --- encodings/alp/src/alp/plugin.rs | 5 +---- encodings/alp/src/alp_rd/array.rs | 21 +++++++++++------- encodings/bytebool/src/array.rs | 2 +- encodings/elias-fano/src/unary/tests.rs | 9 +------- encodings/fastlanes/src/rle/vtable/mod.rs | 18 ++++++++++----- encodings/parquet-variant/src/vtable.rs | 1 - encodings/zstd/src/zstd_buffers.rs | 2 +- vortex-array/src/array/vtable/mod.rs | 1 - .../src/arrays/constant/vtable/mod.rs | 2 -- vortex-array/src/arrays/interleave/mod.rs | 1 - vortex-array/src/arrays/list/vtable/mod.rs | 11 ++++++---- vortex-array/src/arrays/masked/vtable/mod.rs | 2 -- vortex-array/src/arrays/patched/vtable/mod.rs | 2 -- vortex-array/src/arrays/runend/array.rs | 22 ++++++++++--------- .../src/arrays/runend/array/plugin.rs | 4 ++-- vortex-array/src/arrays/slice/vtable.rs | 1 - vortex-array/src/arrays/varbin/vtable/mod.rs | 20 ++++++++++------- 17 files changed, 62 insertions(+), 62 deletions(-) diff --git a/encodings/alp/src/alp/plugin.rs b/encodings/alp/src/alp/plugin.rs index 4ae02ea455b..0f3db28bc23 100644 --- a/encodings/alp/src/alp/plugin.rs +++ b/encodings/alp/src/alp/plugin.rs @@ -17,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; @@ -56,9 +55,7 @@ impl ArrayPlugin for ALPPatchedPlugin { self.id(), "ALP plugin does not recognize serialized ID" ); - let alp_array = ArrayPlugin::deserialize(&ALP, parts, session)? - .try_downcast::() - .map_err(|_| vortex_err!("ALP plugin should only deserialize vortex.alp"))?; + let alp_array = ArrayPlugin::deserialize(&ALP, 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 47df00cce0d..3c1b3e327d0 100644 --- a/encodings/alp/src/alp_rd/array.rs +++ b/encodings/alp/src/alp_rd/array.rs @@ -523,6 +523,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; @@ -608,12 +609,16 @@ mod test { } .encode_to_vec(); let children: Vec = Vec::new(); - let result = ALPRD.deserialize( - &DType::Primitive(PType::F64, Nullability::NonNullable), - 0, - &metadata, - &[], - &children, + let result = ArrayPlugin::deserialize( + &ALPRD, + ArrayDeserialization::new( + VTable::id(&ALPRD), + &DType::Primitive(PType::F64, Nullability::NonNullable), + 0, + &metadata, + &[], + &children, + ), &SESSION, ); assert!(result.is_err_and(|err| { @@ -638,10 +643,10 @@ mod test { metadata: &ALPRDMetadata, children: Vec, ) -> VortexResult { - vortex_array::ArrayPlugin::deserialize( + ArrayPlugin::deserialize( &ALPRD, 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 4356712d944..d663ce6f1bd 100644 --- a/encodings/bytebool/src/array.rs +++ b/encodings/bytebool/src/array.rs @@ -36,7 +36,6 @@ use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; mod plugin; @@ -298,6 +297,7 @@ mod tests { use vortex_array::serde::SerializedArray; use vortex_array::session::ArraySessionExt; use vortex_buffer::ByteBufferMut; + use vortex_session::VortexSession; use vortex_session::registry::ReadContext; use super::*; diff --git a/encodings/elias-fano/src/unary/tests.rs b/encodings/elias-fano/src/unary/tests.rs index 993687eec07..3ddba58729e 100644 --- a/encodings/elias-fano/src/unary/tests.rs +++ b/encodings/elias-fano/src/unary/tests.rs @@ -294,14 +294,7 @@ fn test_deserialize_rejects() -> VortexResult<()> { let deserialize = |metadata: &[u8], buffers: &[BufferHandle], children: &[ArrayRef]| { ArrayPlugin::deserialize( &Unary, - ArrayDeserialization::new( - VTable::id(&Unary), - &dtype, - 2, - metadata, - buffers, - &children, - ), + ArrayDeserialization::new(VTable::id(&Unary), &dtype, 2, metadata, buffers, &children), &SESSION, ) }; diff --git a/encodings/fastlanes/src/rle/vtable/mod.rs b/encodings/fastlanes/src/rle/vtable/mod.rs index e1fec5b5476..c0e35854537 100644 --- a/encodings/fastlanes/src/rle/vtable/mod.rs +++ b/encodings/fastlanes/src/rle/vtable/mod.rs @@ -276,6 +276,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; @@ -323,12 +325,16 @@ mod tests { ) -> VortexResult<()> { 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, + let result = ArrayPlugin::deserialize( + &RLE, + 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/parquet-variant/src/vtable.rs b/encodings/parquet-variant/src/vtable.rs index 96b89a1ab33..a4d9fbe36c3 100644 --- a/encodings/parquet-variant/src/vtable.rs +++ b/encodings/parquet-variant/src/vtable.rs @@ -15,7 +15,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::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::child_to_validity; use vortex_error::VortexResult; diff --git a/encodings/zstd/src/zstd_buffers.rs b/encodings/zstd/src/zstd_buffers.rs index 0c916ac3bf1..0ee91a9964e 100644 --- a/encodings/zstd/src/zstd_buffers.rs +++ b/encodings/zstd/src/zstd_buffers.rs @@ -40,7 +40,6 @@ use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::ZstdBuffersMetadata; use crate::decompressor::with_decompressor; use crate::validate_frame_content_size; @@ -596,6 +595,7 @@ mod tests { use vortex_session::registry::ReadContext; use super::*; + use crate::ZstdBuffersMetadata; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/array/vtable/mod.rs b/vortex-array/src/array/vtable/mod.rs index 20bbb34a972..4fce425baac 100644 --- a/vortex-array/src/array/vtable/mod.rs +++ b/vortex-array/src/array/vtable/mod.rs @@ -27,7 +27,6 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; - use crate::Array; use crate::ArrayRef; use crate::ArrayView; diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index fcf4fcce79a..5a377b6d77e 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -10,7 +10,6 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::ArrayEq; @@ -49,7 +48,6 @@ use crate::match_each_native_ptype; use crate::match_each_varbin_builder; use crate::scalar::DecimalValue; use crate::scalar::ListScalar; -use crate::scalar::Scalar; use crate::scalar::ScalarValue; pub(crate) mod canonical; mod operations; diff --git a/vortex-array/src/arrays/interleave/mod.rs b/vortex-array/src/arrays/interleave/mod.rs index cdaf8edffb6..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; diff --git a/vortex-array/src/arrays/list/vtable/mod.rs b/vortex-array/src/arrays/list/vtable/mod.rs index 9e41261cba8..2c809f47391 100644 --- a/vortex-array/src/arrays/list/vtable/mod.rs +++ b/vortex-array/src/arrays/list/vtable/mod.rs @@ -4,14 +4,12 @@ use std::hash::Hasher; use std::sync::Arc; -use prost::Message; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_error::vortex_panic; -use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::ArrayEq; @@ -34,7 +32,6 @@ use crate::arrays::listview::list_view_from_list; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::dtype::DType; -use crate::dtype::Nullability; use crate::dtype::PType; use crate::match_each_list_builder; mod operations; @@ -183,6 +180,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; @@ -202,7 +201,11 @@ mod tests { } .encode_to_vec(); let children: Vec = Vec::new(); - let result = List.deserialize(&dtype, 0, &metadata, &[], &children, &SESSION); + let result = ArrayPlugin::deserialize( + &List, + 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/masked/vtable/mod.rs b/vortex-array/src/arrays/masked/vtable/mod.rs index 383e08979fd..51981fafb00 100644 --- a/vortex-array/src/arrays/masked/vtable/mod.rs +++ b/vortex-array/src/arrays/masked/vtable/mod.rs @@ -11,7 +11,6 @@ use vortex_error::VortexResult; 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; use crate::AnyCanonical; @@ -40,7 +39,6 @@ use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; use crate::require_child; use crate::scalar::Scalar; -use crate::validity::Validity; /// A [`Masked`]-encoded Vortex array. pub type MaskedArray = Array; diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs index e528b363205..fc7b6e73690 100644 --- a/vortex-array/src/arrays/patched/vtable/mod.rs +++ b/vortex-array/src/arrays/patched/vtable/mod.rs @@ -1,8 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use prost::Message; - use crate::ArrayEq; use crate::ArrayHash; mod kernels; diff --git a/vortex-array/src/arrays/runend/array.rs b/vortex-array/src/arrays/runend/array.rs index 5ce5210d5d7..bb6fd64d591 100644 --- a/vortex-array/src/arrays/runend/array.rs +++ b/vortex-array/src/arrays/runend/array.rs @@ -7,7 +7,6 @@ use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; -use prost::Message; use vortex_error::VortexExpect as _; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -15,7 +14,6 @@ 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; @@ -31,7 +29,6 @@ use crate::ExecutionCtx; use crate::ExecutionResult; use crate::IntoArray; use crate::TypedArrayRef; -use crate::VortexSessionExecute; use crate::array_slots; use crate::arrays::DecimalArray; use crate::arrays::ListViewArray; @@ -49,7 +46,6 @@ use crate::arrays::runend::ops::find_slice_end_index; use crate::arrays::runend::rules::RULES; use crate::buffer::BufferHandle; use crate::dtype::DType; -use crate::dtype::Nullability; use crate::dtype::PType; use crate::require_child; use crate::validity::Validity; @@ -547,6 +543,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; @@ -906,12 +904,16 @@ mod tests { } .encode_to_vec(); let children: Vec = Vec::new(); - let result = RunEnd.deserialize( - &DType::Primitive(PType::I32, Nullability::NonNullable), - 0, - &metadata, - &[], - &children, + let result = ArrayPlugin::deserialize( + &RunEnd, + 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/runend/array/plugin.rs b/vortex-array/src/arrays/runend/array/plugin.rs index 3907726f882..2de4227fd42 100644 --- a/vortex-array/src/arrays/runend/array/plugin.rs +++ b/vortex-array/src/arrays/runend/array/plugin.rs @@ -46,8 +46,8 @@ impl ArrayPlugin for RunEnd { ); let array = array.as_::(); let metadata = RunEndMetadata { - ends_ptype: PType::try_from(array.ends().dtype()) - .vortex_expect("Must be a valid PType") as i32, + 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, } diff --git a/vortex-array/src/arrays/slice/vtable.rs b/vortex-array/src/arrays/slice/vtable.rs index 7f65b5c5f3a..241dea96031 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -13,7 +13,6 @@ 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; diff --git a/vortex-array/src/arrays/varbin/vtable/mod.rs b/vortex-array/src/arrays/varbin/vtable/mod.rs index d9e3254cfd9..aee7951425c 100644 --- a/vortex-array/src/arrays/varbin/vtable/mod.rs +++ b/vortex-array/src/arrays/varbin/vtable/mod.rs @@ -3,7 +3,6 @@ use std::hash::Hasher; -use prost::Message; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -30,7 +29,6 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::VarBinViewBuilder; use crate::dtype::DType; -use crate::dtype::Nullability; use crate::dtype::PType; use crate::match_each_integer_ptype; use crate::match_each_varbin_builder; @@ -232,6 +230,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; @@ -243,12 +243,16 @@ mod tests { fn test_deserialize_rejects_unknown_offsets_ptype() -> VortexResult<()> { 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, + let result = ArrayPlugin::deserialize( + &VarBin, + ArrayDeserialization::new( + VTable::id(&VarBin), + &DType::Utf8(Nullability::NonNullable), + 0, + &metadata, + &[], + &children, + ), &SESSION, ); assert!(result.is_err_and(|err| { From 6e909c351f0c9915d8a0a0e5e863aa04b2660d16 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 6 Nov 2026 11:25:32 -0500 Subject: [PATCH 4/8] reword plugin docs now that bare vtables cannot register Signed-off-by: Matt Katz --- .../src/decimal_byte_parts/plugin/mod.rs | 4 ++-- encodings/fastlanes/src/bitpacking/plugin/mod.rs | 4 ++-- encodings/fastlanes/src/for_/plugin/mod.rs | 4 ++-- vortex-array/src/array/plugin.rs | 9 +++++++-- vortex-array/src/arrays/decimal/vtable/mod.rs | 6 ++---- vortex-array/src/arrays/decimal/vtable/plugin/mod.rs | 2 +- 6 files changed, 16 insertions(+), 13 deletions(-) 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/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/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/vortex-array/src/array/plugin.rs b/vortex-array/src/array/plugin.rs index d56943ecfc8..33193cea177 100644 --- a/vortex-array/src/array/plugin.rs +++ b/vortex-array/src/array/plugin.rs @@ -132,7 +132,10 @@ pub trait ArrayPlugin: 'static + Send + Sync { /// Serialize `array` to its wire representation. /// - /// This function is called only for arrays whose in-memory encoding matches [`id`](Self::id). + /// The session calls this only for arrays whose in-memory encoding matches [`id`](Self::id). + /// Implementations in this crate check that as well, so a direct caller gets an error rather + /// than a panic from the downcast. + /// /// The returned ID must be declared by [`write_ids`](Self::write_ids). Return /// `Ok(None)` when the array cannot be serialized. fn serialize( @@ -143,7 +146,9 @@ 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 + /// `serialized_id` identifies the exact representation encountered on disk. The session only + /// routes IDs from [`read_ids`](Self::read_ids) here; implementations in this crate check it as + /// well, so a direct caller gets an error rather than a misread 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 diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index e21adb51186..1985468489b 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -3,8 +3,7 @@ //! The [`Decimal`] vtable and the plugin that owns its wire form. //! -//! The vtable does not serialize or deserialize. [`DecimalPlugin`] reads and writes the -//! `vortex.decimal` wire form, including legacy nodes. +//! [`DecimalPlugin`] reads and writes the `vortex.decimal` wire form, including legacy nodes. mod kernel; mod operations; @@ -160,7 +159,6 @@ 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. +/// A session **must** register [`DecimalPlugin`] to read or write decimal arrays. #[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 324036e5b08..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 From 542ad683d9b5e16135af60694fefafa1ccdb59ba Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 6 Nov 2026 11:31:30 -0500 Subject: [PATCH 5/8] implement ArrayPlugin for Decimal instead of a DecimalPlugin struct Signed-off-by: Matt Katz --- vortex-array/src/arrays/decimal/mod.rs | 6 +-- vortex-array/src/arrays/decimal/vtable/mod.rs | 10 +---- .../src/arrays/decimal/vtable/plugin/mod.rs | 42 ++++++++----------- .../arrays/decimal/vtable/plugin/stored.rs | 5 ++- .../arrays/decimal/vtable/plugin/tests/mod.rs | 7 ++-- vortex-array/src/session/mod.rs | 4 +- 6 files changed, 29 insertions(+), 45 deletions(-) diff --git a/vortex-array/src/arrays/decimal/mod.rs b/vortex-array/src/arrays/decimal/mod.rs index 82cd8f95993..3518ab99fdf 100644 --- a/vortex-array/src/arrays/decimal/mod.rs +++ b/vortex-array/src/arrays/decimal/mod.rs @@ -3,9 +3,8 @@ //! Decimal arrays whose logical precision determines their canonical integer width. //! -//! [`DecimalPlugin`] reads and writes the `vortex.decimal` wire form in place of the [`Decimal`] -//! vtable. It reads a legacy node, stored at another width, as the canonical array or as a -//! [`Widen`]. +//! [`Decimal`] reads and writes the `vortex.decimal` wire form. It reads a legacy node, stored at +//! another width, as the canonical array or as a [`Widen`]. //! //! [`Widen`]: crate::arrays::Widen @@ -22,7 +21,6 @@ pub use compute::rules::DecimalMaskedValidityRule; mod vtable; pub use vtable::Decimal; pub use vtable::DecimalArray; -pub use vtable::DecimalPlugin; mod utils; pub use utils::*; diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index 1985468489b..24c2e892a5c 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -1,16 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! The [`Decimal`] vtable and the plugin that owns its wire form. -//! -//! [`DecimalPlugin`] reads and writes the `vortex.decimal` wire form, including legacy nodes. +//! The [`Decimal`] vtable. Its plugin reads and writes the `vortex.decimal` wire form, including +//! legacy nodes. mod kernel; mod operations; - mod plugin; -pub use plugin::DecimalPlugin; - mod validity; use std::hash::Hash; @@ -158,7 +154,5 @@ impl VTable for Decimal { } /// The vtable of [`DecimalArray`]. -/// -/// A session **must** register [`DecimalPlugin`] to read or write decimal arrays. #[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 8fc4bfbd29b..734fcf9be5e 100644 --- a/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs @@ -1,14 +1,22 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Serde for `vortex.decimal` nodes, including legacy nodes. +//! [`ArrayPlugin`] implementation for [`Decimal`], including legacy nodes. //! -//! [`StoredDecimal::is_legacy`] defines a legacy node, which [`DecimalPlugin`] reads as the -//! canonical array or a [`Widen`]. +//! [`Decimal`] writes a [`DecimalArray`] as a `vortex.decimal` node, and reads every such node as +//! the canonical decimal array or a decimal [`Widen`]. [`StoredDecimal::is_legacy`] defines a +//! legacy node. +//! +//! A canonical node keeps its buffer without a copy or a scan, so the buffer can stay on a device, +//! and its valid values are trusted to fit the precision. A legacy node goes through +//! [`WidenArray::from_stored_decimal`], which scans its valid values on the host, so a legacy node +//! in device memory fails to read. //! //! Each legacy read scans the valid values once, because the deserializer attaches stored //! statistics only after the plugin returns. This cost is not limited to old files, because IPC //! and writes without a compressor still write a decimal `Widen` as a legacy node. +//! +//! [`DecimalArray`]: super::DecimalArray mod stored; @@ -58,25 +66,9 @@ impl DecimalMetadata { } } -/// Writes a [`DecimalArray`] as a `vortex.decimal` node, and reads every such node as the -/// canonical decimal array or a decimal [`Widen`]. -/// -/// A canonical node keeps its buffer without a copy or a scan, so the buffer can stay on a device, -/// and its valid values are trusted to fit the precision. A legacy node goes through -/// [`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 implement [`ArrayPlugin`], so the default [`ArraySession`] -/// registers this plugin in its place. -/// -/// [`DecimalArray`]: super::DecimalArray -/// [`ArraySession`]: crate::session::ArraySession -#[derive(Clone, Debug)] -pub struct DecimalPlugin; - -impl ArrayPlugin for DecimalPlugin { +impl ArrayPlugin for Decimal { fn id(&self) -> ArrayId { - VTable::id(&Decimal) + VTable::id(self) } fn serialize( @@ -86,7 +78,7 @@ impl ArrayPlugin for DecimalPlugin { ) -> VortexResult> { let Some(decimal) = array.as_opt::() else { vortex_bail!( - "DecimalPlugin serializes only Decimal arrays, got {}", + "Decimal plugin serializes only Decimal arrays, got {}", array.encoding_id() ); }; @@ -95,7 +87,7 @@ impl ArrayPlugin for DecimalPlugin { }; Ok(Some(ArraySerialization::from_array( - self.id(), + ArrayPlugin::id(self), array, metadata.encode_to_vec(), ))) @@ -108,8 +100,8 @@ impl ArrayPlugin for DecimalPlugin { ) -> VortexResult { vortex_ensure_eq!( parts.serialized_id, - self.id(), - "DecimalPlugin reads only its own serialized ID", + ArrayPlugin::id(self), + "Decimal plugin reads only its own serialized ID", ); let stored = StoredDecimal::deserialize( diff --git a/vortex-array/src/arrays/decimal/vtable/plugin/stored.rs b/vortex-array/src/arrays/decimal/vtable/plugin/stored.rs index efebdf800b3..5b8d3863a0d 100644 --- a/vortex-array/src/arrays/decimal/vtable/plugin/stored.rs +++ b/vortex-array/src/arrays/decimal/vtable/plugin/stored.rs @@ -1,12 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Validation of a `vortex.decimal` node before [`DecimalPlugin`] constructs an array in memory. +//! Validation of a `vortex.decimal` node before the [`Decimal`] plugin constructs an array in +//! memory. //! //! A malformed row count or buffer layout returns an error here instead of failing an assert //! later. //! -//! [`DecimalPlugin`]: super::DecimalPlugin +//! [`Decimal`]: super::Decimal use prost::Message; use vortex_error::VortexResult; diff --git a/vortex-array/src/arrays/decimal/vtable/plugin/tests/mod.rs b/vortex-array/src/arrays/decimal/vtable/plugin/tests/mod.rs index c3000dd962f..4a82c579744 100644 --- a/vortex-array/src/arrays/decimal/vtable/plugin/tests/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/plugin/tests/mod.rs @@ -22,7 +22,6 @@ use vortex_session::VortexSession; use vortex_session::registry::ReadContext; use super::DecimalMetadata; -use super::DecimalPlugin; use crate::ArrayContext; use crate::ArrayRef; use crate::IntoArray; @@ -85,7 +84,7 @@ fn read_node( Ok(ReadNode { array, stored_ptr }) } -/// Deserializes the parts of a node through [`DecimalPlugin`]. +/// Deserializes the parts of a node through the [`Decimal`] plugin. fn deserialize( dtype: &DType, len: usize, @@ -93,10 +92,10 @@ fn deserialize( buffers: &[BufferHandle], children: &[ArrayRef], ) -> VortexResult { - let id = DecimalPlugin.id(); + let id = ArrayPlugin::id(&Decimal); let parts = ArrayDeserialization::new(id, dtype, len, metadata, buffers, &children); - DecimalPlugin.deserialize(parts, &SESSION) + ArrayPlugin::deserialize(&Decimal, parts, &SESSION) } /// Serializes `array` through the session, then decodes it with the dtype `dtype`. diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index 85c61191de7..17f367db8fc 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -27,6 +27,7 @@ use crate::array::ArraySerialization; use crate::arrays::Bool; use crate::arrays::Chunked; use crate::arrays::Constant; +use crate::arrays::Decimal; use crate::arrays::Dict; use crate::arrays::Extension; use crate::arrays::FixedSizeList; @@ -42,7 +43,6 @@ use crate::arrays::Union; use crate::arrays::VarBin; use crate::arrays::VarBinView; use crate::arrays::Variant; -use crate::arrays::decimal::DecimalPlugin; use crate::arrays::widen::WidenPlugin; /// Registry of array encodings. @@ -96,7 +96,7 @@ impl Default for ArraySession { this.register(Null); this.register(Bool); this.register(Primitive); - this.register(DecimalPlugin); + this.register(Decimal); this.register(VarBinView); this.register(ListView); this.register(Map); From c2c0895e25b72e897c0f46cf36f10425eed0ac4b Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 6 Nov 2026 11:58:36 -0500 Subject: [PATCH 6/8] address review: alp downcast error, dict safety note, vtable guide Signed-off-by: Matt Katz --- docs/developer-guide/internals/vtables.md | 8 ++++++-- encodings/alp/src/alp/plugin.rs | 5 ++++- vortex-array/src/arrays/dict/vtable/plugin.rs | 3 +++ 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/docs/developer-guide/internals/vtables.md b/docs/developer-guide/internals/vtables.md index 77283ea012d..e05c7b1270e 100644 --- a/docs/developer-guide/internals/vtables.md +++ b/docs/developer-guide/internals/vtables.md @@ -121,7 +121,9 @@ needs: But operations like deserialization need to construct a new instance from raw bytes -- there is no instance yet. Vortex vtables colocate instance logic (e.g. `serialize`) with type-level logic (e.g. `deserialize`) in the same struct, since the vtable itself is a lightweight value -(often zero-sized) that can be default-constructed. +(often zero-sized) that can be default-constructed. Array vtables are the exception: their serde +lives on a separate `ArrayPlugin` implementation, usually on the vtable type itself, so a vtable +without one is in-memory only. **No associated type access.** A `dyn Trait` erases all associated types. With Vortex vtables, the typed form `Foo` preserves full access to `V::Metadata` and other associated types, @@ -176,7 +178,9 @@ Methods that exist only for erased dispatch (e.g. `as_any`, `metadata_any`, ## Registration and Deserialization -Vtables are registered in the session by their ID. When a serialized value is encountered +Vtables are registered in the session by their ID. An array vtable registers through its +`ArrayPlugin` implementation, and a vtable without one cannot be registered. When a serialized +value is encountered (e.g. an extension dtype in a file footer), the session's registry resolves the ID to the `FooPlugin` and calls `deserialize` to reconstruct the value. The result is returned as the erased form (`FooRef`) so it can be stored without generic parameters. diff --git a/encodings/alp/src/alp/plugin.rs b/encodings/alp/src/alp/plugin.rs index 0f3db28bc23..fe438069314 100644 --- a/encodings/alp/src/alp/plugin.rs +++ b/encodings/alp/src/alp/plugin.rs @@ -17,6 +17,7 @@ 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; @@ -55,7 +56,9 @@ impl ArrayPlugin for ALPPatchedPlugin { self.id(), "ALP plugin does not recognize serialized ID" ); - let alp_array = ArrayPlugin::deserialize(&ALP, parts, session)?.downcast::(); + let alp_array = ArrayPlugin::deserialize(&ALP, parts, session)? + .try_downcast::() + .map_err(|_| vortex_err!("ALP plugin returned an array that is not vortex.alp"))?; // Check if there are interior patches to externalize. let Some(patches) = alp_array.patches() else { diff --git a/vortex-array/src/arrays/dict/vtable/plugin.rs b/vortex-array/src/arrays/dict/vtable/plugin.rs index 1487e256cd8..d7559985911 100644 --- a/vortex-array/src/arrays/dict/vtable/plugin.rs +++ b/vortex-array/src/arrays/dict/vtable/plugin.rs @@ -104,6 +104,9 @@ impl ArrayPlugin for Dict { self.clone(), dtype.clone(), len, + // SAFETY: `Array::try_from_parts` validates the codes and values against the dtype and + // length before the array is exposed. The `all_values_referenced` flag is trusted from + // the metadata, as it was before serde moved here. unsafe { DictData::new_unchecked().set_all_values_referenced(all_values_referenced) }, smallvec![Some(codes), Some(values)], ))? From 29697edd249409528284f1d5eaa27cbdfce950d5 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 6 Nov 2026 12:38:27 -0500 Subject: [PATCH 7/8] state the dict safety invariant and reflow the vtable guide Signed-off-by: Matt Katz --- docs/developer-guide/internals/vtables.md | 7 +++---- vortex-array/src/arrays/dict/vtable/plugin.rs | 4 ++-- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/docs/developer-guide/internals/vtables.md b/docs/developer-guide/internals/vtables.md index e05c7b1270e..873ad57d524 100644 --- a/docs/developer-guide/internals/vtables.md +++ b/docs/developer-guide/internals/vtables.md @@ -180,10 +180,9 @@ Methods that exist only for erased dispatch (e.g. `as_any`, `metadata_any`, Vtables are registered in the session by their ID. An array vtable registers through its `ArrayPlugin` implementation, and a vtable without one cannot be registered. When a serialized -value is encountered -(e.g. an extension dtype in a file footer), the session's registry resolves the ID to the -`FooPlugin` and calls `deserialize` to reconstruct the value. The result is returned as the -erased form (`FooRef`) so it can be stored without generic parameters. +value is encountered (e.g. an extension dtype in a file footer), the session's registry resolves +the ID to the `FooPlugin` and calls `deserialize` to reconstruct the value. The result is returned +as the erased form (`FooRef`) so it can be stored without generic parameters. This pattern -- register plugin by ID, deserialize via plugin, store as erased ref -- is consistent across all vtable-backed types in Vortex. diff --git a/vortex-array/src/arrays/dict/vtable/plugin.rs b/vortex-array/src/arrays/dict/vtable/plugin.rs index d7559985911..742f6c8c6d7 100644 --- a/vortex-array/src/arrays/dict/vtable/plugin.rs +++ b/vortex-array/src/arrays/dict/vtable/plugin.rs @@ -105,8 +105,8 @@ impl ArrayPlugin for Dict { dtype.clone(), len, // SAFETY: `Array::try_from_parts` validates the codes and values against the dtype and - // length before the array is exposed. The `all_values_referenced` flag is trusted from - // the metadata, as it was before serde moved here. + // length before the array is exposed. `all_values_referenced` is the writer's claim + // about the array it serialized, carried in the metadata. unsafe { DictData::new_unchecked().set_all_values_referenced(all_values_referenced) }, smallvec![Some(codes), Some(values)], ))? From 503854f6a3c4eb53b3c6637eece57f927e24fe63 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 6 Nov 2026 12:38:27 -0500 Subject: [PATCH 8/8] use directory modules for vtables with a plugin and merge the alp plugins Signed-off-by: Matt Katz --- .../alp/src/alp/{array.rs => array/mod.rs} | 1 + encodings/alp/src/alp/array/plugin.rs | 232 +++++++++++++++- encodings/alp/src/alp/mod.rs | 3 +- encodings/alp/src/alp/plugin.rs | 248 ------------------ .../alp/src/alp_rd/{array.rs => array/mod.rs} | 0 .../bytebool/src/{array.rs => array/mod.rs} | 0 .../src/{array.rs => array/mod.rs} | 0 encodings/fsst/src/{array.rs => array/mod.rs} | 0 .../onpair/src/{array.rs => array/mod.rs} | 0 .../src/{vtable.rs => vtable/mod.rs} | 0 encodings/pco/src/{array.rs => array/mod.rs} | 0 .../sequence/src/{array.rs => array/mod.rs} | 0 .../zigzag/src/{array.rs => array/mod.rs} | 0 encodings/zstd/src/{array.rs => array/mod.rs} | 0 .../{zstd_buffers.rs => zstd_buffers/mod.rs} | 0 .../arrays/runend/{array.rs => array/mod.rs} | 0 .../arrays/py/{vtable.rs => vtable/mod.rs} | 0 17 files changed, 233 insertions(+), 251 deletions(-) rename encodings/alp/src/alp/{array.rs => array/mod.rs} (99%) delete mode 100644 encodings/alp/src/alp/plugin.rs rename encodings/alp/src/alp_rd/{array.rs => array/mod.rs} (100%) rename encodings/bytebool/src/{array.rs => array/mod.rs} (100%) rename encodings/datetime-parts/src/{array.rs => array/mod.rs} (100%) rename encodings/fsst/src/{array.rs => array/mod.rs} (100%) rename encodings/onpair/src/{array.rs => array/mod.rs} (100%) rename encodings/parquet-variant/src/{vtable.rs => vtable/mod.rs} (100%) rename encodings/pco/src/{array.rs => array/mod.rs} (100%) rename encodings/sequence/src/{array.rs => array/mod.rs} (100%) rename encodings/zigzag/src/{array.rs => array/mod.rs} (100%) rename encodings/zstd/src/{array.rs => array/mod.rs} (100%) rename encodings/zstd/src/{zstd_buffers.rs => zstd_buffers/mod.rs} (100%) rename vortex-array/src/arrays/runend/{array.rs => array/mod.rs} (100%) rename vortex-python/src/arrays/py/{vtable.rs => vtable/mod.rs} (100%) diff --git a/encodings/alp/src/alp/array.rs b/encodings/alp/src/alp/array/mod.rs similarity index 99% rename from encodings/alp/src/alp/array.rs rename to encodings/alp/src/alp/array/mod.rs index 48967c719f4..3ac70d6871d 100644 --- a/encodings/alp/src/alp/array.rs +++ b/encodings/alp/src/alp/array/mod.rs @@ -50,6 +50,7 @@ use crate::alp::decompress::execute_decompress; use crate::alp::rules::RULES; mod plugin; +pub(crate) use plugin::ALPPatchedPlugin; /// A [`ALP`]-encoded Vortex array. pub type ALPArray = Array; diff --git a/encodings/alp/src/alp/array/plugin.rs b/encodings/alp/src/alp/array/plugin.rs index b4af77efd67..a50f2467299 100644 --- a/encodings/alp/src/alp/array/plugin.rs +++ b/encodings/alp/src/alp/array/plugin.rs @@ -1,7 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! [`ArrayPlugin`] implementation for [`ALP`]. +//! [`ArrayPlugin`] implementations for [`ALP`]. +//! +//! [`ALP`] itself reads and writes the current `vortex.alp` form. [`ALPPatchedPlugin`] reads an +//! ALP node with interior patches as a `PatchedArray` wrapping a patchless ALP array, which keeps +//! previously written datasets readable at no cost. use prost::Message; use vortex_array::Array; @@ -14,16 +18,19 @@ use vortex_array::ArraySerialization; use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Patched; use vortex_array::dtype::DType; use vortex_array::dtype::PType; use vortex_array::patches::Patches; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; use vortex_session::VortexSession; use super::ALP; use super::ALPArrayExt; +use super::ALPArrayOwnedExt; use super::ALPData; use super::ALPMetadata; use crate::alp::Exponents; @@ -125,3 +132,226 @@ impl ArrayPlugin for ALP { .into_array()) } } + +/// Custom deserialization plugin that converts an ALP array with interior +/// patches into a PatchedArray holding an ALP array. +#[derive(Debug, Clone)] +pub(crate) struct ALPPatchedPlugin; + +impl ArrayPlugin for ALPPatchedPlugin { + fn id(&self) -> ArrayId { + // We reuse the existing `ALP` ID so that we can take over its + // deserialization pathway. + // TODO(joe): dedup method name + VTable::id(&ALP) + } + + fn serialize( + &self, + array: &ArrayRef, + session: &VortexSession, + ) -> VortexResult> { + // Delegate to ALP's metadata serde + ArrayPlugin::serialize(&ALP, array, session) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + self.id(), + "ALP plugin does not recognize serialized ID" + ); + let alp_array = ArrayPlugin::deserialize(&ALP, parts, session)? + .try_downcast::() + .map_err(|_| vortex_err!("ALP plugin returned an array that is not vortex.alp"))?; + + // Check if there are interior patches to externalize. + let Some(patches) = alp_array.patches() else { + return Ok(alp_array.into_array()); + }; + + // Extract components and create a new ALP array without patches. + let (encoded, exponents, _) = alp_array.into_parts(); + + let alp_without_patches = ALP::try_new(encoded, exponents, None)?.into_array(); + + let patched = Patched::from_array_and_patches( + alp_without_patches, + &patches, + &mut session.create_execution_ctx(), + )?; + + Ok(patched.into_array()) + } + + fn is_supported_encoding(&self, id: &ArrayId) -> bool { + // TODO(joe): dedup method name + id == VTable::id(&Patched) || id == VTable::id(&ALP) + } +} + +#[cfg(test)] +mod tests { + use std::f64::consts::PI; + use std::sync::LazyLock; + + use vortex_array::ArrayDeserialization; + use vortex_array::ArrayPlugin; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PatchedArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::patched::PatchedArraySlotsExt; + use vortex_array::assert_arrays_eq; + use vortex_array::buffer::BufferHandle; + use vortex_array::session::ArraySessionExt; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use vortex_session::VortexSession; + + use super::ALPPatchedPlugin; + use crate::ALP; + use crate::ALPArray; + use crate::ALPArrayExt; + use crate::alp_encode; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + session.arrays().register(ALPPatchedPlugin); + session + }); + + #[test] + fn test_decode_alp_patches() -> VortexResult<()> { + // Create values where some don't encode cleanly with ALP, causing patches. + // PI doesn't encode cleanly. + let values: Vec = (0..100) + .map(|i| if i % 4 == 3 { PI } else { i as f64 }) + .collect(); + + let parray = PrimitiveArray::from_iter(values); + let alp_encoded = alp_encode(parray.as_view(), None, &mut SESSION.create_execution_ctx())?; + + assert!( + alp_encoded.patches().is_some(), + "Expected ALP array to have patches" + ); + + let array = alp_encoded.as_array(); + + let serialization = SESSION.array_serialize(array)?.unwrap(); + let children = array.children(); + let buffers = array + .buffers() + .into_iter() + .map(BufferHandle::new_host) + .collect::>(); + + let deserialized = ALPPatchedPlugin.deserialize( + ArrayDeserialization::new( + ALPPatchedPlugin.id(), + array.dtype(), + array.len(), + &serialization.metadata, + &buffers, + &children, + ), + &SESSION, + )?; + + assert_arrays_eq!(deserialized, parray, &mut SESSION.create_execution_ctx()); + + let patched: PatchedArray = deserialized + .try_downcast() + .map_err(|a| vortex_err!("Expected Patched, got {}", a.encoding_id()))?; + + let inner_alp: ALPArray = patched + .inner() + .clone() + .try_downcast() + .map_err(|a| vortex_err!("Expected inner ALP, got {}", a.encoding_id()))?; + + assert!( + inner_alp.patches().is_none(), + "Inner ALP should NOT have patches" + ); + + Ok(()) + } + + #[test] + fn alp_without_patches_stays_alp() -> VortexResult<()> { + // Values that encode cleanly without patches. + let values: Vec = (0..100).map(|i| i as f64).collect(); + let parray = PrimitiveArray::from_iter(values); + let alp_encoded = alp_encode(parray.as_view(), None, &mut SESSION.create_execution_ctx())?; + + assert!( + alp_encoded.patches().is_none(), + "Expected ALP array without patches" + ); + + let array = alp_encoded.as_array(); + + let serialization = SESSION.array_serialize(array)?.unwrap(); + let children = array.children(); + let buffers = array + .buffers() + .into_iter() + .map(BufferHandle::new_host) + .collect::>(); + + let deserialized = ALPPatchedPlugin.deserialize( + ArrayDeserialization::new( + ALPPatchedPlugin.id(), + array.dtype(), + array.len(), + &serialization.metadata, + &buffers, + &children, + ), + &SESSION, + )?; + + assert_arrays_eq!(deserialized, parray, &mut SESSION.create_execution_ctx()); + + let result = deserialized + .try_downcast::() + .map_err(|a| vortex_err!("Expected deserialized ALP, got {}", a.encoding_id()))?; + + assert!(result.patches().is_none(), "Result should not have patches"); + + Ok(()) + } + + #[test] + #[should_panic(expected = "index out of bounds")] + fn primitive_array_returns_error() { + let array = PrimitiveArray::from_iter([1.0f64, 2.0, 3.0]).into_array(); + + let serialization = SESSION.array_serialize(&array).unwrap().unwrap(); + let children = array.children(); + let buffers = array + .buffers() + .into_iter() + .map(BufferHandle::new_host) + .collect::>(); + + // This panics because PrimitiveArray has no children and ALP requires encoded child. + let _result = ALPPatchedPlugin.deserialize( + ArrayDeserialization::new( + ALPPatchedPlugin.id(), + array.dtype(), + array.len(), + &serialization.metadata, + &buffers, + &children, + ), + &SESSION, + ); + } +} diff --git a/encodings/alp/src/alp/mod.rs b/encodings/alp/src/alp/mod.rs index ea6f03ad038..cb1d5c2c64f 100644 --- a/encodings/alp/src/alp/mod.rs +++ b/encodings/alp/src/alp/mod.rs @@ -15,10 +15,9 @@ mod compress; pub(crate) mod compute; mod decompress; mod ops; -mod plugin; mod rules; -pub(crate) use plugin::ALPPatchedPlugin; +pub(crate) use array::ALPPatchedPlugin; pub(crate) fn initialize(session: &VortexSession) { rules::initialize(session); diff --git a/encodings/alp/src/alp/plugin.rs b/encodings/alp/src/alp/plugin.rs deleted file mode 100644 index fe438069314..00000000000 --- a/encodings/alp/src/alp/plugin.rs +++ /dev/null @@ -1,248 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! A custom [`ArrayPlugin`] that lets you load in and deserialize an `ALP` array with interior -//! patches as a `PatchedArray` that wraps a patchless `ALP` array. -//! -//! This enables zero-cost backward compatibility with previously written datasets. - -use vortex_array::ArrayDeserialization; -use vortex_array::ArrayId; -use vortex_array::ArrayPlugin; -use vortex_array::ArrayRef; -use vortex_array::ArraySerialization; -use vortex_array::ArrayVTable; -use vortex_array::IntoArray; -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; -use crate::ALPArrayExt; -use crate::ALPArrayOwnedExt; - -/// Custom deserialization plugin that converts an ALP array with interior -/// patches into a PatchedArray holding an ALP array. -#[derive(Debug, Clone)] -pub(crate) struct ALPPatchedPlugin; - -impl ArrayPlugin for ALPPatchedPlugin { - fn id(&self) -> ArrayId { - // We reuse the existing `ALP` ID so that we can take over its - // deserialization pathway. - // TODO(joe): dedup method name - ArrayVTable::id(&ALP) - } - - fn serialize( - &self, - array: &ArrayRef, - session: &VortexSession, - ) -> VortexResult> { - // Delegate to ALP's metadata serde - ArrayPlugin::serialize(&ALP, array, session) - } - - fn deserialize( - &self, - parts: ArrayDeserialization<'_>, - session: &VortexSession, - ) -> VortexResult { - vortex_ensure_eq!( - parts.serialized_id, - self.id(), - "ALP plugin does not recognize serialized ID" - ); - let alp_array = ArrayPlugin::deserialize(&ALP, parts, session)? - .try_downcast::() - .map_err(|_| vortex_err!("ALP plugin returned an array that is not vortex.alp"))?; - - // Check if there are interior patches to externalize. - let Some(patches) = alp_array.patches() else { - return Ok(alp_array.into_array()); - }; - - // Extract components and create a new ALP array without patches. - let (encoded, exponents, _) = alp_array.into_parts(); - - let alp_without_patches = ALP::try_new(encoded, exponents, None)?.into_array(); - - let patched = Patched::from_array_and_patches( - alp_without_patches, - &patches, - &mut session.create_execution_ctx(), - )?; - - Ok(patched.into_array()) - } - - fn is_supported_encoding(&self, id: &ArrayId) -> bool { - // TODO(joe): dedup method name - id == ArrayVTable::id(&Patched) || id == ArrayVTable::id(&ALP) - } -} - -#[cfg(test)] -mod tests { - use std::f64::consts::PI; - use std::sync::LazyLock; - - use vortex_array::ArrayDeserialization; - use vortex_array::ArrayPlugin; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::PatchedArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::patched::PatchedArraySlotsExt; - use vortex_array::assert_arrays_eq; - use vortex_array::buffer::BufferHandle; - use vortex_array::session::ArraySessionExt; - use vortex_error::VortexResult; - use vortex_error::vortex_err; - use vortex_session::VortexSession; - - use super::ALPPatchedPlugin; - use crate::ALP; - use crate::ALPArray; - use crate::ALPArrayExt; - use crate::alp_encode; - - static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty(); - session.arrays().register(ALPPatchedPlugin); - session - }); - - #[test] - fn test_decode_alp_patches() -> VortexResult<()> { - // Create values where some don't encode cleanly with ALP, causing patches. - // PI doesn't encode cleanly. - let values: Vec = (0..100) - .map(|i| if i % 4 == 3 { PI } else { i as f64 }) - .collect(); - - let parray = PrimitiveArray::from_iter(values); - let alp_encoded = alp_encode(parray.as_view(), None, &mut SESSION.create_execution_ctx())?; - - assert!( - alp_encoded.patches().is_some(), - "Expected ALP array to have patches" - ); - - let array = alp_encoded.as_array(); - - let serialization = SESSION.array_serialize(array)?.unwrap(); - let children = array.children(); - let buffers = array - .buffers() - .into_iter() - .map(BufferHandle::new_host) - .collect::>(); - - let deserialized = ALPPatchedPlugin.deserialize( - ArrayDeserialization::new( - ALPPatchedPlugin.id(), - array.dtype(), - array.len(), - &serialization.metadata, - &buffers, - &children, - ), - &SESSION, - )?; - - assert_arrays_eq!(deserialized, parray, &mut SESSION.create_execution_ctx()); - - let patched: PatchedArray = deserialized - .try_downcast() - .map_err(|a| vortex_err!("Expected Patched, got {}", a.encoding_id()))?; - - let inner_alp: ALPArray = patched - .inner() - .clone() - .try_downcast() - .map_err(|a| vortex_err!("Expected inner ALP, got {}", a.encoding_id()))?; - - assert!( - inner_alp.patches().is_none(), - "Inner ALP should NOT have patches" - ); - - Ok(()) - } - - #[test] - fn alp_without_patches_stays_alp() -> VortexResult<()> { - // Values that encode cleanly without patches. - let values: Vec = (0..100).map(|i| i as f64).collect(); - let parray = PrimitiveArray::from_iter(values); - let alp_encoded = alp_encode(parray.as_view(), None, &mut SESSION.create_execution_ctx())?; - - assert!( - alp_encoded.patches().is_none(), - "Expected ALP array without patches" - ); - - let array = alp_encoded.as_array(); - - let serialization = SESSION.array_serialize(array)?.unwrap(); - let children = array.children(); - let buffers = array - .buffers() - .into_iter() - .map(BufferHandle::new_host) - .collect::>(); - - let deserialized = ALPPatchedPlugin.deserialize( - ArrayDeserialization::new( - ALPPatchedPlugin.id(), - array.dtype(), - array.len(), - &serialization.metadata, - &buffers, - &children, - ), - &SESSION, - )?; - - assert_arrays_eq!(deserialized, parray, &mut SESSION.create_execution_ctx()); - - let result = deserialized - .try_downcast::() - .map_err(|a| vortex_err!("Expected deserialized ALP, got {}", a.encoding_id()))?; - - assert!(result.patches().is_none(), "Result should not have patches"); - - Ok(()) - } - - #[test] - #[should_panic(expected = "index out of bounds")] - fn primitive_array_returns_error() { - let array = PrimitiveArray::from_iter([1.0f64, 2.0, 3.0]).into_array(); - - let serialization = SESSION.array_serialize(&array).unwrap().unwrap(); - let children = array.children(); - let buffers = array - .buffers() - .into_iter() - .map(BufferHandle::new_host) - .collect::>(); - - // This panics because PrimitiveArray has no children and ALP requires encoded child. - let _result = ALPPatchedPlugin.deserialize( - ArrayDeserialization::new( - ALPPatchedPlugin.id(), - array.dtype(), - array.len(), - &serialization.metadata, - &buffers, - &children, - ), - &SESSION, - ); - } -} diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array/mod.rs similarity index 100% rename from encodings/alp/src/alp_rd/array.rs rename to encodings/alp/src/alp_rd/array/mod.rs diff --git a/encodings/bytebool/src/array.rs b/encodings/bytebool/src/array/mod.rs similarity index 100% rename from encodings/bytebool/src/array.rs rename to encodings/bytebool/src/array/mod.rs diff --git a/encodings/datetime-parts/src/array.rs b/encodings/datetime-parts/src/array/mod.rs similarity index 100% rename from encodings/datetime-parts/src/array.rs rename to encodings/datetime-parts/src/array/mod.rs diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array/mod.rs similarity index 100% rename from encodings/fsst/src/array.rs rename to encodings/fsst/src/array/mod.rs diff --git a/encodings/onpair/src/array.rs b/encodings/onpair/src/array/mod.rs similarity index 100% rename from encodings/onpair/src/array.rs rename to encodings/onpair/src/array/mod.rs diff --git a/encodings/parquet-variant/src/vtable.rs b/encodings/parquet-variant/src/vtable/mod.rs similarity index 100% rename from encodings/parquet-variant/src/vtable.rs rename to encodings/parquet-variant/src/vtable/mod.rs diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array/mod.rs similarity index 100% rename from encodings/pco/src/array.rs rename to encodings/pco/src/array/mod.rs diff --git a/encodings/sequence/src/array.rs b/encodings/sequence/src/array/mod.rs similarity index 100% rename from encodings/sequence/src/array.rs rename to encodings/sequence/src/array/mod.rs diff --git a/encodings/zigzag/src/array.rs b/encodings/zigzag/src/array/mod.rs similarity index 100% rename from encodings/zigzag/src/array.rs rename to encodings/zigzag/src/array/mod.rs diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array/mod.rs similarity index 100% rename from encodings/zstd/src/array.rs rename to encodings/zstd/src/array/mod.rs diff --git a/encodings/zstd/src/zstd_buffers.rs b/encodings/zstd/src/zstd_buffers/mod.rs similarity index 100% rename from encodings/zstd/src/zstd_buffers.rs rename to encodings/zstd/src/zstd_buffers/mod.rs diff --git a/vortex-array/src/arrays/runend/array.rs b/vortex-array/src/arrays/runend/array/mod.rs similarity index 100% rename from vortex-array/src/arrays/runend/array.rs rename to vortex-array/src/arrays/runend/array/mod.rs diff --git a/vortex-python/src/arrays/py/vtable.rs b/vortex-python/src/arrays/py/vtable/mod.rs similarity index 100% rename from vortex-python/src/arrays/py/vtable.rs rename to vortex-python/src/arrays/py/vtable/mod.rs