diff --git a/vortex-array/src/arrays/decimal/mod.rs b/vortex-array/src/arrays/decimal/mod.rs index a806957e62d..82cd8f95993 100644 --- a/vortex-array/src/arrays/decimal/mod.rs +++ b/vortex-array/src/arrays/decimal/mod.rs @@ -1,23 +1,32 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Decimal arrays whose logical precision determines their canonical integer width. +//! +//! [`DecimalPlugin`] reads and writes the `vortex.decimal` wire form in place of the [`Decimal`] +//! vtable. It reads a legacy node, stored at another width, as the canonical array or as a +//! [`Widen`]. +//! +//! [`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..9153448d4c0 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -1,48 +1,49 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! The [`Decimal`] vtable and the plugin that owns its wire form. +//! +//! The vtable does not serialize or deserialize. [`DecimalPlugin`] reads and writes the +//! `vortex.decimal` wire form, including legacy nodes. + +mod kernel; +mod operations; + +mod plugin; +pub use plugin::DecimalPlugin; + +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; @@ -50,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); @@ -125,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( @@ -166,40 +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 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 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, - )) + vortex_bail!("Decimal deserialization requires DecimalPlugin") } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { @@ -230,182 +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 { - 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(()) - } -} 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..04925656780 --- /dev/null +++ b/vortex-array/src/arrays/decimal/vtable/plugin/mod.rs @@ -0,0 +1,158 @@ +// 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, 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 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::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`]. +/// +/// A canonical node keeps its buffer without a copy or a scan, so the buffer can stay on a device, +/// and its valid values are trusted to fit the precision. A legacy node goes through +/// [`WidenArray::from_stored_decimal`], which scans its valid values on the host, so a legacy node +/// in device memory fails to read. +/// +/// The [`Decimal`] vtable does not 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; + +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> { + 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( + &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/stored.rs b/vortex-array/src/arrays/decimal/vtable/plugin/stored.rs new file mode 100644 index 00000000000..efebdf800b3 --- /dev/null +++ b/vortex-array/src/arrays/decimal/vtable/plugin/stored.rs @@ -0,0 +1,116 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Validation of a `vortex.decimal` node before [`DecimalPlugin`] constructs an array in memory. +//! +//! A malformed row count or buffer layout returns an error here 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::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; +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, + /// 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. + 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: 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(|| { + 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, + len, + 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`](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 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, self.len, data, slots)) + } +} 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..c6b1e532026 --- /dev/null +++ b/vortex-array/src/arrays/decimal/vtable/plugin/tests/device.rs @@ -0,0 +1,156 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! 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::deserialize; +use super::metadata; +use super::stored_values; +use crate::ArrayRef; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::Decimal; +use crate::arrays::DictArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::decimal::DecimalArrayExt; +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 and fails every copy to the host, so a read +/// that copies it fails. +#[derive(Debug, PartialEq, Eq, Hash)] +struct TestDeviceBuffer { + bytes: ByteBuffer, +} + +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 { + 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 never copies device buffers to the host") + } + + fn slice(&self, range: Range) -> Arc { + Arc::new(Self { + bytes: self.bytes.slice(range), + }) + } + + 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 })) +} + +/// 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(BoolArray::try_new_from_handle(bits, 0, 3, Validity::NonNullable)?.into_array()) +} + +/// 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 = DictArray::try_new( + codes.into_array(), + BoolArray::from_iter([true, false]).into_array(), + )?; + + Ok(validity.into_array()) +} + +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 canonical_node_keeps_device_buffer() -> VortexResult<()> { + let dtype = DType::Decimal(DecimalDType::new(18, 2), Nullability::Nullable); + + let canonical = deserialize( + &dtype, + 3, + &metadata(DecimalType::I64), + &[on_device(buffer![-99i64, 7, 99].into_byte_buffer())], + &[device_validity()?], + )?; + + assert!(canonical.as_::().buffer_handle().is_on_device()); + + 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), + &[values], + &[validity()?], + ); + + assert!(result.is_err_and(|error| error.to_string().contains("must be in host memory"))); + + 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..c3000dd962f --- /dev/null +++ b/vortex-array/src/arrays/decimal/vtable/plugin/tests/mod.rs @@ -0,0 +1,382 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! 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. + +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_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; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array::ArrayDeserialization; +use crate::array::ArrayPlugin; +use crate::arrays::BoolArray; +use crate::arrays::Decimal; +use crate::arrays::DecimalArray; +use crate::arrays::Widen; +use crate::arrays::decimal::DecimalArrayExt; +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::dtype::i256; +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) +} + +/// 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, + } + .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(()) +} + +#[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]), + "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(()) +} + +#[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 decoded = serialize_and_decode(&array, array.dtype())?; + + assert_eq!(decoded.as_::().values_type(), values_type); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); + + 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 decoded = serialize_and_decode( + &written, + &DType::Decimal(read_dtype, Nullability::NonNullable), + )?; + + assert!(decoded.is::()); + assert_arrays_eq!( + decoded, + canonical_reference(read_dtype, &values, None)?, + &mut 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..9cff25943df 100644 --- a/vortex-array/src/arrays/widen/encode/decimal.rs +++ b/vortex-array/src/arrays/widen/encode/decimal.rs @@ -99,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. /// - /// 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, 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: @@ -119,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>, @@ -150,7 +149,12 @@ impl WidenArray { return Ok(array.into_array()); } - let values = values.try_into_host_sync()?; + // 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); 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); 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 ); } }