From 8aeddf41a2e623dce31915643593d9c0b3f744d9 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 5 Nov 2026 09:43:23 -0500 Subject: [PATCH 1/6] Read legacy vortex.decimal nodes into the canonical array or a Widen Signed-off-by: Connor Tsui --- vortex-array/src/arrays/decimal/mod.rs | 18 +- vortex-array/src/arrays/decimal/vtable/mod.rs | 252 ++----------- .../src/arrays/decimal/vtable/plugin/mod.rs | 113 ++++++ .../decimal/vtable/plugin/tests/device.rs | 235 +++++++++++++ .../arrays/decimal/vtable/plugin/tests/mod.rs | 330 ++++++++++++++++++ .../src/arrays/decimal/vtable/stored.rs | 114 ++++++ .../src/arrays/decimal/vtable/tests.rs | 131 +++++++ vortex-array/src/arrays/fixed_width/vtable.rs | 19 +- .../src/arrays/widen/encode/decimal.rs | 48 ++- .../src/arrays/widen/tests/decimal_stored.rs | 4 + vortex-array/src/session/mod.rs | 4 +- 11 files changed, 1039 insertions(+), 229 deletions(-) create mode 100644 vortex-array/src/arrays/decimal/vtable/plugin/mod.rs create mode 100644 vortex-array/src/arrays/decimal/vtable/plugin/tests/device.rs create mode 100644 vortex-array/src/arrays/decimal/vtable/plugin/tests/mod.rs create mode 100644 vortex-array/src/arrays/decimal/vtable/stored.rs create mode 100644 vortex-array/src/arrays/decimal/vtable/tests.rs diff --git a/vortex-array/src/arrays/decimal/mod.rs b/vortex-array/src/arrays/decimal/mod.rs index a806957e62d..c42dbe4ff73 100644 --- a/vortex-array/src/arrays/decimal/mod.rs +++ b/vortex-array/src/arrays/decimal/mod.rs @@ -1,23 +1,31 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Decimal arrays whose logical precision determines their canonical integer width. +//! +//! The [`Decimal`] vtable reads only nodes at the canonical width. [`DecimalPlugin`] also reads +//! legacy `vortex.decimal` nodes at any other width, as the canonical array or as a [`Widen`]. +//! +//! [`Widen`]: crate::arrays::Widen + mod array; pub use array::DecimalArrayExt; pub use array::DecimalArraySlotsExt; pub use array::DecimalData; pub use array::DecimalDataParts; pub use array::DecimalSlots; -pub use vtable::DecimalArray; pub(crate) mod compute; +pub use compute::rules::DecimalMaskedValidityRule; mod vtable; -pub use compute::rules::DecimalMaskedValidityRule; pub use vtable::Decimal; +pub use vtable::DecimalArray; +pub use vtable::DecimalPlugin; + +mod utils; +pub use utils::*; pub(crate) fn initialize(session: &vortex_session::VortexSession) { vtable::initialize(session); } - -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 64691a6cb20..e92791c8a2d 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -1,48 +1,56 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! The [`Decimal`] vtable and the two readers of the `vortex.decimal` wire form. +//! +//! [`Decimal`] reads only nodes at the canonical width of their precision, and [`DecimalPlugin`] +//! also reads legacy nodes at any other width. Both readers validate a node with [`StoredDecimal`] +//! first. + +mod kernel; +mod operations; + +mod plugin; +pub use plugin::DecimalPlugin; + +mod stored; +use stored::StoredDecimal; + +mod validity; + +use std::hash::Hash; use std::hash::Hasher; use prost::Message; -use vortex_buffer::Alignment; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use crate::ArrayParts; use crate::ArrayRef; +use crate::EqMode; use crate::ExecutionCtx; use crate::ExecutionResult; use crate::array::Array; +use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; use crate::arrays::decimal::DecimalData; +use crate::arrays::decimal::array::DecimalSlots; +use crate::arrays::decimal::compute::rules::RULES; use crate::arrays::fixed_width::vtable as fixed_width; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::DecimalBuilder; use crate::dtype::DType; use crate::dtype::DecimalType; -use crate::dtype::NativeDecimalType; -use crate::match_each_decimal_value_type; -use crate::serde::ArrayChildren; -mod kernel; -mod operations; -mod validity; - -use std::hash::Hash; - -use vortex_session::registry::CachedId; - -use crate::EqMode; -use crate::array::ArrayId; -use crate::arrays::decimal::array::DecimalSlots; -use crate::arrays::decimal::compute::rules::RULES; use crate::hash::ArrayEq; use crate::hash::ArrayHash; +use crate::serde::ArrayChildren; + /// A [`Decimal`]-encoded Vortex array. pub type DecimalArray = Array; @@ -173,33 +181,15 @@ impl VTable for Decimal { children: &dyn ArrayChildren, _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)?; - - let Some(decimal_dtype) = dtype.as_decimal_opt() else { - vortex_bail!("Expected Decimal dtype, got {:?}", dtype) - }; + let stored = StoredDecimal::deserialize(dtype, len, metadata, buffers, children)?; + vortex_ensure_eq!( + stored.values_type, + stored.decimal_dtype.canonical_values_type(), + "Decimal reads only the canonical width of {}, and DecimalPlugin reads legacy widths", + stored.decimal_dtype, + ); - let slots = DecimalData::make_slots(&validity, len); - 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, values_type, *decimal_dtype) - })?; - Ok(ArrayParts::new( - self.clone(), - dtype.clone(), - len, - data, - slots, - )) + stored.into_array_parts() } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { @@ -234,178 +224,4 @@ impl VTable for Decimal { pub struct Decimal; #[cfg(test)] -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::validity::Validity; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - - #[test] - fn test_array_serde() { - let array = DecimalArray::new( - buffer![100i128, 200i128, 300i128, 400i128, 500i128], - DecimalDType::new(10, 2), - Validity::NonNullable, - ); - let dtype = array.dtype().clone(); - - let array_ctx = ArrayContext::empty(); - let out = array - .into_array() - .serialize(&array_ctx, &SESSION, &SerializeOptions::default()) - .unwrap(); - // Concat into a single buffer - let mut concat = ByteBufferMut::empty(); - for buf in out { - concat.extend_from_slice(buf.as_ref()); - } - - let concat = concat.freeze(); - - let parts = SerializedArray::try_from(concat).unwrap(); - let decoded = parts - .decode(&dtype, 5, &ReadContext::new(array_ctx.to_ids()), &SESSION) - .unwrap(); - assert!(decoded.is::()); - assert_arrays_eq!( - decoded, - DecimalArray::from_iter([100i128, 200, 300, 400, 500], DecimalDType::new(10, 2)), - &mut SESSION.create_execution_ctx() - ); - } - - #[test] - fn test_nullable_decimal_serde_roundtrip() { - let mut ctx = SESSION.create_execution_ctx(); - let array = DecimalArray::new( - buffer![1234567i32, 0i32, -9999999i32], - DecimalDType::new(7, 3), - Validity::from_iter([true, false, true]), - ); - 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()) - .unwrap(); - let mut concat = ByteBufferMut::empty(); - for buf in out { - concat.extend_from_slice(buf.as_ref()); - } - - let parts = SerializedArray::try_from(concat.freeze()).unwrap(); - let decoded = parts - .decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &SESSION) - .unwrap(); - - 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, - ); - - let expected = - format!("DecimalMetadata values_type must be a known DecimalType, got {raw}"); - assert!(result.is_err_and(|err| err.to_string().contains(&expected))); - } - - #[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(()) - } -} +mod tests; diff --git a/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs b/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs new file mode 100644 index 00000000000..3b14bd5528c --- /dev/null +++ b/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs @@ -0,0 +1,113 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Serde for `vortex.decimal` nodes, including legacy nodes. +//! +//! [`StoredDecimal::is_legacy`] defines a legacy node. [`Decimal`] rejects legacy nodes, and +//! [`DecimalPlugin`] reads them as the canonical array or a [`Widen`]. +//! +//! 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. + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use super::Decimal; +use super::StoredDecimal; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array::Array; +use crate::array::ArrayDeserialization; +use crate::array::ArrayId; +use crate::array::ArrayPlugin; +use crate::array::ArraySerialization; +use crate::array::VTable; +use crate::arrays::Widen; +use crate::arrays::WidenArray; +use crate::arrays::decimal::DecimalDataParts; +use crate::expr::stats::Precision; + +/// Reads every `vortex.decimal` 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 copies it to the host to scan its valid values. +/// +/// The default [`ArraySession`] registers this plugin in place of [`Decimal`], and the plugin +/// writes arrays as [`Decimal`] does. +/// +/// [`ArraySession`]: crate::session::ArraySession +#[derive(Clone, Debug)] +pub struct DecimalPlugin; + +impl ArrayPlugin for DecimalPlugin { + fn id(&self) -> ArrayId { + VTable::id(&Decimal) + } + + fn read_ids(&self) -> Vec { + vec![self.id()] + } + + fn write_ids(&self) -> Vec { + vec![self.id()] + } + + fn serialize( + &self, + array: &ArrayRef, + session: &VortexSession, + ) -> VortexResult> { + ArrayPlugin::serialize(&Decimal, array, session) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!( + parts.serialized_id, + self.id(), + "DecimalPlugin reads only its own serialized ID", + ); + + let stored = StoredDecimal::deserialize( + parts.dtype, + parts.len, + parts.metadata, + parts.buffers, + parts.children, + )?; + + if !stored.is_legacy() { + return Ok(Array::::try_from_parts(stored.into_array_parts()?)?.into_array()); + } + + // This is a legacy node, whose width a `Decimal` leaf no longer allows, so the valid values + // pick the canonical array or a `Widen`. + let parts = DecimalDataParts { + decimal_dtype: stored.decimal_dtype, + values: stored.values, + values_type: stored.values_type, + validity: stored.validity, + }; + + // Stored statistics are attached only after this returns, so no range is known yet. + WidenArray::from_stored_decimal( + parts, + Precision::Absent, + &mut session.create_execution_ctx(), + ) + } + + fn is_supported_encoding(&self, id: &ArrayId) -> bool { + *id == VTable::id(&Decimal) || *id == VTable::id(&Widen) + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/arrays/decimal/vtable/plugin/tests/device.rs b/vortex-array/src/arrays/decimal/vtable/plugin/tests/device.rs new file mode 100644 index 00000000000..f26632e6352 --- /dev/null +++ b/vortex-array/src/arrays/decimal/vtable/plugin/tests/device.rs @@ -0,0 +1,235 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Host copies of canonical and legacy nodes backed by device buffers. + +use std::any::Any; +use std::ops::Range; +use std::sync::Arc; + +use futures::future::BoxFuture; +use vortex_buffer::Alignment; +use vortex_buffer::ByteBuffer; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; + +use super::SESSION; +use super::canonical_reference; +use super::decimal_leaf; +use super::deserialize; +use super::metadata; +use super::stored_values; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::arrays::BoolArray; +use crate::arrays::Decimal; +use crate::arrays::DictArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::decimal::DecimalArrayExt; +use crate::assert_arrays_eq; +use crate::buffer::BufferHandle; +use crate::buffer::DeviceBuffer; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::DecimalType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::validity::Validity; + +/// A device buffer that holds its bytes on the host. A copy to the host succeeds unless +/// `fails_host_copy` is set. +#[derive(Debug, PartialEq, Eq, Hash)] +struct TestDeviceBuffer { + bytes: ByteBuffer, + fails_host_copy: bool, +} + +impl DeviceBuffer for TestDeviceBuffer { + fn as_any(&self) -> &dyn Any { + self + } + + fn len(&self) -> usize { + self.bytes.len() + } + + fn alignment(&self) -> Alignment { + self.bytes.alignment() + } + + fn copy_to_host_sync(&self, alignment: Alignment) -> VortexResult { + if self.fails_host_copy { + vortex_bail!("no host copy"); + } + + Ok(ByteBuffer::copy_from_aligned( + self.bytes.as_slice(), + alignment, + )) + } + + fn copy_to_host( + &self, + _alignment: Alignment, + ) -> VortexResult>> { + vortex_bail!("a decimal read copies device buffers synchronously") + } + + fn slice(&self, range: Range) -> Arc { + Arc::new(Self { + bytes: self.bytes.slice(range), + fails_host_copy: self.fails_host_copy, + }) + } + + fn aligned(self: Arc, alignment: Alignment) -> VortexResult> { + assert!(self.alignment().is_aligned_to(alignment)); + Ok(self) + } +} + +fn on_device(bytes: ByteBuffer) -> BufferHandle { + BufferHandle::new_device(Arc::new(TestDeviceBuffer { + bytes, + fails_host_copy: false, + })) +} + +fn on_failing_device(bytes: ByteBuffer) -> BufferHandle { + BufferHandle::new_device(Arc::new(TestDeviceBuffer { + bytes, + fails_host_copy: true, + })) +} + +/// A canonical node keeps its device buffer, and a legacy node and its validity move to the host. +#[test] +fn device_buffers_read_on_host() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let dtype = DType::Decimal(DecimalDType::new(18, 2), Nullability::Nullable); + let validity = [true, false, true]; + let validity_bits = on_device(ByteBuffer::copy_from([0b101u8])); + let validity_child = + BoolArray::try_new_from_handle(validity_bits, 0, 3, Validity::NonNullable)?; + let children = [validity_child.into_array()]; + + let canonical_values = on_device(buffer![-99i64, 7, 99].into_byte_buffer()); + let canonical = deserialize( + &dtype, + 3, + &metadata(DecimalType::I64), + &[canonical_values], + &children, + )?; + assert!(canonical.as_::().buffer_handle().is_on_device()); + + let legacy_values = on_device(buffer![-99i32, 7, 99].into_byte_buffer()); + let legacy = deserialize( + &dtype, + 3, + &metadata(DecimalType::I32), + &[legacy_values], + &children, + )?; + assert!(legacy.is_host()); + assert_eq!( + decimal_leaf(&legacy).decimal_dtype(), + DecimalDType::new(2, 2) + ); + assert_arrays_eq!( + legacy, + canonical_reference(DecimalDType::new(18, 2), &[-99, 7, 99], Some(&validity))?, + &mut ctx + ); + + Ok(()) +} + +/// A device buffer one slot below the validity array is also read on the host. +#[test] +fn nested_device_validity_reads_on_host() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let dtype = DType::Decimal(DecimalDType::new(18, 2), Nullability::Nullable); + let codes = PrimitiveArray::from_buffer_handle( + on_device(buffer![0u8, 1, 0].into_byte_buffer()), + PType::U8, + Validity::NonNullable, + ); + let validity_child = DictArray::try_new( + codes.into_array(), + BoolArray::from_iter([true, false]).into_array(), + )? + .into_array(); + assert!(!validity_child.is_host()); + + let legacy = deserialize( + &dtype, + 3, + &metadata(DecimalType::I32), + &[stored_values(DecimalType::I32, &[-99, 7, 99])], + &[validity_child], + )?; + assert!(legacy.is_host()); + assert_eq!( + decimal_leaf(&legacy).decimal_dtype(), + DecimalDType::new(2, 2) + ); + assert_arrays_eq!( + legacy, + canonical_reference( + DecimalDType::new(18, 2), + &[-99, 7, 99], + Some(&[true, false, true]) + )?, + &mut ctx + ); + + Ok(()) +} + +/// A legacy node fails when its values or its validity cannot be copied to the host. A canonical +/// node copies nothing, so it still reads. +#[test] +fn failed_host_copy_fails() -> VortexResult<()> { + let dtype = DType::Decimal(DecimalDType::new(18, 2), Nullability::Nullable); + let device_validity = [BoolArray::try_new_from_handle( + on_failing_device(ByteBuffer::copy_from([0b101u8])), + 0, + 3, + Validity::NonNullable, + )? + .into_array()]; + let host_validity = BoolArray::from_iter([true, false, true]).into_array(); + + let error = deserialize( + &dtype, + 3, + &metadata(DecimalType::I32), + &[stored_values(DecimalType::I32, &[-99, 7, 99])], + &device_validity, + ) + .expect_err("a legacy node with device validity that cannot be copied fails"); + assert!(error.to_string().contains("no host copy")); + + let error = deserialize( + &dtype, + 3, + &metadata(DecimalType::I32), + &[on_failing_device(buffer![-99i32, 7, 99].into_byte_buffer())], + &[host_validity], + ) + .expect_err("a legacy node with device values that cannot be copied fails"); + assert!(error.to_string().contains("no host copy")); + + let canonical = deserialize( + &dtype, + 3, + &metadata(DecimalType::I64), + &[on_failing_device(buffer![-99i64, 7, 99].into_byte_buffer())], + &device_validity, + )?; + assert!(canonical.as_::().buffer_handle().is_on_device()); + + Ok(()) +} diff --git a/vortex-array/src/arrays/decimal/vtable/plugin/tests/mod.rs b/vortex-array/src/arrays/decimal/vtable/plugin/tests/mod.rs new file mode 100644 index 00000000000..4b936f39fa3 --- /dev/null +++ b/vortex-array/src/arrays/decimal/vtable/plugin/tests/mod.rs @@ -0,0 +1,330 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Reads of canonical and legacy `vortex.decimal` nodes. +//! +//! Hand-built nodes cover legacy widths without constructing invalid decimal arrays. Their raw +//! slots include nonzero null payloads, and absent validity denotes a non-nullable node. + +mod device; + +use std::sync::LazyLock; + +use prost::Message; +use rstest::rstest; +use vortex_buffer::Buffer; +use vortex_buffer::ByteBuffer; +use vortex_buffer::ByteBufferMut; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; + +use super::DecimalPlugin; +use crate::ArrayContext; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array::ArrayDeserialization; +use crate::array::ArrayPlugin; +use crate::array::VTable; +use crate::arrays::BoolArray; +use crate::arrays::Decimal; +use crate::arrays::DecimalArray; +use crate::arrays::Widen; +use crate::arrays::decimal::DecimalArrayExt; +use crate::arrays::decimal::vtable::DecimalMetadata; +use crate::arrays::widen::WidenArraySlotsExt; +use crate::assert_arrays_eq; +use crate::buffer::BufferHandle; +use crate::dtype::BigCast; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::DecimalType; +use crate::dtype::Nullability; +use crate::match_each_decimal_value_type; +use crate::serde::SerializeOptions; +use crate::serde::SerializedArray; +use crate::validity::Validity; + +static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + +/// An array read from a node, with the address of the buffer of the node to check its reuse. +struct ReadNode { + array: ArrayRef, + /// The original host allocation address, used only to check buffer identity. + stored_ptr: *const u8, +} + +/// Reads a node that stores `slots` as `values_type` integers, with the dtype `decimal_dtype`. +fn read_node( + values_type: DecimalType, + decimal_dtype: DecimalDType, + slots: &[i128], + validity: Option<&[bool]>, +) -> VortexResult { + let values = stored_values(values_type, slots); + let stored_ptr = values.as_host().as_ptr(); + let (nullability, children) = match validity { + Some(validity) => ( + Nullability::Nullable, + vec![BoolArray::from_iter(validity.iter().copied()).into_array()], + ), + None => (Nullability::NonNullable, Vec::new()), + }; + + let array = deserialize( + &DType::Decimal(decimal_dtype, nullability), + slots.len(), + &metadata(values_type), + &[values], + &children, + )?; + + Ok(ReadNode { array, stored_ptr }) +} + +/// Deserializes the parts of a node through [`DecimalPlugin`]. +fn deserialize( + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &[ArrayRef], +) -> VortexResult { + let id = DecimalPlugin.id(); + let parts = ArrayDeserialization::new(id, dtype, len, metadata, buffers, &children); + + DecimalPlugin.deserialize(parts, &SESSION) +} + +fn metadata(values_type: DecimalType) -> Vec { + DecimalMetadata { + values_type: values_type as i32, + } + .encode_to_vec() +} + +/// Returns `slots` as a host buffer of `values_type` integers. +fn stored_values(values_type: DecimalType, slots: &[i128]) -> BufferHandle { + let bytes = match_each_decimal_value_type!(values_type, |T| { + slots + .iter() + .map(|&slot| ::from(slot).vortex_expect("a slot fits its stored width")) + .collect::>() + .into_byte_buffer() + }); + + BufferHandle::new_host(bytes) +} + +/// Returns the canonical array of the valid values of a node, with zeros in its null slots. +fn canonical_reference( + decimal_dtype: DecimalDType, + slots: &[i128], + validity: Option<&[bool]>, +) -> VortexResult { + let valid_slots = slots + .iter() + .enumerate() + .map(|(row, &slot)| match validity { + Some(validity) if !validity[row] => 0, + _ => slot, + }) + .collect::>(); + let validity = match validity { + Some(validity) => Validity::from_iter(validity.iter().copied()), + None => Validity::NonNullable, + }; + + let canonical_type = decimal_dtype.canonical_values_type(); + + DecimalArray::try_new_handle( + stored_values(canonical_type, &valid_slots), + canonical_type, + decimal_dtype, + validity, + ) +} + +/// Returns the decimal leaf that holds the values of `array`: the child of a `Widen`, or `array`. +fn decimal_leaf(array: &ArrayRef) -> DecimalArray { + let leaf = match array.as_opt::() { + Some(widen) => widen.values().clone(), + None => array.clone(), + }; + + leaf.as_::().into_owned() +} + +/// Checks the encoding, the leaf dtype and width, the reuse of the stored buffer, and the values. +/// +/// The cases cover each result of a read: a canonical node keeps its buffer, and a legacy node +/// reads as a `Widen` or as the canonical array. The tests of [`WidenArray::from_stored_decimal`] +/// cover the rule that picks the child precision. +/// +/// [`WidenArray::from_stored_decimal`]: crate::arrays::WidenArray::from_stored_decimal +#[rstest] +#[case::canonical_reuse( + DecimalType::I32, DecimalDType::new(9, 2), + vec![-999_999_999, 0, 999_999_999], None, + None, true, +)] +#[case::narrow_reuse( + DecimalType::I16, DecimalDType::new(18, 2), + vec![-9_999, 0, 9_999], None, + Some(4), true, +)] +#[case::wider_to_canonical( + DecimalType::I128, DecimalDType::new(5, 2), + vec![-99_999, 0, 99_999], None, + None, false, +)] +#[case::nullable_reuse( + DecimalType::I16, DecimalDType::new(18, 2), + vec![-9_999, 32_767, 9_999], Some(vec![true, false, true]), + Some(4), true, +)] +fn reads_node_at_any_width( + #[case] values_type: DecimalType, + #[case] decimal_dtype: DecimalDType, + #[case] slots: Vec, + #[case] validity: Option>, + #[case] child_precision: Option, + #[case] reuses_buffer: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + let node = read_node(values_type, decimal_dtype, &slots, validity.as_deref())?; + + assert_eq!(node.array.is::(), child_precision.is_some()); + let leaf = decimal_leaf(&node.array); + let leaf_dtype = match child_precision { + Some(precision) => DecimalDType::new(precision, decimal_dtype.scale()), + None => decimal_dtype, + }; + assert_eq!(leaf.decimal_dtype(), leaf_dtype); + assert_eq!(leaf.values_type(), leaf_dtype.canonical_values_type()); + + // A leaf at another width than the stored width is a copy. Two empty buffers can share an + // address, so only a reuse is checked by address. + assert_eq!(leaf.values_type() == values_type, reuses_buffer); + if reuses_buffer { + assert_eq!(leaf.buffer_handle().as_host().as_ptr(), node.stored_ptr); + } + + assert_arrays_eq!( + node.array, + canonical_reference(decimal_dtype, &slots, validity.as_deref())?, + &mut ctx + ); + + Ok(()) +} + +/// Wire length and alignment must be checked before reading the stored integers. +#[rstest] +#[case::short_buffer(18, DecimalType::I32, 4, stored_values(DecimalType::I32, &[1, 2, 3]), + "byte length of 4 i32 values")] +#[case::overflowing_length(18, DecimalType::I16, usize::MAX / 2 + 1, + stored_values(DecimalType::I16, &[]), "fits in usize")] +#[case::misaligned_buffer(9, DecimalType::I32, 2, + BufferHandle::new_host(ByteBuffer::copy_from([0u8; 8])), "alignment")] +fn invalid_buffer_fails( + #[case] precision: u8, + #[case] values_type: DecimalType, + #[case] len: usize, + #[case] values: BufferHandle, + #[case] expected: &str, +) -> VortexResult<()> { + let dtype = DType::Decimal(DecimalDType::new(precision, 2), Nullability::NonNullable); + + let result = deserialize(&dtype, len, &metadata(values_type), &[values], &[]); + + assert!(result.is_err_and(|error| error.to_string().contains(expected))); + + Ok(()) +} + +#[rstest] +#[case::extra_child(Nullability::Nullable, 2, "got 2")] +#[case::non_nullable_node(Nullability::NonNullable, 1, "requires a nullable dtype")] +fn invalid_validity_children_fail( + #[case] nullability: Nullability, + #[case] child_count: usize, + #[case] expected: &str, +) -> VortexResult<()> { + let dtype = DType::Decimal(DecimalDType::new(18, 2), nullability); + let values = stored_values(DecimalType::I32, &[1, 2]); + let validity = BoolArray::from_iter([true, false]).into_array(); + + let result = deserialize( + &dtype, + 2, + &metadata(DecimalType::I32), + &[values], + &vec![validity; child_count], + ); + + assert!(result.is_err_and(|error| error.to_string().contains(expected))); + + Ok(()) +} + +#[test] +fn decimal_vtable_rejects_legacy_node() -> VortexResult<()> { + let dtype = DType::Decimal(DecimalDType::new(18, 2), Nullability::NonNullable); + let buffers = [stored_values(DecimalType::I16, &[1, 2])]; + let children: Vec = Vec::new(); + + let result = VTable::deserialize( + &Decimal, + &dtype, + 2, + &metadata(DecimalType::I16), + &buffers, + &children, + &SESSION, + ); + + assert!(result.is_err_and(|error| { + let message = error.to_string(); + message.contains("DecimalPlugin") && message.contains("left: i16") + })); + + Ok(()) +} + +/// Writes a canonical array of precision 2 and decodes it at precision 10, as a legacy node. +/// +/// This goes through the deserializer, which asserts that the plugin supports the encoding of the +/// result, here a `Widen`. +#[test] +fn decode_legacy_node_through_serialized_array() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = [-99, 0, 99]; + let written = canonical_reference(DecimalDType::new(2, 0), &values, None)?.into_array(); + let read_dtype = DecimalDType::new(10, 0); + + let array_ctx = ArrayContext::empty(); + let mut bytes = ByteBufferMut::empty(); + for buffer in written.serialize(&array_ctx, &SESSION, &SerializeOptions::default())? { + bytes.extend_from_slice(buffer.as_ref()); + } + + let decoded = SerializedArray::try_from(bytes.freeze())?.decode( + &DType::Decimal(read_dtype, Nullability::NonNullable), + written.len(), + &ReadContext::new(array_ctx.to_ids()), + &SESSION, + )?; + + assert!(decoded.is::()); + assert_arrays_eq!( + decoded, + canonical_reference(read_dtype, &values, None)?, + &mut ctx + ); + + Ok(()) +} diff --git a/vortex-array/src/arrays/decimal/vtable/stored.rs b/vortex-array/src/arrays/decimal/vtable/stored.rs new file mode 100644 index 00000000000..124bf274ea7 --- /dev/null +++ b/vortex-array/src/arrays/decimal/vtable/stored.rs @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Validation of decimal wire nodes before either reader constructs an array in memory. +//! +//! Both [`Decimal`] and [`DecimalPlugin`] validate a node here first, so a malformed row count or +//! buffer layout returns an error instead of failing an assert later. +//! +//! [`DecimalPlugin`]: super::DecimalPlugin + +use prost::Message; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; + +use super::Decimal; +use super::DecimalMetadata; +use crate::ArrayParts; +use crate::arrays::decimal::DecimalData; +use crate::arrays::fixed_width::vtable as fixed_width; +use crate::buffer::BufferHandle; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::DecimalType; +use crate::serde::ArrayChildren; +use crate::validity::Validity; + +/// A `vortex.decimal` node as it is stored, with its parts checked against its dtype and length. +/// +/// The wire metadata records the value width independently of the logical precision. +pub(super) struct StoredDecimal { + /// The logical precision and scale supplied by the containing format. + pub(super) decimal_dtype: DecimalDType, + /// The stored width of `values`. + pub(super) values_type: DecimalType, + /// One `values_type` integer per row, aligned to its width, in host or device memory. + pub(super) values: BufferHandle, + /// The validity, with one entry per row and the nullability of the dtype. + pub(super) validity: Validity, +} + +impl StoredDecimal { + /// Decodes the parts of a node with `len` rows of the decimal dtype `dtype`. + /// + /// # Errors + /// + /// Returns an error if `dtype` is not a decimal dtype, if the metadata names no known width, + /// if the node does not have exactly one buffer, if the buffer does not hold `len` values of + /// the stored width or is not aligned to it, or if the validity children do not match the + /// nullability of `dtype`. + pub(super) fn deserialize( + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + ) -> VortexResult { + let Some(&decimal_dtype) = dtype.as_decimal_opt() else { + vortex_bail!("Decimal node requires a decimal dtype, got {dtype}"); + }; + let values_type = DecimalMetadata::decode(metadata)?.get_values_type()?; + let values = fixed_width::single_buffer(buffers)?; + + // NB: every later read takes the row count from the buffer, so a buffer of another length + // must fail here. A read that returns another row count panics at the length assert of + // the deserializer. + let byte_len = len.checked_mul(values_type.byte_width()).ok_or_else(|| { + vortex_err!( + "Decimal node of {values_type} values must have a byte length that fits in usize, \ + got length {len}" + ) + })?; + vortex_ensure_eq!( + values.len(), + byte_len, + "Decimal buffer must have the byte length of {len} {values_type} values", + ); + + DecimalData::validate(&values, values_type)?; + + let validity = fixed_width::deserialize_validity(dtype.nullability(), len, children)?; + + Ok(Self { + decimal_dtype, + values_type, + values, + validity, + }) + } + + /// Returns whether this is a legacy node, whose stored width is not the canonical width of its + /// precision. + /// + /// Writers used to store a decimal at any width that held its values. A decimal array must now + /// use the canonical width of its precision, and a smaller width must be stored as a [`Widen`] + /// over a lower-precision child. + /// + /// [`Widen`]: crate::arrays::Widen + pub(super) fn is_legacy(&self) -> bool { + self.values_type != self.decimal_dtype.canonical_values_type() + } + + /// Returns the parts of a [`DecimalArray`](super::DecimalArray) at the stored width, with the + /// same buffer. + pub(super) fn into_array_parts(self) -> VortexResult> { + let dtype = DType::Decimal(self.decimal_dtype, self.validity.nullability()); + let len = self.values.len() / self.values_type.byte_width(); + let slots = DecimalData::make_slots(&self.validity, len); + let data = DecimalData::try_new_handle(self.values, self.values_type, self.decimal_dtype)?; + + Ok(ArrayParts::new(Decimal, dtype, len, data, slots)) + } +} diff --git a/vortex-array/src/arrays/decimal/vtable/tests.rs b/vortex-array/src/arrays/decimal/vtable/tests.rs new file mode 100644 index 00000000000..804ad718dc5 --- /dev/null +++ b/vortex-array/src/arrays/decimal/vtable/tests.rs @@ -0,0 +1,131 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Canonical decimal wire validation and round trips. +//! +//! These tests call [`Decimal`] directly as well as through session serialization, so +//! [`DecimalPlugin`] cannot hide a regression in the validation of the vtable. +//! +//! [`DecimalPlugin`]: crate::arrays::decimal::DecimalPlugin + +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::validity::Validity; + +static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + +/// Serializes through the session, then decodes using the original logical dtype and length. +fn round_trip(array: &ArrayRef) -> VortexResult { + let array_ctx = ArrayContext::empty(); + let mut bytes = ByteBufferMut::empty(); + + for buffer in array.serialize(&array_ctx, &SESSION, &SerializeOptions::default())? { + bytes.extend_from_slice(buffer.as_ref()); + } + + SerializedArray::try_from(bytes.freeze())?.decode( + array.dtype(), + array.len(), + &ReadContext::new(array_ctx.to_ids()), + &SESSION, + ) +} + +#[rstest] +#[case::first_unused(6)] +#[case::negative(-1)] +#[case::max(i32::MAX)] +fn test_deserialize_rejects_unknown_values_type(#[case] raw: i32) -> VortexResult<()> { + 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, + ); + + let expected = format!("DecimalMetadata values_type must be a known DecimalType, got {raw}"); + assert!(result.is_err_and(|err| err.to_string().contains(&expected))); + + Ok(()) +} + +#[rstest] +#[case::i32_nullable( + DecimalArray::new( + buffer![1_234_567i32, 0, -9_999_999], + DecimalDType::new(7, 3), + Validity::from_iter([true, false, true]), + ), + DecimalType::I32 +)] +#[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 array = array.into_array(); + let decoded = round_trip(&array)?; + + 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/arrays/fixed_width/vtable.rs b/vortex-array/src/arrays/fixed_width/vtable.rs index b0994013a66..0ec8f11293b 100644 --- a/vortex-array/src/arrays/fixed_width/vtable.rs +++ b/vortex-array/src/arrays/fixed_width/vtable.rs @@ -3,6 +3,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_panic; @@ -30,6 +31,13 @@ pub(crate) fn single_buffer(buffers: &[BufferHandle]) -> VortexResult VortexResult { match children.len() { 0 => Ok(Validity::from(nullability)), - 1 => Ok(Validity::Array(children.get(0, &Validity::DTYPE, len)?)), + 1 => { + // A non-nullable node never has a validity child. Accepting one would build an array + // whose validity does not match its dtype. + vortex_ensure!( + nullability.is_nullable(), + "A validity child requires a nullable dtype, got a non-nullable dtype" + ); + + Ok(Validity::Array(children.get(0, &Validity::DTYPE, len)?)) + } child_count => vortex_bail!("Expected 0 or 1 child, got {child_count}"), } } diff --git a/vortex-array/src/arrays/widen/encode/decimal.rs b/vortex-array/src/arrays/widen/encode/decimal.rs index dc7ed5e0db6..cf5985a21b5 100644 --- a/vortex-array/src/arrays/widen/encode/decimal.rs +++ b/vortex-array/src/arrays/widen/encode/decimal.rs @@ -30,6 +30,7 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; +use crate::array::ArraySlots; use crate::arrays::DecimalArray; use crate::arrays::WidenArray; use crate::arrays::decimal::DecimalArrayExt; @@ -99,9 +100,9 @@ impl WidenArray { /// We reuse buffers at the selected width and convert other widths without rescaling. /// A narrowing copy writes zeros in null slots. A reused buffer keeps its null slots. /// - /// Scanning and conversion use CPU kernels, so noncanonical buffers are brought to the host. - /// Their validity **must** already be executable on the host because these kernels need a mask - /// to exclude null payloads. + /// The scan and the conversion run on the CPU, and the scan reads the validity to skip null + /// slots. A CUDA read can leave the values and the validity in device memory, so we copy both + /// to the host first, and the result stays there. /// /// Readers and decoders may already know the range. For noncanonical storage, `known_range` /// lets us reuse it: @@ -151,6 +152,7 @@ impl WidenArray { } let values = values.try_into_host_sync()?; + let validity = validity_to_host(validity)?; let mask = validity.execute_mask(len, ctx)?; match_each_decimal_value_type!(values_type, |T| { let values = Buffer::::from_byte_buffer(values); @@ -160,6 +162,46 @@ impl WidenArray { } } +/// Returns `validity` with every buffer of its array copied to the host. +fn validity_to_host(validity: Validity) -> VortexResult { + match validity { + Validity::Array(array) => Ok(Validity::Array(array_to_host(array)?)), + Validity::NonNullable | Validity::AllValid | Validity::AllInvalid => Ok(validity), + } +} + +/// Returns `array` with the same encoding and every buffer of its tree on the host. +/// +/// # Errors +/// +/// Returns an error if a copy to the host fails, or if an encoding with a device buffer cannot +/// replace its buffers. +fn array_to_host(array: ArrayRef) -> VortexResult { + if array.is_host() { + return Ok(array); + } + + let slots = array + .slots() + .iter() + .map(|slot| slot.clone().map(array_to_host).transpose()) + .collect::>()?; + // SAFETY: each slot holds the same values, and only the placement of its buffers changes. + let array = unsafe { array.with_slots(slots) }?; + + let buffers = array.buffer_handles(); + if buffers.iter().all(BufferHandle::is_on_host) { + return Ok(array); + } + + let buffers = buffers + .into_iter() + .map(|buffer| Ok(BufferHandle::new_host(buffer.try_into_host_sync()?))) + .collect::>>()?; + // SAFETY: each buffer is a host copy of the same bytes, so only the placement changes. + unsafe { array.with_buffers(buffers) } +} + /// Chooses child precision from a valid-value range and the parent scale. /// /// We keep the parent scale so values do not need rescaling. The child precision must therefore diff --git a/vortex-array/src/arrays/widen/tests/decimal_stored.rs b/vortex-array/src/arrays/widen/tests/decimal_stored.rs index efc37b5b600..75a039686bf 100644 --- a/vortex-array/src/arrays/widen/tests/decimal_stored.rs +++ b/vortex-array/src/arrays/widen/tests/decimal_stored.rs @@ -141,6 +141,10 @@ fn validity_array(valid: &[bool]) -> Validity { vec![Some(-128)], DecimalType::I8, DecimalDType::new(10, 0), Some(3), false, )] +#[case::negative_scale( + vec![Some(-99), Some(99)], + DecimalType::I32, DecimalDType::new(18, -3), Some(2), false, +)] #[case::i128_narrowing_copy( vec![Some(-99), Some(99)], DecimalType::I128, DecimalDType::new(18, 2), Some(2), false, diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index 96be71b0454..62c500f459d 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -27,7 +27,6 @@ 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; @@ -44,6 +43,7 @@ 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. @@ -97,7 +97,7 @@ impl Default for ArraySession { this.register(Null); this.register(Bool); this.register(Primitive); - this.register(Decimal); + this.register(DecimalPlugin); this.register(VarBinView); this.register(ListView); this.register(Map); From 797d187c54628ed083af16f5f9be7b7fd504276e Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 5 Nov 2026 09:43:25 -0500 Subject: [PATCH 2/6] Read legacy decimal fixtures as canonical arrays or Widens Signed-off-by: Connor Tsui --- .../src/fixtures/arrays/synthetic/mod.rs | 43 +++++++++++++------ 1 file changed, 31 insertions(+), 12 deletions(-) diff --git a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs index 5396729b4cf..c9924f71c0b 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs @@ -24,8 +24,10 @@ mod tests { use vortex_array::arrays::DecimalArray; use vortex_array::arrays::Struct; use vortex_array::arrays::StructArray; + use vortex_array::arrays::Widen; use vortex_array::arrays::decimal::DecimalArrayExt; use vortex_array::arrays::struct_::StructArrayExt; + use vortex_array::arrays::widen::WidenArraySlotsExt; use vortex_array::assert_arrays_eq; use vortex_array::dtype::BigCast; use vortex_array::match_each_decimal_value_type; @@ -97,10 +99,11 @@ mod tests { } #[test] - fn legacy_decimal_fixtures_keep_stored_widths() -> VortexResult<()> { + fn legacy_decimal_fixtures_read_canonical_or_widen() -> VortexResult<()> { for fixture in legacy_decimal_fixtures() { let mut ctx = SESSION.create_execution_ctx(); let written = fixture.build_written(&mut ctx)?; + let expected = fixture.build(&mut ctx)?.execute::(&mut ctx)?; let bytes = adapter::write_file_to_bytes(written.clone())?; let read_back = adapter::read_file(bytes)?.execute::(&mut ctx)?; let written = written.as_::(); @@ -133,19 +136,35 @@ mod tests { "{context}", ); - let read_field = read_back - .unmasked_field_by_name(column.name)? - .clone() - .execute::(&mut ctx)?; + // The `WidenArray::from_stored_decimal` tests pin the child precision, so this + // checks only that every leaf is canonical. + let read_field = read_back.unmasked_field_by_name(column.name)?; + let leaf = match read_field.as_opt::() { + Some(widen) => widen.values().clone(), + None => read_field.clone(), + }; + let leaf = leaf.as_::(); + let read_context = format!("{context} after a file round trip"); assert_eq!( - read_field.values_type(), - column.stored_type, - "{context} after a file round trip", + leaf.values_type(), + leaf.decimal_dtype().canonical_values_type(), + "{read_context}", ); - assert_eq!( - raw_slots(&read_field), - column.slots, - "{context} after a file round trip", + + // A column at the canonical width reads as it was written, with its raw slots. + if column.stored_type == canonical_type { + assert!(!read_field.is::(), "{read_context}"); + assert_eq!( + raw_slots(&leaf.into_owned()), + column.slots, + "{read_context}" + ); + } + + assert_arrays_eq!( + read_field, + expected.unmasked_field_by_name(column.name)?, + &mut ctx ); } } From 5a19487056983f4ea72fc4aaf8c5dc44eeef23db Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 5 Nov 2026 10:04:42 -0500 Subject: [PATCH 3/6] Move vortex.decimal serde from the Decimal vtable into DecimalPlugin Serde is moving out of vtables, so the Decimal vtable now returns None from serialize and bails in deserialize, as FoR and BitPacked do. DecimalPlugin owns the metadata, the serializer, and StoredDecimal, and the vtable serde tests move into the plugin tests. Signed-off-by: Connor Tsui --- vortex-array/src/arrays/decimal/mod.rs | 5 +- vortex-array/src/arrays/decimal/vtable/mod.rs | 71 ++-------- .../src/arrays/decimal/vtable/plugin/mod.rs | 60 ++++++-- .../decimal/vtable/{ => plugin}/stored.rs | 12 +- .../arrays/decimal/vtable/plugin/tests/mod.rs | 114 ++++++++++----- .../src/arrays/decimal/vtable/tests.rs | 131 ------------------ 6 files changed, 159 insertions(+), 234 deletions(-) rename vortex-array/src/arrays/decimal/vtable/{ => plugin}/stored.rs (91%) delete mode 100644 vortex-array/src/arrays/decimal/vtable/tests.rs diff --git a/vortex-array/src/arrays/decimal/mod.rs b/vortex-array/src/arrays/decimal/mod.rs index c42dbe4ff73..82cd8f95993 100644 --- a/vortex-array/src/arrays/decimal/mod.rs +++ b/vortex-array/src/arrays/decimal/mod.rs @@ -3,8 +3,9 @@ //! Decimal arrays whose logical precision determines their canonical integer width. //! -//! The [`Decimal`] vtable reads only nodes at the canonical width. [`DecimalPlugin`] also reads -//! legacy `vortex.decimal` nodes at any other width, as the canonical array or as a [`Widen`]. +//! [`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`]. //! //! [`Widen`]: crate::arrays::Widen diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index e92791c8a2d..9153448d4c0 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -1,11 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! The [`Decimal`] vtable and the two readers of the `vortex.decimal` wire form. +//! The [`Decimal`] vtable and the plugin that owns its wire form. //! -//! [`Decimal`] reads only nodes at the canonical width of their precision, and [`DecimalPlugin`] -//! also reads legacy nodes at any other width. Both readers validate a node with [`StoredDecimal`] -//! first. +//! The vtable does not serialize or deserialize. [`DecimalPlugin`] reads and writes the +//! `vortex.decimal` wire form, including legacy nodes. mod kernel; mod operations; @@ -13,19 +12,14 @@ mod operations; mod plugin; pub use plugin::DecimalPlugin; -mod stored; -use stored::StoredDecimal; - mod validity; use std::hash::Hash; use std::hash::Hasher; -use prost::Message; 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 vortex_session::registry::CachedId; @@ -46,7 +40,6 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::DecimalBuilder; use crate::dtype::DType; -use crate::dtype::DecimalType; use crate::hash::ArrayEq; use crate::hash::ArrayHash; use crate::serde::ArrayChildren; @@ -58,28 +51,6 @@ pub(crate) fn initialize(session: &VortexSession) { kernel::initialize(session); } -/// Metadata for the existing `vortex.decimal` wire representation. -#[derive(prost::Message)] -pub struct DecimalMetadata { - /// Stored [`DecimalType`] discriminant, retained for wire compatibility. - /// - /// The legacy format records the actual width separately from the logical dtype. Readers - /// need this field even when the in-memory representation uses canonical decimal leaves. - #[prost(int32, tag = "1")] - pub(super) values_type: i32, -} - -impl DecimalMetadata { - fn get_values_type(&self) -> VortexResult { - DecimalType::try_from(self.values_type).map_err(|_| { - vortex_err!( - "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); @@ -133,15 +104,10 @@ impl VTable for Decimal { } fn serialize( - array: ArrayView<'_, Self>, + _array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { - Ok(Some( - DecimalMetadata { - values_type: array.values_type() as i32, - } - .encode_to_vec(), - )) + Ok(None) } fn validate( @@ -174,22 +140,14 @@ impl VTable for Decimal { fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + _dtype: &DType, + _len: usize, + _metadata: &[u8], + _buffers: &[BufferHandle], + _children: &dyn ArrayChildren, _session: &VortexSession, ) -> VortexResult> { - let stored = StoredDecimal::deserialize(dtype, len, metadata, buffers, children)?; - vortex_ensure_eq!( - stored.values_type, - stored.decimal_dtype.canonical_values_type(), - "Decimal reads only the canonical width of {}, and DecimalPlugin reads legacy widths", - stored.decimal_dtype, - ); - - stored.into_array_parts() + vortex_bail!("Decimal deserialization requires DecimalPlugin") } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { @@ -220,8 +178,9 @@ impl VTable for Decimal { } } +/// The vtable of [`DecimalArray`]. +/// +/// This vtable does not serialize or deserialize. A session **must** register [`DecimalPlugin`] +/// in its place to read or write decimal arrays. #[derive(Clone, Debug)] pub struct Decimal; - -#[cfg(test)] -mod tests; diff --git a/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs b/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs index 3b14bd5528c..ff0ca53a350 100644 --- a/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs @@ -3,19 +3,24 @@ //! Serde for `vortex.decimal` nodes, including legacy nodes. //! -//! [`StoredDecimal::is_legacy`] defines a legacy node. [`Decimal`] rejects legacy nodes, and -//! [`DecimalPlugin`] reads them as the canonical array or a [`Widen`]. +//! [`StoredDecimal::is_legacy`] defines a legacy node, which [`DecimalPlugin`] reads as the +//! canonical array or a [`Widen`]. //! //! 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. +mod stored; + +use prost::Message; 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 self::stored::StoredDecimal; use super::Decimal; -use super::StoredDecimal; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; @@ -28,17 +33,20 @@ use crate::array::VTable; use crate::arrays::Widen; use crate::arrays::WidenArray; use crate::arrays::decimal::DecimalDataParts; +use crate::dtype::DecimalType; use crate::expr::stats::Precision; -/// Reads every `vortex.decimal` node as the canonical decimal array or a decimal [`Widen`]. +/// 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 copies it to the host to scan its valid values. /// -/// The default [`ArraySession`] registers this plugin in place of [`Decimal`], and the plugin -/// writes arrays as [`Decimal`] does. +/// The [`Decimal`] vtable does not serialize or deserialize, so the default [`ArraySession`] +/// registers this plugin in its place. /// +/// [`DecimalArray`]: super::DecimalArray /// [`ArraySession`]: crate::session::ArraySession #[derive(Clone, Debug)] pub struct DecimalPlugin; @@ -59,9 +67,23 @@ impl ArrayPlugin for DecimalPlugin { fn serialize( &self, array: &ArrayRef, - session: &VortexSession, + _session: &VortexSession, ) -> VortexResult> { - ArrayPlugin::serialize(&Decimal, array, session) + let Some(decimal) = array.as_opt::() else { + vortex_bail!( + "DecimalPlugin serializes only Decimal arrays, got {}", + array.encoding_id() + ); + }; + let metadata = DecimalMetadata { + values_type: decimal.values_type() as i32, + }; + + Ok(Some(ArraySerialization::from_array( + self.id(), + array, + metadata.encode_to_vec(), + ))) } fn deserialize( @@ -109,5 +131,27 @@ impl ArrayPlugin for DecimalPlugin { } } +/// The metadata of a `vortex.decimal` node. +#[derive(prost::Message)] +struct DecimalMetadata { + /// The [`DecimalType`] discriminant of the stored width. + /// + /// The width is recorded apart from the dtype, because a legacy node can store a width other + /// than the canonical width of its precision. + #[prost(int32, tag = "1")] + values_type: i32, +} + +impl DecimalMetadata { + fn get_values_type(&self) -> VortexResult { + DecimalType::try_from(self.values_type).map_err(|_| { + vortex_err!( + "DecimalMetadata values_type must be a known DecimalType, got {}", + self.values_type + ) + }) + } +} + #[cfg(test)] mod tests; diff --git a/vortex-array/src/arrays/decimal/vtable/stored.rs b/vortex-array/src/arrays/decimal/vtable/plugin/stored.rs similarity index 91% rename from vortex-array/src/arrays/decimal/vtable/stored.rs rename to vortex-array/src/arrays/decimal/vtable/plugin/stored.rs index 124bf274ea7..666871767d3 100644 --- a/vortex-array/src/arrays/decimal/vtable/stored.rs +++ b/vortex-array/src/arrays/decimal/vtable/plugin/stored.rs @@ -1,10 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Validation of decimal wire nodes before either reader constructs an array in memory. +//! Validation of a `vortex.decimal` node before [`DecimalPlugin`] constructs an array in memory. //! -//! Both [`Decimal`] and [`DecimalPlugin`] validate a node here first, so a malformed row count or -//! buffer layout returns an error instead of failing an assert later. +//! A malformed row count or buffer layout returns an error here instead of failing an assert +//! later. //! //! [`DecimalPlugin`]: super::DecimalPlugin @@ -14,9 +14,9 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; -use super::Decimal; use super::DecimalMetadata; use crate::ArrayParts; +use crate::arrays::Decimal; use crate::arrays::decimal::DecimalData; use crate::arrays::fixed_width::vtable as fixed_width; use crate::buffer::BufferHandle; @@ -101,8 +101,8 @@ impl StoredDecimal { self.values_type != self.decimal_dtype.canonical_values_type() } - /// Returns the parts of a [`DecimalArray`](super::DecimalArray) at the stored width, with the - /// same buffer. + /// Returns the parts of a [`DecimalArray`](crate::arrays::DecimalArray) at the stored width, + /// with the same buffer. pub(super) fn into_array_parts(self) -> VortexResult> { let dtype = DType::Decimal(self.decimal_dtype, self.validity.nullability()); let len = self.values.len() / self.values_type.byte_width(); 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 4b936f39fa3..c3000dd962f 100644 --- a/vortex-array/src/arrays/decimal/vtable/plugin/tests/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/plugin/tests/mod.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Reads of canonical and legacy `vortex.decimal` nodes. +//! Serde of canonical and legacy `vortex.decimal` nodes. //! //! Hand-built nodes cover legacy widths without constructing invalid decimal arrays. Their raw //! slots include nonzero null payloads, and absent validity denotes a non-nullable node. @@ -15,11 +15,13 @@ use rstest::rstest; use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; +use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::ReadContext; +use super::DecimalMetadata; use super::DecimalPlugin; use crate::ArrayContext; use crate::ArrayRef; @@ -27,13 +29,11 @@ use crate::IntoArray; use crate::VortexSessionExecute; use crate::array::ArrayDeserialization; use crate::array::ArrayPlugin; -use crate::array::VTable; use crate::arrays::BoolArray; use crate::arrays::Decimal; use crate::arrays::DecimalArray; use crate::arrays::Widen; use crate::arrays::decimal::DecimalArrayExt; -use crate::arrays::decimal::vtable::DecimalMetadata; use crate::arrays::widen::WidenArraySlotsExt; use crate::assert_arrays_eq; use crate::buffer::BufferHandle; @@ -42,6 +42,7 @@ use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::DecimalType; use crate::dtype::Nullability; +use crate::dtype::i256; use crate::match_each_decimal_value_type; use crate::serde::SerializeOptions; use crate::serde::SerializedArray; @@ -98,6 +99,22 @@ fn deserialize( DecimalPlugin.deserialize(parts, &SESSION) } +/// Serializes `array` through the session, then decodes it with the dtype `dtype`. +fn serialize_and_decode(array: &ArrayRef, dtype: &DType) -> VortexResult { + let array_ctx = ArrayContext::empty(); + let mut bytes = ByteBufferMut::empty(); + for buffer in array.serialize(&array_ctx, &SESSION, &SerializeOptions::default())? { + bytes.extend_from_slice(buffer.as_ref()); + } + + SerializedArray::try_from(bytes.freeze())?.decode( + dtype, + array.len(), + &ReadContext::new(array_ctx.to_ids()), + &SESSION, + ) +} + fn metadata(values_type: DecimalType) -> Vec { DecimalMetadata { values_type: values_type as i32, @@ -222,6 +239,23 @@ fn reads_node_at_any_width( Ok(()) } +#[rstest] +#[case::first_unused(6)] +#[case::negative(-1)] +#[case::max(i32::MAX)] +fn unknown_values_type_fails(#[case] raw: i32) -> VortexResult<()> { + let dtype = DType::Decimal(DecimalDType::new(2, 0), Nullability::NonNullable); + let metadata = DecimalMetadata { values_type: raw }.encode_to_vec(); + let values = BufferHandle::new_host(ByteBuffer::zeroed(4)); + + let result = deserialize(&dtype, 4, &metadata, &[values], &[]); + + let expected = format!("DecimalMetadata values_type must be a known DecimalType, got {raw}"); + assert!(result.is_err_and(|error| error.to_string().contains(&expected))); + + Ok(()) +} + /// Wire length and alignment must be checked before reading the stored integers. #[rstest] #[case::short_buffer(18, DecimalType::I32, 4, stored_values(DecimalType::I32, &[1, 2, 3]), @@ -271,26 +305,52 @@ fn invalid_validity_children_fail( Ok(()) } -#[test] -fn decimal_vtable_rejects_legacy_node() -> VortexResult<()> { - let dtype = DType::Decimal(DecimalDType::new(18, 2), Nullability::NonNullable); - let buffers = [stored_values(DecimalType::I16, &[1, 2])]; - let children: Vec = Vec::new(); +#[rstest] +#[case::i32_nullable( + DecimalArray::new( + buffer![1_234_567i32, 0, -9_999_999], + DecimalDType::new(7, 3), + Validity::from_iter([true, false, true]), + ), + DecimalType::I32 +)] +#[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 round_trips_each_values_type( + #[case] array: DecimalArray, + #[case] values_type: DecimalType, +) -> VortexResult<()> { + let array = array.into_array(); - let result = VTable::deserialize( - &Decimal, - &dtype, - 2, - &metadata(DecimalType::I16), - &buffers, - &children, - &SESSION, - ); + let decoded = serialize_and_decode(&array, array.dtype())?; - assert!(result.is_err_and(|error| { - let message = error.to_string(); - message.contains("DecimalPlugin") && message.contains("left: i16") - })); + assert_eq!(decoded.as_::().values_type(), values_type); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -306,17 +366,9 @@ fn decode_legacy_node_through_serialized_array() -> VortexResult<()> { let written = canonical_reference(DecimalDType::new(2, 0), &values, None)?.into_array(); let read_dtype = DecimalDType::new(10, 0); - let array_ctx = ArrayContext::empty(); - let mut bytes = ByteBufferMut::empty(); - for buffer in written.serialize(&array_ctx, &SESSION, &SerializeOptions::default())? { - bytes.extend_from_slice(buffer.as_ref()); - } - - let decoded = SerializedArray::try_from(bytes.freeze())?.decode( + let decoded = serialize_and_decode( + &written, &DType::Decimal(read_dtype, Nullability::NonNullable), - written.len(), - &ReadContext::new(array_ctx.to_ids()), - &SESSION, )?; assert!(decoded.is::()); diff --git a/vortex-array/src/arrays/decimal/vtable/tests.rs b/vortex-array/src/arrays/decimal/vtable/tests.rs deleted file mode 100644 index 804ad718dc5..00000000000 --- a/vortex-array/src/arrays/decimal/vtable/tests.rs +++ /dev/null @@ -1,131 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Canonical decimal wire validation and round trips. -//! -//! These tests call [`Decimal`] directly as well as through session serialization, so -//! [`DecimalPlugin`] cannot hide a regression in the validation of the vtable. -//! -//! [`DecimalPlugin`]: crate::arrays::decimal::DecimalPlugin - -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::validity::Validity; - -static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - -/// Serializes through the session, then decodes using the original logical dtype and length. -fn round_trip(array: &ArrayRef) -> VortexResult { - let array_ctx = ArrayContext::empty(); - let mut bytes = ByteBufferMut::empty(); - - for buffer in array.serialize(&array_ctx, &SESSION, &SerializeOptions::default())? { - bytes.extend_from_slice(buffer.as_ref()); - } - - SerializedArray::try_from(bytes.freeze())?.decode( - array.dtype(), - array.len(), - &ReadContext::new(array_ctx.to_ids()), - &SESSION, - ) -} - -#[rstest] -#[case::first_unused(6)] -#[case::negative(-1)] -#[case::max(i32::MAX)] -fn test_deserialize_rejects_unknown_values_type(#[case] raw: i32) -> VortexResult<()> { - 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, - ); - - let expected = format!("DecimalMetadata values_type must be a known DecimalType, got {raw}"); - assert!(result.is_err_and(|err| err.to_string().contains(&expected))); - - Ok(()) -} - -#[rstest] -#[case::i32_nullable( - DecimalArray::new( - buffer![1_234_567i32, 0, -9_999_999], - DecimalDType::new(7, 3), - Validity::from_iter([true, false, true]), - ), - DecimalType::I32 -)] -#[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 array = array.into_array(); - let decoded = round_trip(&array)?; - - assert_eq!(decoded.as_::().values_type(), values_type); - assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); - - Ok(()) -} From a8658955acd20c1d16f6082aad783aae1fcd0b45 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 5 Nov 2026 10:16:22 -0500 Subject: [PATCH 4/6] Build canonical decimal leaves with the recorded row count StoredDecimal now keeps the row count from the containing format, so into_array_parts uses it directly instead of deriving it from the buffer. The byte length check stays, because a legacy read still takes its row count from the buffer. Signed-off-by: Connor Tsui --- .../src/arrays/decimal/vtable/plugin/stored.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/vortex-array/src/arrays/decimal/vtable/plugin/stored.rs b/vortex-array/src/arrays/decimal/vtable/plugin/stored.rs index 666871767d3..efebdf800b3 100644 --- a/vortex-array/src/arrays/decimal/vtable/plugin/stored.rs +++ b/vortex-array/src/arrays/decimal/vtable/plugin/stored.rs @@ -34,6 +34,8 @@ pub(super) struct StoredDecimal { pub(super) decimal_dtype: DecimalDType, /// The stored width of `values`. pub(super) values_type: DecimalType, + /// The row count supplied by the containing format, which `values` and `validity` match. + len: usize, /// One `values_type` integer per row, aligned to its width, in host or device memory. pub(super) values: BufferHandle, /// The validity, with one entry per row and the nullability of the dtype. @@ -62,7 +64,7 @@ impl StoredDecimal { let values_type = DecimalMetadata::decode(metadata)?.get_values_type()?; let values = fixed_width::single_buffer(buffers)?; - // NB: every later read takes the row count from the buffer, so a buffer of another length + // NB: a legacy read takes the row count from the buffer, so a buffer of another length // must fail here. A read that returns another row count panics at the length assert of // the deserializer. let byte_len = len.checked_mul(values_type.byte_width()).ok_or_else(|| { @@ -84,6 +86,7 @@ impl StoredDecimal { Ok(Self { decimal_dtype, values_type, + len, values, validity, }) @@ -105,10 +108,9 @@ impl StoredDecimal { /// with the same buffer. pub(super) fn into_array_parts(self) -> VortexResult> { let dtype = DType::Decimal(self.decimal_dtype, self.validity.nullability()); - let len = self.values.len() / self.values_type.byte_width(); - let slots = DecimalData::make_slots(&self.validity, len); + let slots = DecimalData::make_slots(&self.validity, self.len); let data = DecimalData::try_new_handle(self.values, self.values_type, self.decimal_dtype)?; - Ok(ArrayParts::new(Decimal, dtype, len, data, slots)) + Ok(ArrayParts::new(Decimal, dtype, self.len, data, slots)) } } From 9856e12cbce52b78a3d5ba81538a1ebe167ccc77 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 5 Nov 2026 13:31:51 -0500 Subject: [PATCH 5/6] Reject legacy decimal nodes in device memory instead of copying them Choosing the child precision of a legacy node scans its valid values on the CPU. Copying device buffers to the host for that scan mixes host and device execution, so from_stored_decimal now returns an error for noncanonical values or validity in device memory. The CUDA executor can read legacy nodes itself later. A canonical node still keeps its device buffer. Signed-off-by: Connor Tsui --- .../src/arrays/decimal/vtable/plugin/mod.rs | 3 +- .../decimal/vtable/plugin/tests/device.rs | 189 +++++------------- .../src/arrays/widen/encode/decimal.rs | 58 +----- 3 files changed, 67 insertions(+), 183 deletions(-) diff --git a/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs b/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs index ff0ca53a350..009750d469e 100644 --- a/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs @@ -41,7 +41,8 @@ use crate::expr::stats::Precision; /// /// 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 copies it to the host to scan its valid values. +/// [`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`] /// registers this plugin in its place. diff --git a/vortex-array/src/arrays/decimal/vtable/plugin/tests/device.rs b/vortex-array/src/arrays/decimal/vtable/plugin/tests/device.rs index f26632e6352..c6b1e532026 100644 --- a/vortex-array/src/arrays/decimal/vtable/plugin/tests/device.rs +++ b/vortex-array/src/arrays/decimal/vtable/plugin/tests/device.rs @@ -1,33 +1,30 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Host copies of canonical and legacy nodes backed by device buffers. +//! Reads of canonical and legacy nodes backed by device buffers. use std::any::Any; use std::ops::Range; use std::sync::Arc; use futures::future::BoxFuture; +use rstest::rstest; use vortex_buffer::Alignment; use vortex_buffer::ByteBuffer; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use super::SESSION; -use super::canonical_reference; -use super::decimal_leaf; use super::deserialize; use super::metadata; use super::stored_values; +use crate::ArrayRef; use crate::IntoArray; -use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::Decimal; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; use crate::arrays::decimal::DecimalArrayExt; -use crate::assert_arrays_eq; use crate::buffer::BufferHandle; use crate::buffer::DeviceBuffer; use crate::dtype::DType; @@ -37,12 +34,11 @@ use crate::dtype::Nullability; use crate::dtype::PType; use crate::validity::Validity; -/// A device buffer that holds its bytes on the host. A copy to the host succeeds unless -/// `fails_host_copy` is set. +/// A device buffer that holds its bytes on the host and fails every copy to the host, so a read +/// that copies it fails. #[derive(Debug, PartialEq, Eq, Hash)] struct TestDeviceBuffer { bytes: ByteBuffer, - fails_host_copy: bool, } impl DeviceBuffer for TestDeviceBuffer { @@ -58,28 +54,20 @@ impl DeviceBuffer for TestDeviceBuffer { self.bytes.alignment() } - fn copy_to_host_sync(&self, alignment: Alignment) -> VortexResult { - if self.fails_host_copy { - vortex_bail!("no host copy"); - } - - Ok(ByteBuffer::copy_from_aligned( - self.bytes.as_slice(), - alignment, - )) + fn copy_to_host_sync(&self, _alignment: Alignment) -> VortexResult { + vortex_bail!("a decimal read never copies device buffers to the host") } fn copy_to_host( &self, _alignment: Alignment, ) -> VortexResult>> { - vortex_bail!("a decimal read copies device buffers synchronously") + vortex_bail!("a decimal read never copies device buffers to the host") } fn slice(&self, range: Range) -> Arc { Arc::new(Self { bytes: self.bytes.slice(range), - fails_host_copy: self.fails_host_copy, }) } @@ -90,146 +78,79 @@ impl DeviceBuffer for TestDeviceBuffer { } fn on_device(bytes: ByteBuffer) -> BufferHandle { - BufferHandle::new_device(Arc::new(TestDeviceBuffer { - bytes, - fails_host_copy: false, - })) -} - -fn on_failing_device(bytes: ByteBuffer) -> BufferHandle { - BufferHandle::new_device(Arc::new(TestDeviceBuffer { - bytes, - fails_host_copy: true, - })) + BufferHandle::new_device(Arc::new(TestDeviceBuffer { bytes })) } -/// A canonical node keeps its device buffer, and a legacy node and its validity move to the host. -#[test] -fn device_buffers_read_on_host() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let dtype = DType::Decimal(DecimalDType::new(18, 2), Nullability::Nullable); - let validity = [true, false, true]; - let validity_bits = on_device(ByteBuffer::copy_from([0b101u8])); - let validity_child = - BoolArray::try_new_from_handle(validity_bits, 0, 3, Validity::NonNullable)?; - let children = [validity_child.into_array()]; - - let canonical_values = on_device(buffer![-99i64, 7, 99].into_byte_buffer()); - let canonical = deserialize( - &dtype, - 3, - &metadata(DecimalType::I64), - &[canonical_values], - &children, - )?; - assert!(canonical.as_::().buffer_handle().is_on_device()); - - let legacy_values = on_device(buffer![-99i32, 7, 99].into_byte_buffer()); - let legacy = deserialize( - &dtype, - 3, - &metadata(DecimalType::I32), - &[legacy_values], - &children, - )?; - assert!(legacy.is_host()); - assert_eq!( - decimal_leaf(&legacy).decimal_dtype(), - DecimalDType::new(2, 2) - ); - assert_arrays_eq!( - legacy, - canonical_reference(DecimalDType::new(18, 2), &[-99, 7, 99], Some(&validity))?, - &mut ctx - ); +/// Returns the validity of `[true, false, true]` as a bool array in device memory. +fn device_validity() -> VortexResult { + let bits = on_device(ByteBuffer::copy_from([0b101u8])); - Ok(()) + Ok(BoolArray::try_new_from_handle(bits, 0, 3, Validity::NonNullable)?.into_array()) } -/// A device buffer one slot below the validity array is also read on the host. -#[test] -fn nested_device_validity_reads_on_host() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let dtype = DType::Decimal(DecimalDType::new(18, 2), Nullability::Nullable); +/// Returns the validity of `[true, false, true]` as a dict whose codes are in device memory. +fn nested_device_validity() -> VortexResult { let codes = PrimitiveArray::from_buffer_handle( on_device(buffer![0u8, 1, 0].into_byte_buffer()), PType::U8, Validity::NonNullable, ); - let validity_child = DictArray::try_new( + let validity = DictArray::try_new( codes.into_array(), BoolArray::from_iter([true, false]).into_array(), - )? - .into_array(); - assert!(!validity_child.is_host()); - - let legacy = deserialize( - &dtype, - 3, - &metadata(DecimalType::I32), - &[stored_values(DecimalType::I32, &[-99, 7, 99])], - &[validity_child], )?; - assert!(legacy.is_host()); - assert_eq!( - decimal_leaf(&legacy).decimal_dtype(), - DecimalDType::new(2, 2) - ); - assert_arrays_eq!( - legacy, - canonical_reference( - DecimalDType::new(18, 2), - &[-99, 7, 99], - Some(&[true, false, true]) - )?, - &mut ctx - ); - Ok(()) + Ok(validity.into_array()) } -/// A legacy node fails when its values or its validity cannot be copied to the host. A canonical -/// node copies nothing, so it still reads. +fn host_validity() -> VortexResult { + Ok(BoolArray::from_iter([true, false, true]).into_array()) +} + +/// A canonical node keeps its device buffer, because its read never scans the values. #[test] -fn failed_host_copy_fails() -> VortexResult<()> { +fn canonical_node_keeps_device_buffer() -> VortexResult<()> { let dtype = DType::Decimal(DecimalDType::new(18, 2), Nullability::Nullable); - let device_validity = [BoolArray::try_new_from_handle( - on_failing_device(ByteBuffer::copy_from([0b101u8])), - 0, - 3, - Validity::NonNullable, - )? - .into_array()]; - let host_validity = BoolArray::from_iter([true, false, true]).into_array(); - let error = deserialize( + let canonical = deserialize( &dtype, 3, - &metadata(DecimalType::I32), - &[stored_values(DecimalType::I32, &[-99, 7, 99])], - &device_validity, - ) - .expect_err("a legacy node with device validity that cannot be copied fails"); - assert!(error.to_string().contains("no host copy")); + &metadata(DecimalType::I64), + &[on_device(buffer![-99i64, 7, 99].into_byte_buffer())], + &[device_validity()?], + )?; + + assert!(canonical.as_::().buffer_handle().is_on_device()); - let error = deserialize( + Ok(()) +} + +/// A legacy node fails when its values or any buffer of its validity are in device memory, because +/// choosing its width scans the values on the host. +#[rstest] +#[case::device_values(true, host_validity)] +#[case::device_validity(false, device_validity)] +#[case::nested_device_validity(false, nested_device_validity)] +fn legacy_node_in_device_memory_fails( + #[case] values_on_device: bool, + #[case] validity: fn() -> VortexResult, +) -> VortexResult<()> { + let dtype = DType::Decimal(DecimalDType::new(18, 2), Nullability::Nullable); + let values = if values_on_device { + on_device(buffer![-99i32, 7, 99].into_byte_buffer()) + } else { + stored_values(DecimalType::I32, &[-99, 7, 99]) + }; + + let result = deserialize( &dtype, 3, &metadata(DecimalType::I32), - &[on_failing_device(buffer![-99i32, 7, 99].into_byte_buffer())], - &[host_validity], - ) - .expect_err("a legacy node with device values that cannot be copied fails"); - assert!(error.to_string().contains("no host copy")); + &[values], + &[validity()?], + ); - let canonical = deserialize( - &dtype, - 3, - &metadata(DecimalType::I64), - &[on_failing_device(buffer![-99i64, 7, 99].into_byte_buffer())], - &device_validity, - )?; - assert!(canonical.as_::().buffer_handle().is_on_device()); + assert!(result.is_err_and(|error| error.to_string().contains("must be in host memory"))); Ok(()) } diff --git a/vortex-array/src/arrays/widen/encode/decimal.rs b/vortex-array/src/arrays/widen/encode/decimal.rs index cf5985a21b5..9cff25943df 100644 --- a/vortex-array/src/arrays/widen/encode/decimal.rs +++ b/vortex-array/src/arrays/widen/encode/decimal.rs @@ -30,7 +30,6 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; -use crate::array::ArraySlots; use crate::arrays::DecimalArray; use crate::arrays::WidenArray; use crate::arrays::decimal::DecimalArrayExt; @@ -100,9 +99,8 @@ impl WidenArray { /// We reuse buffers at the selected width and convert other widths without rescaling. /// A narrowing copy writes zeros in null slots. A reused buffer keeps its null slots. /// - /// The scan and the conversion run on the CPU, and the scan reads the validity to skip null - /// slots. A CUDA read can leave the values and the validity in device memory, so we copy both - /// to the host first, and the result stays there. + /// The scan and the conversion run on the CPU, so noncanonical values and their validity + /// **must** be in host memory. This function does not copy device memory to the host. /// /// Readers and decoders may already know the range. For noncanonical storage, `known_range` /// lets us reuse it: @@ -120,8 +118,8 @@ impl WidenArray { /// # Errors /// /// Rejects malformed buffer layout or validity length. For noncanonical storage, also rejects - /// failed host copies, reversed ranges, and range bounds outside the logical precision. A - /// narrowing copy rejects valid values outside the selected precision. + /// values or validity in device memory, reversed ranges, and range bounds outside the logical + /// precision. A narrowing copy rejects valid values outside the selected precision. pub fn from_stored_decimal( parts: DecimalDataParts, known_range: Precision>, @@ -151,8 +149,12 @@ impl WidenArray { return Ok(array.into_array()); } - let values = values.try_into_host_sync()?; - let validity = validity_to_host(validity)?; + // The scan below reads host memory, so device memory would panic there. + vortex_ensure!( + values.is_on_host() && validity.as_array().is_none_or(ArrayRef::is_host), + "Decimal storage at a noncanonical width must be in host memory, got device memory" + ); + let values = values.unwrap_host(); let mask = validity.execute_mask(len, ctx)?; match_each_decimal_value_type!(values_type, |T| { let values = Buffer::::from_byte_buffer(values); @@ -162,46 +164,6 @@ impl WidenArray { } } -/// Returns `validity` with every buffer of its array copied to the host. -fn validity_to_host(validity: Validity) -> VortexResult { - match validity { - Validity::Array(array) => Ok(Validity::Array(array_to_host(array)?)), - Validity::NonNullable | Validity::AllValid | Validity::AllInvalid => Ok(validity), - } -} - -/// Returns `array` with the same encoding and every buffer of its tree on the host. -/// -/// # Errors -/// -/// Returns an error if a copy to the host fails, or if an encoding with a device buffer cannot -/// replace its buffers. -fn array_to_host(array: ArrayRef) -> VortexResult { - if array.is_host() { - return Ok(array); - } - - let slots = array - .slots() - .iter() - .map(|slot| slot.clone().map(array_to_host).transpose()) - .collect::>()?; - // SAFETY: each slot holds the same values, and only the placement of its buffers changes. - let array = unsafe { array.with_slots(slots) }?; - - let buffers = array.buffer_handles(); - if buffers.iter().all(BufferHandle::is_on_host) { - return Ok(array); - } - - let buffers = buffers - .into_iter() - .map(|buffer| Ok(BufferHandle::new_host(buffer.try_into_host_sync()?))) - .collect::>>()?; - // SAFETY: each buffer is a host copy of the same bytes, so only the placement changes. - unsafe { array.with_buffers(buffers) } -} - /// Chooses child precision from a valid-value range and the parent scale. /// /// We keep the parent scale so values do not need rescaling. The child precision must therefore From 0dcf448bf6c525df07a2b83e2dc87c9399380cd8 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 5 Nov 2026 13:45:47 -0500 Subject: [PATCH 6/6] Define DecimalMetadata above DecimalPlugin Signed-off-by: Connor Tsui --- .../src/arrays/decimal/vtable/plugin/mod.rs | 44 +++++++++---------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs b/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs index 009750d469e..04925656780 100644 --- a/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs @@ -36,6 +36,28 @@ use crate::arrays::decimal::DecimalDataParts; use crate::dtype::DecimalType; use crate::expr::stats::Precision; +/// The metadata of a `vortex.decimal` node. +#[derive(prost::Message)] +struct DecimalMetadata { + /// The [`DecimalType`] discriminant of the stored width. + /// + /// The width is recorded apart from the dtype, because a legacy node can store a width other + /// than the canonical width of its precision. + #[prost(int32, tag = "1")] + values_type: i32, +} + +impl DecimalMetadata { + fn get_values_type(&self) -> VortexResult { + DecimalType::try_from(self.values_type).map_err(|_| { + vortex_err!( + "DecimalMetadata values_type must be a known DecimalType, got {}", + self.values_type + ) + }) + } +} + /// Writes a [`DecimalArray`] as a `vortex.decimal` node, and reads every such node as the /// canonical decimal array or a decimal [`Widen`]. /// @@ -132,27 +154,5 @@ impl ArrayPlugin for DecimalPlugin { } } -/// The metadata of a `vortex.decimal` node. -#[derive(prost::Message)] -struct DecimalMetadata { - /// The [`DecimalType`] discriminant of the stored width. - /// - /// The width is recorded apart from the dtype, because a legacy node can store a width other - /// than the canonical width of its precision. - #[prost(int32, tag = "1")] - values_type: i32, -} - -impl DecimalMetadata { - fn get_values_type(&self) -> VortexResult { - DecimalType::try_from(self.values_type).map_err(|_| { - vortex_err!( - "DecimalMetadata values_type must be a known DecimalType, got {}", - self.values_type - ) - }) - } -} - #[cfg(test)] mod tests;