From 864694f9014a719af1b7a14fd449499f8d4fbfef Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 23 Oct 2026 15:58:39 -0400 Subject: [PATCH 1/3] fix(array): decode decimal values_type strictly Signed-off-by: Connor Tsui --- vortex-array/src/arrays/decimal/vtable/mod.rs | 109 +++++++++++++++++- vortex-array/src/test_harness/mod.rs | 16 +++ 2 files changed, 123 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index fe4d8ddeea7..488da7a9f98 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -9,6 +9,7 @@ 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::ArrayParts; @@ -56,6 +57,17 @@ pub struct DecimalMetadata { pub(super) values_type: i32, } +impl DecimalMetadata { + fn get_values_type(&self) -> VortexResult { + DecimalType::try_from(self.values_type).map_err(|_| { + vortex_err!( + Serde: "DecimalMetadata values_type must be a known DecimalType, got {}", + self.values_type + ) + }) + } +} + impl ArrayHash for DecimalData { fn array_hash(&self, state: &mut H, accuracy: EqMode) { self.values.array_hash(state, accuracy); @@ -158,6 +170,7 @@ impl VTable for Decimal { _session: &VortexSession, ) -> VortexResult> { let metadata = DecimalMetadata::decode(metadata)?; + let values_type = metadata.get_values_type()?; let values = fixed_width::single_buffer(buffers)?; let validity = fixed_width::deserialize_validity(dtype.nullability(), len, children)?; @@ -167,14 +180,14 @@ impl VTable for Decimal { }; let slots = DecimalData::make_slots(&validity, len); - let data = match_each_decimal_value_type!(metadata.values_type(), |D| { + let data = match_each_decimal_value_type!(values_type, |D| { // Check and reinterpret-cast the buffer vortex_ensure!( values.is_aligned_to(Alignment::of::()), "DecimalArray buffer not aligned for values type {:?}", D::DECIMAL_TYPE ); - DecimalData::try_new_handle(values, metadata.values_type(), *decimal_dtype) + DecimalData::try_new_handle(values, values_type, *decimal_dtype) })?; Ok(ArrayParts::new( self.clone(), @@ -220,20 +233,33 @@ pub struct Decimal; mod tests { use std::sync::LazyLock; + use prost::Message; + use rstest::rstest; + use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; + use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::ReadContext; + use super::DecimalMetadata; use crate::ArrayContext; + use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array::VTable; use crate::arrays::Decimal; use crate::arrays::DecimalArray; use crate::assert_arrays_eq; + use crate::buffer::BufferHandle; + use crate::dtype::DType; use crate::dtype::DecimalDType; + use crate::dtype::DecimalType; + use crate::dtype::Nullability; + use crate::dtype::i256; use crate::serde::SerializeOptions; use crate::serde::SerializedArray; + use crate::test_harness::assert_serde_error; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -301,4 +327,83 @@ mod tests { assert_arrays_eq!(decoded, array, &mut ctx); } + + #[rstest] + #[case::first_unused(6)] + #[case::negative(-1)] + #[case::max(i32::MAX)] + fn test_deserialize_rejects_unknown_values_type(#[case] raw: i32) { + let metadata = DecimalMetadata { values_type: raw }.encode_to_vec(); + let buffers = [BufferHandle::new_host(ByteBuffer::zeroed(4))]; + let children: Vec = Vec::new(); + + let result = Decimal.deserialize( + &DType::Decimal(DecimalDType::new(2, 0), Nullability::NonNullable), + 4, + &metadata, + &buffers, + &children, + &SESSION, + ); + + assert_serde_error( + result, + &format!("DecimalMetadata values_type must be a known DecimalType, got {raw}"), + ); + } + + #[rstest] + #[case::i8( + DecimalArray::from_iter([1i8, -2, 3], DecimalDType::new(2, 0)), + DecimalType::I8 + )] + #[case::i16( + DecimalArray::from_iter([1i16, -2, 3], DecimalDType::new(4, 0)), + DecimalType::I16 + )] + #[case::i32( + DecimalArray::from_iter([1i32, -2, 3], DecimalDType::new(9, 2)), + DecimalType::I32 + )] + #[case::i64( + DecimalArray::from_iter([1i64, -2, 3], DecimalDType::new(18, 2)), + DecimalType::I64 + )] + #[case::i128( + DecimalArray::from_iter([1i128, -2, 3], DecimalDType::new(38, 2)), + DecimalType::I128 + )] + #[case::i256( + DecimalArray::from_iter( + [i256::from_i128(1), i256::from_i128(-2), i256::from_i128(3)], + DecimalDType::new(76, 2), + ), + DecimalType::I256 + )] + fn test_serde_each_values_type( + #[case] array: DecimalArray, + #[case] values_type: DecimalType, + ) -> VortexResult<()> { + let dtype = array.dtype().clone(); + let len = array.len(); + + let array_ctx = ArrayContext::empty(); + let out = array.clone().into_array().serialize( + &array_ctx, + &SESSION, + &SerializeOptions::default(), + )?; + let mut concat = ByteBufferMut::empty(); + for buf in out { + concat.extend_from_slice(buf.as_ref()); + } + + let parts = SerializedArray::try_from(concat.freeze())?; + let decoded = parts.decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &SESSION)?; + + assert_eq!(decoded.as_::().values_type(), values_type); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); + + Ok(()) + } } diff --git a/vortex-array/src/test_harness/mod.rs b/vortex-array/src/test_harness/mod.rs index d3a6b829e80..6f2e38e882f 100644 --- a/vortex-array/src/test_harness/mod.rs +++ b/vortex-array/src/test_harness/mod.rs @@ -6,6 +6,7 @@ use std::io::Write; use goldenfile::Mint; use goldenfile::differs::binary_diff; use itertools::Itertools; +use vortex_error::VortexError; use vortex_error::VortexResult; use crate::ExecutionCtx; @@ -27,6 +28,21 @@ pub fn check_metadata(name: &str, metadata: &[u8]) { f.write_all(metadata).unwrap(); } +/// Asserts that `result` is a [`VortexError::Serde`] error whose message is exactly `expected`. +/// +/// A read path reports malformed stored data as a `Serde` error. A substring match on the +/// message cannot show that the error has this category or that it ends with the right raw +/// value, so tests of rejected stored values use this helper instead. +#[track_caller] +#[expect(clippy::panic)] +pub fn assert_serde_error(result: VortexResult, expected: &str) { + match result { + Err(VortexError::Serde(message, _)) => assert_eq!(&*message, expected), + Err(err) => panic!("expected a Serde error, got {err}"), + Ok(_) => panic!("expected a Serde error, got Ok"), + } +} + /// Outputs the indices of the true values in a BoolArray pub fn to_int_indices(indices_bits: BoolArray, ctx: &mut ExecutionCtx) -> VortexResult> { let buffer = indices_bits.to_bit_buffer(); From 268c43fa75c8aa064f0a3c46d0417aff70cd9f6a Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 23 Oct 2026 16:00:05 -0400 Subject: [PATCH 2/3] fix(array): decode protobuf PType strictly Signed-off-by: Connor Tsui --- vortex-array/src/dtype/serde/flatbuffers.rs | 14 +++++- vortex-array/src/dtype/serde/proto.rs | 51 ++++++++++++++++++++- 2 files changed, 63 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/dtype/serde/flatbuffers.rs b/vortex-array/src/dtype/serde/flatbuffers.rs index baabcf10ce4..ffd662ecd32 100644 --- a/vortex-array/src/dtype/serde/flatbuffers.rs +++ b/vortex-array/src/dtype/serde/flatbuffers.rs @@ -556,7 +556,10 @@ impl TryFrom for PType { fb::PType::F16 => Self::F16, fb::PType::F32 => Self::F32, fb::PType::F64 => Self::F64, - _ => vortex_bail!(Serde: "Unknown PType variant"), + _ => vortex_bail!( + Serde: "flatbuffer Primitive ptype must be a known PType, got {}", + value.0 + ), }) } } @@ -580,6 +583,7 @@ mod test { use crate::flatbuffers::FlatBuffer; use crate::flatbuffers::WriteFlatBuffer; use crate::flatbuffers::WriteFlatBufferExt; + use crate::test_harness::assert_serde_error; fn roundtrip_dtype(dtype: DType) { let bytes = dtype.write_flatbuffer_bytes().unwrap(); @@ -851,4 +855,12 @@ mod test { assert!(result.is_err()); assert!(result.unwrap_err().to_string().contains("length mismatch")); } + + #[test] + fn test_unknown_ptype_from_flatbuffer() { + assert_serde_error( + PType::try_from(fb::PType(11)), + "flatbuffer Primitive ptype must be a known PType, got 11", + ); + } } diff --git a/vortex-array/src/dtype/serde/proto.rs b/vortex-array/src/dtype/serde/proto.rs index 99a3fd7801f..019fd0b408a 100644 --- a/vortex-array/src/dtype/serde/proto.rs +++ b/vortex-array/src/dtype/serde/proto.rs @@ -33,7 +33,15 @@ impl DType { { DtypeType::Null(_) => Ok(Self::Null), DtypeType::Bool(b) => Ok(Self::Bool(b.nullable.into())), - DtypeType::Primitive(p) => Ok(Self::Primitive(p.r#type().into(), p.nullable.into())), + DtypeType::Primitive(p) => { + let ptype = pb::PType::try_from(p.r#type).map_err(|_| { + vortex_err!( + Serde: "protobuf Primitive type must be a known PType, got {}", + p.r#type + ) + })?; + Ok(Self::Primitive(ptype.into(), p.nullable.into())) + } DtypeType::Decimal(d) => Ok(Self::Decimal( DecimalDType::try_new( d.precision @@ -286,6 +294,7 @@ impl TryFrom<&pb::FieldPath> for FieldPath { mod tests { use std::sync::Arc; + use rstest::rstest; use vortex_session::VortexSession; use super::*; @@ -302,6 +311,7 @@ mod tests { use crate::dtype::test::SESSION; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; + use crate::test_harness::assert_serde_error; fn round_trip_dtype(dtype: &DType) -> DType { let pb_dtype = pb::DType::try_from(dtype).unwrap(); @@ -623,6 +633,45 @@ mod tests { } } + #[rstest] + #[case::first_unused(11)] + #[case::negative(-1)] + #[case::max(i32::MAX)] + fn test_from_proto_rejects_unknown_ptype(#[case] raw: i32) { + let pb_dtype = pb::DType { + dtype_type: Some(DtypeType::Primitive(pb::Primitive { + r#type: raw, + nullable: false, + })), + }; + + assert_serde_error( + DType::from_proto(&pb_dtype, &SESSION), + &format!("protobuf Primitive type must be a known PType, got {raw}"), + ); + } + + #[test] + fn test_from_proto_rejects_unknown_ptype_in_struct_field() { + let pb_dtype = pb::DType { + dtype_type: Some(DtypeType::Struct(pb::Struct { + names: vec!["typed_value".to_string()], + dtypes: vec![pb::DType { + dtype_type: Some(DtypeType::Primitive(pb::Primitive { + r#type: 11, + nullable: true, + })), + }], + nullable: false, + })), + }; + + assert_serde_error( + DType::from_proto(&pb_dtype, &SESSION), + "protobuf Primitive type must be a known PType, got 11", + ); + } + #[test] fn test_invalid_decimal_from_proto() { // Test precision that doesn't fit in u8 From 12ad928f2562ef0537729847bcd9858f254237bf Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 23 Oct 2026 16:08:47 -0400 Subject: [PATCH 3/3] fix: decode PType metadata fields strictly Signed-off-by: Connor Tsui --- encodings/alp/src/alp_rd/array.rs | 50 +++++++++++- encodings/fastlanes/src/rle/vtable/mod.rs | 77 ++++++++++++++++-- encodings/runend/src/array.rs | 49 ++++++++++- vortex-array/src/arrays/dict/array.rs | 46 +++++++++++ vortex-array/src/arrays/dict/vtable/mod.rs | 3 +- vortex-array/src/arrays/list/vtable/mod.rs | 62 +++++++++++++- .../src/arrays/listview/vtable/mod.rs | 81 ++++++++++++++++++- vortex-array/src/arrays/varbin/vtable/mod.rs | 59 +++++++++++++- vortex-layout/src/layouts/dict/mod.rs | 65 ++++++++++++++- vortex-layout/src/layouts/list/mod.rs | 70 +++++++++++++++- 10 files changed, 544 insertions(+), 18 deletions(-) diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array.rs index 71a450585a6..2d433fd2957 100644 --- a/encodings/alp/src/alp_rd/array.rs +++ b/encodings/alp/src/alp_rd/array.rs @@ -72,6 +72,17 @@ pub struct ALPRDMetadata { patches: Option, } +impl ALPRDMetadata { + fn get_left_parts_ptype(&self) -> VortexResult { + PType::try_from(self.left_parts_ptype).map_err(|_| { + vortex_err!( + Serde: "ALPRDMetadata left_parts_ptype must be a known PType, got {}", + self.left_parts_ptype + ) + }) + } +} + impl ArrayHash for ALPRDData { fn array_hash(&self, state: &mut H, accuracy: EqMode) { self.left_parts_dictionary.array_hash(state, accuracy); @@ -172,6 +183,7 @@ impl VTable for ALPRD { 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 {}", @@ -179,7 +191,7 @@ impl VTable for ALPRD { ); } - let left_parts_dtype = DType::Primitive(metadata.left_parts_ptype(), dtype.nullability()); + 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] @@ -565,15 +577,22 @@ mod test { use prost::Message; use rstest::rstest; + use vortex_array::ArrayRef; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; + use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::patches::PatchesMetadata; + use vortex_array::test_harness::assert_serde_error; use vortex_array::test_harness::check_metadata; + use vortex_array::vtable::VTable; + use vortex_error::VortexResult; use vortex_session::VortexSession; + use super::ALPRD; use super::ALPRDMetadata; use crate::ALPRDFloat; use crate::RDEncoderExt; @@ -633,4 +652,33 @@ mod test { .encode_to_vec(), ); } + + #[test] + fn test_deserialize_rejects_unknown_left_parts_ptype() -> VortexResult<()> { + let metadata = ALPRDMetadata { + left_parts_ptype: 11, + ..Default::default() + } + .encode_to_vec(); + let children: Vec = Vec::new(); + assert_serde_error( + ALPRD.deserialize( + &DType::Primitive(PType::F64, Nullability::NonNullable), + 0, + &metadata, + &[], + &children, + &SESSION, + ), + "ALPRDMetadata left_parts_ptype must be a known PType, got 11", + ); + + let known = ALPRDMetadata { + left_parts_ptype: PType::U32 as i32, + ..Default::default() + }; + assert_eq!(known.get_left_parts_ptype()?, PType::U32); + + Ok(()) + } } diff --git a/encodings/fastlanes/src/rle/vtable/mod.rs b/encodings/fastlanes/src/rle/vtable/mod.rs index 9e4e11a8a12..9f0163261ab 100644 --- a/encodings/fastlanes/src/rle/vtable/mod.rs +++ b/encodings/fastlanes/src/rle/vtable/mod.rs @@ -26,6 +26,7 @@ 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; @@ -61,6 +62,26 @@ pub struct RLEMetadata { pub offset: u64, } +impl RLEMetadata { + fn get_indices_ptype(&self) -> VortexResult { + PType::try_from(self.indices_ptype).map_err(|_| { + vortex_err!( + Serde: "RLEMetadata indices_ptype must be a known PType, got {}", + self.indices_ptype + ) + }) + } + + fn get_values_idx_offsets_ptype(&self) -> VortexResult { + PType::try_from(self.values_idx_offsets_ptype).map_err(|_| { + vortex_err!( + Serde: "RLEMetadata values_idx_offsets_ptype must be a known PType, got {}", + self.values_idx_offsets_ptype + ) + }) + } +} + impl ArrayHash for RLEData { fn array_hash(&self, state: &mut H, _accuracy: EqMode) { self.offset.hash(state); @@ -161,12 +182,14 @@ impl VTable for RLE { 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 metadata = RLEMetadata::decode(metadata)?; let values = children.get( 0, &DType::Primitive(dtype.as_ptype(), Nullability::NonNullable), @@ -175,16 +198,13 @@ impl VTable for RLE { let indices = children.get( 1, - &DType::Primitive(metadata.indices_ptype(), dtype.nullability()), + &DType::Primitive(indices_ptype, dtype.nullability()), usize::try_from(metadata.indices_len)?, )?; let values_idx_offsets = children.get( 2, - &DType::Primitive( - metadata.values_idx_offsets_ptype(), - Nullability::NonNullable, - ), + &DType::Primitive(values_idx_offsets_ptype, Nullability::NonNullable), usize::try_from(metadata.values_idx_offsets_len)?, )?; @@ -328,9 +348,19 @@ fn validate_parts( #[cfg(test)] mod tests { use prost::Message; + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::test_harness::assert_serde_error; use vortex_array::test_harness::check_metadata; + use vortex_array::vtable::VTable; + use vortex_error::VortexResult; + use super::RLE; use super::RLEMetadata; + use crate::test::SESSION; #[cfg_attr(miri, ignore)] #[test] @@ -348,4 +378,39 @@ mod tests { .encode_to_vec(), ); } + + #[rstest] + #[case::indices_ptype( + |raw| RLEMetadata { indices_ptype: raw, ..Default::default() }, + RLEMetadata::get_indices_ptype, + "RLEMetadata indices_ptype must be a known PType, got 11", + )] + #[case::values_idx_offsets_ptype( + |raw| RLEMetadata { values_idx_offsets_ptype: raw, ..Default::default() }, + RLEMetadata::get_values_idx_offsets_ptype, + "RLEMetadata values_idx_offsets_ptype must be a known PType, got 11", + )] + fn test_deserialize_rejects_unknown_ptype( + #[case] with_ptype: fn(i32) -> RLEMetadata, + #[case] get_ptype: fn(&RLEMetadata) -> VortexResult, + #[case] expected: &str, + ) -> VortexResult<()> { + let metadata = with_ptype(11).encode_to_vec(); + let children: Vec = Vec::new(); + assert_serde_error( + RLE.deserialize( + &DType::Primitive(PType::U32, Nullability::NonNullable), + 0, + &metadata, + &[], + &children, + &SESSION, + ), + expected, + ); + + assert_eq!(get_ptype(&with_ptype(PType::U32 as i32))?, PType::U32); + + Ok(()) + } } diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index 787a5a02ecd..68b6f901edd 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -43,6 +43,7 @@ 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; @@ -69,6 +70,17 @@ pub struct RunEndMetadata { pub offset: u64, } +impl RunEndMetadata { + fn get_ends_ptype(&self) -> VortexResult { + PType::try_from(self.ends_ptype).map_err(|_| { + vortex_err!( + Serde: "RunEndMetadata ends_ptype must be a known PType, got {}", + self.ends_ptype + ) + }) + } +} + impl ArrayHash for RunEndData { fn array_hash(&self, state: &mut H, _accuracy: EqMode) { self.offset.hash(state); @@ -152,7 +164,8 @@ impl VTable for RunEnd { session: &VortexSession, ) -> VortexResult> { let metadata = RunEndMetadata::decode(metadata)?; - let ends_dtype = DType::Primitive(metadata.ends_ptype(), Nullability::NonNullable); + 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)?; @@ -575,6 +588,8 @@ mod tests { use std::sync::Arc; use std::sync::LazyLock; + use prost::Message; + use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::DecimalArray; @@ -590,11 +605,14 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::i256; + use vortex_array::test_harness::assert_serde_error; use vortex_array::validity::Validity; + use vortex_array::vtable::VTable; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; + use super::RunEndMetadata; use crate::RunEnd; static SESSION: LazyLock = LazyLock::new(|| { @@ -934,4 +952,33 @@ mod tests { assert_arrays_eq!(arr, expected, &mut ctx); Ok(()) } + + #[test] + fn test_deserialize_rejects_unknown_ends_ptype() -> VortexResult<()> { + let metadata = RunEndMetadata { + ends_ptype: 11, + ..Default::default() + } + .encode_to_vec(); + let children: Vec = Vec::new(); + assert_serde_error( + RunEnd.deserialize( + &DType::Primitive(PType::I32, Nullability::NonNullable), + 0, + &metadata, + &[], + &children, + &SESSION, + ), + "RunEndMetadata ends_ptype must be a known PType, got 11", + ); + + let known = RunEndMetadata { + ends_ptype: PType::U32 as i32, + ..Default::default() + }; + assert_eq!(known.get_ends_ptype()?, PType::U32); + + Ok(()) + } } diff --git a/vortex-array/src/arrays/dict/array.rs b/vortex-array/src/arrays/dict/array.rs index 79f5ac8be33..fa7832674e3 100644 --- a/vortex-array/src/arrays/dict/array.rs +++ b/vortex-array/src/arrays/dict/array.rs @@ -11,6 +11,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_mask::AllOr; use crate::ArrayRef; @@ -42,6 +43,17 @@ pub struct DictMetadata { pub(super) all_values_referenced: Option, } +impl DictMetadata { + pub(super) fn get_codes_ptype(&self) -> VortexResult { + PType::try_from(self.codes_ptype).map_err(|_| { + vortex_err!( + Serde: "DictMetadata codes_ptype must be a known PType, got {}", + self.codes_ptype + ) + }) + } +} + #[array_slots(Dict)] pub struct DictSlots { /// The codes array mapping each element to a dictionary entry. @@ -358,6 +370,7 @@ impl Array { mod test { use std::sync::LazyLock; + use prost::Message; use rand::RngExt; use rand::SeedableRng; use rand::distr::Distribution; @@ -371,10 +384,13 @@ mod test { use vortex_mask::AllOr; use vortex_session::VortexSession; + use super::DictMetadata; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array::VTable; use crate::arrays::ChunkedArray; + use crate::arrays::Dict; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinViewArray; @@ -387,6 +403,7 @@ mod test { use crate::dtype::Nullability::NonNullable; use crate::dtype::PType; use crate::dtype::UnsignedPType; + use crate::test_harness::assert_serde_error; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -582,6 +599,35 @@ mod test { ); } + #[test] + fn test_deserialize_rejects_unknown_codes_ptype() -> VortexResult<()> { + let metadata = DictMetadata { + codes_ptype: 11, + ..Default::default() + } + .encode_to_vec(); + let children: Vec = Vec::new(); + assert_serde_error( + Dict.deserialize( + &DType::Primitive(PType::I32, NonNullable), + 0, + &metadata, + &[], + &children, + &SESSION, + ), + "DictMetadata codes_ptype must be a known PType, got 11", + ); + + let known = DictMetadata { + codes_ptype: PType::U32 as i32, + ..Default::default() + }; + assert_eq!(known.get_codes_ptype()?, PType::U32); + + Ok(()) + } + /// Covers the early exit taken once every value is referenced, including when that happens /// after the first chunk, alongside cases that never saturate. #[rstest::rstest] diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index a2ba1487a59..e320c741c4a 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/vortex-array/src/arrays/dict/vtable/mod.rs @@ -165,6 +165,7 @@ impl VTable for Dict { _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 {}", @@ -177,7 +178,7 @@ impl VTable for Dict { // If no `is_nullable_codes` metadata use the nullability of the values // (and whole array) as before. .unwrap_or_else(|| dtype.nullability()); - let codes_dtype = DType::Primitive(metadata.codes_ptype(), codes_nullable); + let codes_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); diff --git a/vortex-array/src/arrays/list/vtable/mod.rs b/vortex-array/src/arrays/list/vtable/mod.rs index d04dd4bca07..0f7c73be6d3 100644 --- a/vortex-array/src/arrays/list/vtable/mod.rs +++ b/vortex-array/src/arrays/list/vtable/mod.rs @@ -9,6 +9,7 @@ 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; @@ -53,6 +54,17 @@ pub struct ListMetadata { offset_ptype: i32, } +impl ListMetadata { + fn get_offset_ptype(&self) -> VortexResult { + PType::try_from(self.offset_ptype).map_err(|_| { + vortex_err!( + Serde: "ListMetadata offset_ptype must be a known PType, got {}", + self.offset_ptype + ) + }) + } +} + impl ArrayHash for ListData { fn array_hash(&self, _state: &mut H, _accuracy: EqMode) {} } @@ -155,6 +167,7 @@ impl VTable for List { 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 { @@ -175,7 +188,7 @@ impl VTable for List { let offsets = children.get( 1, - &DType::Primitive(metadata.offset_ptype(), Nullability::NonNullable), + &DType::Primitive(offset_ptype, Nullability::NonNullable), len + 1, )?; @@ -225,3 +238,50 @@ impl VTable for List { #[derive(Clone, Debug)] pub struct List; + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::LazyLock; + + use prost::Message; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use super::List; + use super::ListMetadata; + use crate::ArrayRef; + use crate::array::VTable; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::test_harness::assert_serde_error; + + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + + #[test] + fn test_deserialize_rejects_unknown_offset_ptype() -> VortexResult<()> { + let dtype = DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + Nullability::NonNullable, + ); + let metadata = ListMetadata { + offset_ptype: 11, + ..Default::default() + } + .encode_to_vec(); + let children: Vec = Vec::new(); + assert_serde_error( + List.deserialize(&dtype, 0, &metadata, &[], &children, &SESSION), + "ListMetadata offset_ptype must be a known PType, got 11", + ); + + let known = ListMetadata { + offset_ptype: PType::U32 as i32, + ..Default::default() + }; + assert_eq!(known.get_offset_ptype()?, PType::U32); + + Ok(()) + } +} diff --git a/vortex-array/src/arrays/listview/vtable/mod.rs b/vortex-array/src/arrays/listview/vtable/mod.rs index 9337309c796..3f95fae0809 100644 --- a/vortex-array/src/arrays/listview/vtable/mod.rs +++ b/vortex-array/src/arrays/listview/vtable/mod.rs @@ -11,6 +11,7 @@ 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; @@ -63,6 +64,26 @@ pub struct ListViewMetadata { size_ptype: i32, } +impl ListViewMetadata { + fn get_offset_ptype(&self) -> VortexResult { + PType::try_from(self.offset_ptype).map_err(|_| { + vortex_err!( + Serde: "ListViewMetadata offset_ptype must be a known PType, got {}", + self.offset_ptype + ) + }) + } + + fn get_size_ptype(&self) -> VortexResult { + PType::try_from(self.size_ptype).map_err(|_| { + vortex_err!( + Serde: "ListViewMetadata size_ptype must be a known PType, got {}", + self.size_ptype + ) + }) + } +} + impl ArrayHash for ListViewData { fn array_hash(&self, state: &mut H, _accuracy: EqMode) { self.is_zero_copy_to_list().hash(state); @@ -164,6 +185,8 @@ impl VTable for ListView { 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" @@ -195,14 +218,14 @@ impl VTable for ListView { // Get offsets with proper type from metadata. let offsets = children.get( 1, - &DType::Primitive(metadata.offset_ptype(), Nullability::NonNullable), + &DType::Primitive(offset_ptype, Nullability::NonNullable), len, )?; // Get sizes with proper type from metadata. let sizes = children.get( 2, - &DType::Primitive(metadata.size_ptype(), Nullability::NonNullable), + &DType::Primitive(size_ptype, Nullability::NonNullable), len, )?; @@ -257,3 +280,57 @@ impl VTable for ListView { PARENT_RULES.evaluate(array, parent, child_idx) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::LazyLock; + + use prost::Message; + use rstest::rstest; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use super::ListView; + use super::ListViewMetadata; + use crate::ArrayRef; + use crate::array::VTable; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::test_harness::assert_serde_error; + + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + + #[rstest] + #[case::offset_ptype( + |raw| ListViewMetadata { offset_ptype: raw, ..Default::default() }, + ListViewMetadata::get_offset_ptype, + "ListViewMetadata offset_ptype must be a known PType, got 11", + )] + #[case::size_ptype( + |raw| ListViewMetadata { size_ptype: raw, ..Default::default() }, + ListViewMetadata::get_size_ptype, + "ListViewMetadata size_ptype must be a known PType, got 11", + )] + fn test_deserialize_rejects_unknown_ptype( + #[case] with_ptype: fn(i32) -> ListViewMetadata, + #[case] get_ptype: fn(&ListViewMetadata) -> VortexResult, + #[case] expected: &str, + ) -> VortexResult<()> { + let dtype = DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + Nullability::NonNullable, + ); + let metadata = with_ptype(11).encode_to_vec(); + let children: Vec = Vec::new(); + assert_serde_error( + ListView.deserialize(&dtype, 0, &metadata, &[], &children, &SESSION), + expected, + ); + + assert_eq!(get_ptype(&with_ptype(PType::U32 as i32))?, PType::U32); + + Ok(()) + } +} diff --git a/vortex-array/src/arrays/varbin/vtable/mod.rs b/vortex-array/src/arrays/varbin/vtable/mod.rs index 74d8812565f..783724309ad 100644 --- a/vortex-array/src/arrays/varbin/vtable/mod.rs +++ b/vortex-array/src/arrays/varbin/vtable/mod.rs @@ -9,6 +9,7 @@ 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::registry::CachedId; @@ -63,6 +64,17 @@ pub struct VarBinMetadata { pub(crate) offsets_ptype: i32, } +impl VarBinMetadata { + fn get_offsets_ptype(&self) -> VortexResult { + PType::try_from(self.offsets_ptype).map_err(|_| { + vortex_err!( + Serde: "VarBinMetadata offsets_ptype must be a known PType, got {}", + self.offsets_ptype + ) + }) + } +} + impl ArrayHash for VarBinData { fn array_hash(&self, state: &mut H, accuracy: EqMode) { self.bytes().array_hash(state, accuracy); @@ -166,6 +178,7 @@ impl VTable for VarBin { 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 { @@ -177,7 +190,7 @@ impl VTable for VarBin { let offsets = children.get( 0, - &DType::Primitive(metadata.offsets_ptype(), Nullability::NonNullable), + &DType::Primitive(offsets_ptype, Nullability::NonNullable), len + 1, )?; @@ -271,3 +284,47 @@ fn append_to_varbinview( #[derive(Clone, Debug)] pub struct VarBin; + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use prost::Message; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use super::VarBin; + use super::VarBinMetadata; + use crate::ArrayRef; + use crate::array::VTable; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::test_harness::assert_serde_error; + + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + + #[test] + fn test_deserialize_rejects_unknown_offsets_ptype() -> VortexResult<()> { + let metadata = VarBinMetadata { offsets_ptype: 11 }.encode_to_vec(); + let children: Vec = Vec::new(); + assert_serde_error( + VarBin.deserialize( + &DType::Utf8(Nullability::NonNullable), + 0, + &metadata, + &[], + &children, + &SESSION, + ), + "VarBinMetadata offsets_ptype must be a known PType, got 11", + ); + + let known = VarBinMetadata { + offsets_ptype: PType::U32 as i32, + }; + assert_eq!(known.get_offsets_ptype()?, PType::U32); + + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/dict/mod.rs b/vortex-layout/src/layouts/dict/mod.rs index eee88f916d3..c53fa5fc2d6 100644 --- a/vortex-layout/src/layouts/dict/mod.rs +++ b/vortex-layout/src/layouts/dict/mod.rs @@ -15,6 +15,7 @@ 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; @@ -71,6 +72,7 @@ impl VTable for Dict { args: &LayoutDeserializeArgs<'_>, metadata: &DictLayoutMetadata, ) -> VortexResult { + let codes_ptype = metadata.get_codes_ptype()?; vortex_ensure_eq!( args.children.nchildren(), 2, @@ -80,7 +82,7 @@ impl VTable for Dict { .is_nullable_codes .map(Nullability::from) .unwrap_or_else(|| args.dtype.nullability()); - let codes_dtype = DType::Primitive(metadata.codes_ptype(), codes_nullable); + let codes_dtype = DType::Primitive(codes_ptype, codes_nullable); args.children.child(0, args.dtype)?; let codes = args.children.child(1, &codes_dtype)?; vortex_ensure_eq!( @@ -176,4 +178,65 @@ impl DictLayoutMetadata { metadata.set_codes_ptype(codes_ptype); metadata } + + fn get_codes_ptype(&self) -> VortexResult { + PType::try_from(self.codes_ptype).map_err(|_| { + vortex_err!( + Serde: "DictLayoutMetadata codes_ptype must be a known PType, got {}", + self.codes_ptype + ) + }) + } +} + +#[cfg(test)] +mod tests { + use prost::Message; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::test_harness::assert_serde_error; + use vortex_error::VortexResult; + use vortex_session::registry::ReadContext; + + use super::Dict; + use super::DictLayoutMetadata; + use crate::LayoutBuildContext; + use crate::LayoutVTablePlugin; + use crate::layout_children; + use crate::test::SESSION; + + #[test] + fn test_build_rejects_unknown_codes_ptype() -> VortexResult<()> { + let metadata = DictLayoutMetadata { + codes_ptype: 11, + ..Default::default() + } + .encode_to_vec(); + let children = layout_children(vec![]); + let read_ctx = ReadContext::new([]); + let build_ctx = LayoutBuildContext { + session: &SESSION, + array_read_ctx: &read_ctx, + }; + assert_serde_error( + LayoutVTablePlugin::build( + &Dict, + &DType::Primitive(PType::I32, Nullability::NonNullable), + 0, + &metadata, + vec![], + children.as_ref(), + &build_ctx, + ), + "DictLayoutMetadata codes_ptype must be a known PType, got 11", + ); + + assert_eq!( + DictLayoutMetadata::new(PType::U32).get_codes_ptype()?, + PType::U32 + ); + + Ok(()) + } } diff --git a/vortex-layout/src/layouts/list/mod.rs b/vortex-layout/src/layouts/list/mod.rs index f5c8c618b15..7988b45d051 100644 --- a/vortex-layout/src/layouts/list/mod.rs +++ b/vortex-layout/src/layouts/list/mod.rs @@ -78,13 +78,14 @@ impl VTable for List { args: &LayoutDeserializeArgs<'_>, metadata: &ListLayoutMetadata, ) -> VortexResult { + let offsets_ptype = metadata.get_offsets_ptype()?; ListLayout::validate_children(args.dtype, args.children.nchildren())?; let elements_dtype = args .dtype .as_list_element_opt() .ok_or_else(|| vortex_err!("ListLayout requires a List dtype, got {}", args.dtype))?; args.children.child(ELEMENTS_CHILD_INDEX, elements_dtype)?; - let offsets_dtype = DType::Primitive(metadata.offsets_ptype(), Nullability::NonNullable); + let offsets_dtype = DType::Primitive(offsets_ptype, Nullability::NonNullable); let offsets = args.children.child(OFFSETS_CHILD_INDEX, &offsets_dtype)?; vortex_ensure_eq!( offsets.row_count().saturating_sub(1), @@ -101,9 +102,7 @@ impl VTable for List { "List validity row count does not match parent" ); } - Ok(ListData { - offsets_ptype: metadata.offsets_ptype(), - }) + Ok(ListData { offsets_ptype }) } fn nslots(_layout: &Layout) -> usize { @@ -237,4 +236,67 @@ impl ListLayoutMetadata { metadata.set_offsets_ptype(offsets_ptype); metadata } + + fn get_offsets_ptype(&self) -> VortexResult { + PType::try_from(self.offsets_ptype).map_err(|_| { + vortex_err!( + Serde: "ListLayoutMetadata offsets_ptype must be a known PType, got {}", + self.offsets_ptype + ) + }) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use prost::Message; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::test_harness::assert_serde_error; + use vortex_error::VortexResult; + use vortex_session::registry::ReadContext; + + use super::List; + use super::ListLayoutMetadata; + use crate::LayoutBuildContext; + use crate::LayoutVTablePlugin; + use crate::layout_children; + use crate::test::SESSION; + + #[test] + fn test_build_rejects_unknown_offsets_ptype() -> VortexResult<()> { + let dtype = DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + Nullability::NonNullable, + ); + let metadata = ListLayoutMetadata { offsets_ptype: 11 }.encode_to_vec(); + let children = layout_children(vec![]); + let read_ctx = ReadContext::new([]); + let build_ctx = LayoutBuildContext { + session: &SESSION, + array_read_ctx: &read_ctx, + }; + assert_serde_error( + LayoutVTablePlugin::build( + &List, + &dtype, + 0, + &metadata, + vec![], + children.as_ref(), + &build_ctx, + ), + "ListLayoutMetadata offsets_ptype must be a known PType, got 11", + ); + + assert_eq!( + ListLayoutMetadata::new(PType::U32).get_offsets_ptype()?, + PType::U32 + ); + + Ok(()) + } }