diff --git a/encodings/fastlanes/tests/narrow.rs b/encodings/fastlanes/tests/narrow.rs new file mode 100644 index 00000000000..c99da2dbf81 --- /dev/null +++ b/encodings/fastlanes/tests/narrow.rs @@ -0,0 +1,69 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Narrow operations on compressed integer children. +//! +//! BitPacked children retain their stored width through selection and comparison. Canonical +//! execution widens them to Narrow's logical dtype. + +#![cfg(test)] + +use std::sync::LazyLock; + +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::fns::sum::sum; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::Narrow; +use vortex_array::arrays::NarrowArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::narrow::NarrowArraySlotsExt; +use vortex_array::assert_arrays_eq; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_fastlanes::BitPacked; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); + +#[test] +fn test_narrow_bitpacked_child() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let packed = BitPacked::encode(&buffer![0u8, 1, 3, 7].into_array(), 3, &mut ctx)?; + let array = NarrowArray::try_new(packed.into_array(), PType::U64.into())?.into_array(); + let selected = array.filter(Mask::from_iter([true, false, true, true]))?; + + assert!(selected.is::()); + assert_eq!(selected.as_::().values().dtype(), &PType::U8.into()); + assert_arrays_eq!(selected, buffer![0u64, 3, 7].into_array(), &mut ctx); + assert_arrays_eq!( + array.take(buffer![3u32, 1].into_array())?, + buffer![7u64, 1].into_array(), + &mut ctx + ); + assert_arrays_eq!( + array.binary( + ConstantArray::new(Scalar::from(3u64), 4).into_array(), + Operator::Lt + )?, + BoolArray::from_iter([true, true, false, false]), + &mut ctx + ); + assert_eq!(sum(&array, &mut ctx)?, Scalar::from(11u64).into_nullable()); + assert_eq!( + array.execute::(&mut ctx)?.as_slice::(), + &[0, 1, 3, 7] + ); + + Ok(()) +} diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index 65b00d81031..cf280b79516 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -47,6 +47,7 @@ use crate::arrays::chunked::compute::aggregate::ChunkedArrayAggregate; use crate::arrays::dict::compute::is_constant::DictIsConstantKernel; use crate::arrays::dict::compute::is_sorted::DictIsSortedKernel; use crate::arrays::dict::compute::min_max::DictMinMaxKernel; +use crate::arrays::narrow; use crate::dtype::DType; /// Session state for aggregate functions and encoding-specific aggregate kernels. @@ -122,6 +123,8 @@ impl Default for AggregateFnSession { this.register_aggregate_kernel(Dict.id(), Some(IsConstant.id()), &DictIsConstantKernel); this.register_aggregate_kernel(Dict.id(), Some(IsSorted.id()), &DictIsSortedKernel); + narrow::register_aggregate_kernels(&this); + // Register the built-in grouped aggregate kernels. this.register_grouped_kernel(Count.id(), &CountGroupedKernel); this.register_grouped_encoding_kernel( diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index f96eebead65..d366b1065a7 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -86,6 +86,10 @@ pub mod masked; pub use masked::Masked; pub use masked::MaskedArray; +pub mod narrow; +pub use narrow::Narrow; +pub use narrow::NarrowArray; + pub mod null; pub use null::Null; pub use null::NullArray; @@ -146,6 +150,7 @@ pub(crate) fn initialize(session: &VortexSession) { list::initialize(session); listview::initialize(session); map::initialize(session); + narrow::initialize(session); patched::initialize(session); primitive::initialize(session); struct_::initialize(session); diff --git a/vortex-array/src/arrays/narrow/aggregates.rs b/vortex-array/src/arrays/narrow/aggregates.rs new file mode 100644 index 00000000000..1e64fd95cba --- /dev/null +++ b/vortex-array/src/arrays/narrow/aggregates.rs @@ -0,0 +1,90 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Aggregate Narrow values without widening the input buffer. +//! +//! Supported aggregates preserve integer ordering or use wide sum states. Partial states are cast +//! to the logical dtype so later batches observe the same boundaries as canonical integers. + +use std::sync::LazyLock; + +use vortex_error::VortexResult; + +use super::Narrow; +use super::NarrowArraySlotsExt; +use crate::ArrayRef; +use crate::ArrayVTable; +use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateFnId; +use crate::aggregate_fn::AggregateFnRef; +use crate::aggregate_fn::AggregateFnVTable; +use crate::aggregate_fn::fns::all_non_null::AllNonNull; +use crate::aggregate_fn::fns::all_null::AllNull; +use crate::aggregate_fn::fns::count::Count; +use crate::aggregate_fn::fns::first::First; +use crate::aggregate_fn::fns::is_constant::IsConstant; +use crate::aggregate_fn::fns::is_sorted::IsSorted; +use crate::aggregate_fn::fns::last::Last; +use crate::aggregate_fn::fns::max::Max; +use crate::aggregate_fn::fns::min::Min; +use crate::aggregate_fn::fns::min_max::MinMax; +use crate::aggregate_fn::fns::null_count::NullCount; +use crate::aggregate_fn::fns::sum::Sum; +use crate::aggregate_fn::fns::sum_v2::SumV2; +use crate::aggregate_fn::kernels::DynAggregateKernel; +use crate::aggregate_fn::session::AggregateFnSession; +use crate::scalar::Scalar; + +pub(crate) fn register_aggregate_kernels(session: &AggregateFnSession) { + // One fallback avoids cloning the registry for each supported aggregate. + session.register_aggregate_kernel(Narrow.id(), None::, &NarrowAggregateKernel); +} + +static SUPPORTED_AGGREGATES: LazyLock<[AggregateFnId; 13]> = LazyLock::new(|| { + [ + MinMax.id(), + Min.id(), + Max.id(), + Sum.id(), + SumV2.id(), + IsConstant.id(), + IsSorted.id(), + First.id(), + Last.id(), + Count.id(), + NullCount.id(), + AllNull.id(), + AllNonNull.id(), + ] +}); + +#[derive(Debug)] +struct NarrowAggregateKernel; + +impl DynAggregateKernel for NarrowAggregateKernel { + fn aggregate( + &self, + aggregate_fn: &AggregateFnRef, + batch: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if !SUPPORTED_AGGREGATES.contains(&aggregate_fn.id()) { + return Ok(None); + } + + let Some(array) = batch.as_opt::() else { + return Ok(None); + }; + + // These aggregates preserve integer ordering and accumulate sums at i64/u64 width. + // Cast the partial state too, since it can contain extrema or boundary values. + let mut accumulator = aggregate_fn.accumulator(array.values().dtype())?; + accumulator.accumulate(array.values(), ctx)?; + let partial = accumulator.partial_scalar()?; + let Some(dtype) = aggregate_fn.state_dtype(batch.dtype()) else { + return Ok(None); + }; + + Ok(Some(partial.cast(&dtype)?)) + } +} diff --git a/vortex-array/src/arrays/narrow/compare.rs b/vortex-array/src/arrays/narrow/compare.rs new file mode 100644 index 00000000000..8cd34b642d4 --- /dev/null +++ b/vortex-array/src/arrays/narrow/compare.rs @@ -0,0 +1,122 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compare Narrow integers at a sufficient stored width. +//! +//! Two children use their wider stored dtype. Constants outside the stored range yield a known +//! comparison result while preserving the array validity. + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use super::Narrow; +use super::NarrowArraySlotsExt; +use crate::ArrayRef; +use crate::ArrayView; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; +use crate::arrays::Primitive; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::binary::Binary; +use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; +use crate::scalar_fn::fns::binary::CompareKernel; +use crate::scalar_fn::fns::binary::execute_compare; +use crate::scalar_fn::fns::operators::CompareOperator; + +pub(super) fn initialize(session: &VortexSession) { + session.kernels().register_execute_parent_kernel( + Binary.id(), + Narrow, + CompareExecuteAdaptor(Narrow), + ); +} + +impl CompareKernel for Narrow { + fn compare( + lhs: ArrayView<'_, Self>, + rhs: &ArrayRef, + operator: CompareOperator, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if !lhs.dtype().eq_ignore_nullability(rhs.dtype()) { + return Ok(None); + } + + if let Some(rhs) = rhs.as_opt::() { + let ptype = if lhs.values().dtype().as_ptype().byte_width() + >= rhs.values().dtype().as_ptype().byte_width() + { + lhs.values().dtype().as_ptype() + } else { + rhs.values().dtype().as_ptype() + }; + let lhs_dtype = DType::Primitive(ptype, lhs.dtype().nullability()); + let rhs_dtype = DType::Primitive(ptype, rhs.dtype().nullability()); + + return compare_values( + &lhs.values().cast(lhs_dtype)?, + &rhs.values().cast(rhs_dtype)?, + operator, + ctx, + ) + .map(Some); + } + + let Some(constant) = rhs.as_constant() else { + return Ok(None); + }; + + let nullability = lhs.dtype().nullability() | rhs.dtype().nullability(); + let storage_dtype = lhs.values().dtype().with_nullability(nullability); + if let Ok(value) = constant.cast(&storage_dtype) { + return compare_values( + lhs.values(), + &ConstantArray::new(value, lhs.len()).into_array(), + operator, + ctx, + ) + .map(Some); + } + + // An integer outside the child's type range is ordered against every stored value. + let below = lhs.dtype().is_signed_int() && i64::try_from(&constant)? < 0; + let result = match operator { + CompareOperator::Eq => false, + CompareOperator::NotEq => true, + CompareOperator::Lt | CompareOperator::Lte => !below, + CompareOperator::Gt | CompareOperator::Gte => below, + }; + let validity = lhs + .values() + .validity()? + .cast_nullability(nullability, lhs.len(), ctx)?; + + Ok(Some( + BoolArray::try_new( + BitBuffer::full_in(result, lhs.len(), ctx.allocator().clone()), + validity, + )? + .into_array(), + )) + } +} + +fn compare_values( + lhs: &ArrayRef, + rhs: &ArrayRef, + operator: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if lhs.is::() && (rhs.is::() || rhs.is::()) { + execute_compare(lhs, rhs, operator, ctx) + } else { + lhs.binary(rhs.clone(), operator.into()) + } +} diff --git a/vortex-array/src/arrays/narrow/encoding.rs b/vortex-array/src/arrays/narrow/encoding.rs new file mode 100644 index 00000000000..26307c514c4 --- /dev/null +++ b/vortex-array/src/arrays/narrow/encoding.rs @@ -0,0 +1,69 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Select the stored width of an integer array. +//! +//! Bounds exclude null payloads. Narrowing changes the physical representation while preserving +//! the logical dtype and validity. + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use super::NarrowArray; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::fns::min_max::min_max; +use crate::arrays::PrimitiveArray; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::PType; + +impl NarrowArray { + /// Stores an integer array at the smallest fitting width while retaining its logical dtype. + /// + /// Uses the non-null bounds and preserves signedness and validity. Empty and all-null arrays + /// use the smallest width of the same signedness. Returns the input unchanged if no narrower + /// type fits. Floating-point arrays are rejected. + pub fn encode(array: PrimitiveArray, ctx: &mut ExecutionCtx) -> VortexResult { + let ptype = array.ptype(); + vortex_ensure!( + ptype.is_int(), + "Narrow requires integer values, got {ptype}" + ); + + if ptype.byte_width() == 1 { + return Ok(array.into_array()); + } + + let bounds = min_max(array.as_ref(), ctx, NumericalAggregateOpts::default())?; + let candidates = if ptype.is_signed_int() { + [PType::I8, PType::I16, PType::I32] + } else { + [PType::U8, PType::U16, PType::U32] + }; + + for candidate in candidates { + if candidate.byte_width() >= ptype.byte_width() { + break; + } + + let storage_dtype = DType::Primitive(candidate, array.dtype().nullability()); + if bounds.as_ref().is_some_and(|bounds| { + bounds.min.cast(&storage_dtype).is_err() || bounds.max.cast(&storage_dtype).is_err() + }) { + continue; + } + + let values = array + .as_ref() + .cast(storage_dtype)? + .execute::(ctx)?; + + return Ok(Self::try_new(values.into_array(), array.dtype().clone())?.into_array()); + } + + Ok(array.into_array()) + } +} diff --git a/vortex-array/src/arrays/narrow/mod.rs b/vortex-array/src/arrays/narrow/mod.rs new file mode 100644 index 00000000000..980530db737 --- /dev/null +++ b/vortex-array/src/arrays/narrow/mod.rs @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Integer arrays that retain their logical dtype while storing a narrower integer child. +//! +//! [`NarrowArray`] stores no buffers of its own and derives validity from its child. Selection, +//! comparison, and aggregate operations can use that child without widening to the logical dtype. +//! Canonical execution returns a widening cast to +//! [`PrimitiveArray`](crate::arrays::PrimitiveArray). + +mod aggregates; +pub(crate) use aggregates::register_aggregate_kernels; + +mod compare; +mod encoding; +mod rules; +mod vtable; + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use crate::Array; +use crate::ArrayParts; +use crate::ArrayRef; +use crate::EmptyArrayData; +use crate::array_slots; +use crate::dtype::DType; +use crate::dtype::PType; + +/// An integer encoding with a narrower child of the same signedness. +#[derive(Clone, Debug)] +pub struct Narrow; + +/// An integer array whose logical dtype is wider than its stored values. +pub type NarrowArray = Array; + +/// The child of a [`NarrowArray`]. +#[array_slots(Narrow)] +pub struct NarrowSlots { + /// Integer values with the same signedness and nullability as the logical dtype. + #[slot(0)] + pub values: ArrayRef, +} + +impl NarrowArray { + /// Wraps an integer child with a wider logical dtype without reading its values. + /// + /// Both dtypes must be integers with the same signedness and nullability. The child must be + /// strictly narrower than `dtype` and may use any encoding. Nested Narrow arrays are flattened. + pub fn try_new(mut values: ArrayRef, dtype: DType) -> VortexResult { + validate_dtypes(values.dtype(), &dtype)?; + + while let Some(inner) = values.as_opt::() { + values = inner.values().clone(); + } + + let len = values.len(); + Self::try_from_parts(ArrayParts::new( + Narrow, + dtype, + len, + EmptyArrayData, + NarrowSlots { values }.into_slots(), + )) + } +} + +pub(super) fn validate_dtypes(storage: &DType, logical: &DType) -> VortexResult<()> { + let storage_ptype = PType::try_from(storage)?; + let logical_ptype = PType::try_from(logical)?; + vortex_ensure!( + storage_ptype.is_int() && logical_ptype.is_int(), + "Narrow requires integer dtypes, got {storage} and {logical}" + ); + vortex_ensure_eq!(storage_ptype.is_signed_int(), logical_ptype.is_signed_int()); + vortex_ensure!( + storage_ptype.byte_width() < logical_ptype.byte_width(), + "Narrow requires storage narrower than {logical}, got {storage}" + ); + vortex_ensure_eq!(storage.nullability(), logical.nullability()); + + Ok(()) +} + +pub(crate) fn initialize(session: &VortexSession) { + compare::initialize(session); +} + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/arrays/narrow/rules.rs b/vortex-array/src/arrays/narrow/rules.rs new file mode 100644 index 00000000000..c3002658454 --- /dev/null +++ b/vortex-array/src/arrays/narrow/rules.rs @@ -0,0 +1,108 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Push selection and compatible scalar operations into Narrow values. +//! +//! Each reduction preserves the logical integer width and derives result nullability from the +//! transformed child. Operations that require wider values fall back to canonical execution. + +use std::ops::Range; + +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use super::Narrow; +use super::NarrowArray; +use super::NarrowArraySlotsExt; +use crate::ArrayRef; +use crate::ArrayView; +use crate::IntoArray; +use crate::arrays::dict::TakeReduce; +use crate::arrays::dict::TakeReduceAdaptor; +use crate::arrays::filter::FilterReduce; +use crate::arrays::filter::FilterReduceAdaptor; +use crate::arrays::slice::SliceReduce; +use crate::arrays::slice::SliceReduceAdaptor; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::optimizer::rules::ParentRuleSet; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::cast::CastReduce; +use crate::scalar_fn::fns::cast::CastReduceAdaptor; +use crate::scalar_fn::fns::fill_null::FillNullReduce; +use crate::scalar_fn::fns::fill_null::FillNullReduceAdaptor; +use crate::scalar_fn::fns::mask::MaskReduce; +use crate::scalar_fn::fns::mask::MaskReduceAdaptor; + +pub(super) const RULES: ParentRuleSet = ParentRuleSet::new(&[ + ParentRuleSet::lift(&CastReduceAdaptor(Narrow)), + ParentRuleSet::lift(&FillNullReduceAdaptor(Narrow)), + ParentRuleSet::lift(&FilterReduceAdaptor(Narrow)), + ParentRuleSet::lift(&MaskReduceAdaptor(Narrow)), + ParentRuleSet::lift(&SliceReduceAdaptor(Narrow)), + ParentRuleSet::lift(&TakeReduceAdaptor(Narrow)), +]); + +fn rewrap(array: ArrayView<'_, Narrow>, values: ArrayRef) -> VortexResult> { + let dtype = array.dtype().with_nullability(values.dtype().nullability()); + Ok(Some(NarrowArray::try_new(values, dtype)?.into_array())) +} + +impl SliceReduce for Narrow { + fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { + rewrap(array, array.values().slice(range)?) + } +} + +impl FilterReduce for Narrow { + fn filter(array: ArrayView<'_, Self>, mask: &Mask) -> VortexResult> { + rewrap(array, array.values().filter(mask.clone())?) + } +} + +impl TakeReduce for Narrow { + fn take(array: ArrayView<'_, Self>, indices: &ArrayRef) -> VortexResult> { + rewrap(array, array.values().take(indices.clone())?) + } +} + +impl MaskReduce for Narrow { + fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult> { + rewrap(array, array.values().clone().mask(mask.clone())?) + } +} + +impl CastReduce for Narrow { + fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { + if !dtype.is_int() || dtype.is_signed_int() != array.dtype().is_signed_int() { + return Ok(None); + } + + let storage_dtype = array.values().dtype().with_nullability(dtype.nullability()); + if dtype.as_ptype().byte_width() <= storage_dtype.as_ptype().byte_width() { + return array.values().cast(dtype.clone()).map(Some); + } + + Ok(Some( + NarrowArray::try_new(array.values().cast(storage_dtype)?, dtype.clone())?.into_array(), + )) + } +} + +impl FillNullReduce for Narrow { + fn fill_null( + array: ArrayView<'_, Self>, + fill_value: &Scalar, + ) -> VortexResult> { + let storage_dtype = array + .values() + .dtype() + .with_nullability(fill_value.dtype().nullability()); + let Ok(fill_value) = fill_value.cast(&storage_dtype) else { + // A valid logical value may require wider storage. Let canonical execution handle it. + return Ok(None); + }; + + rewrap(array, array.values().fill_null(fill_value)?) + } +} diff --git a/vortex-array/src/arrays/narrow/tests.rs b/vortex-array/src/arrays/narrow/tests.rs new file mode 100644 index 00000000000..b6e7da65519 --- /dev/null +++ b/vortex-array/src/arrays/narrow/tests.rs @@ -0,0 +1,461 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_buffer::Buffer; +use vortex_buffer::ByteBufferMut; +use vortex_buffer::buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::registry::ReadContext; + +use super::Narrow; +use super::NarrowArray; +use super::NarrowArraySlotsExt; +use crate::ArrayContext; +use crate::ArrayParts; +use crate::ArrayRef; +use crate::EmptyArrayData; +use crate::IntoArray; +use crate::VTable; +use crate::VortexSessionExecute; +use crate::aggregate_fn::AggregateFnVTableExt; +use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::fns::is_constant::IsConstant; +use crate::aggregate_fn::fns::is_sorted::IsSorted; +use crate::aggregate_fn::fns::is_sorted::IsSortedOptions; +use crate::aggregate_fn::fns::min_max::MinMax; +use crate::aggregate_fn::fns::sum::Sum; +use crate::aggregate_fn::fns::sum_v2::SumV2; +use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; +use crate::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; +use crate::aggregate_fn::session::AggregateFnSessionExt; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::DictArray; +use crate::arrays::PrimitiveArray; +use crate::assert_arrays_eq; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::optimizer::ArrayOptimizer; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::operators::Operator; +use crate::serde::SerializeOptions; +use crate::serde::SerializedArray; +use crate::validity::Validity; + +fn roundtrip(array: &ArrayRef) -> VortexResult { + let session = array_session(); + let array_ctx = ArrayContext::empty(); + let buffers = array.serialize(&array_ctx, &session, &SerializeOptions::default())?; + let mut bytes = ByteBufferMut::empty(); + for buffer in buffers { + 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::i8_limits(vec![-128i64, 127], PType::I8)] +#[case::i16_lower(vec![-129i64, 0], PType::I16)] +#[case::i16_upper(vec![0i64, 128], PType::I16)] +#[case::i32_lower(vec![-32769i64, 0], PType::I32)] +#[case::i32_upper(vec![0i64, 32768], PType::I32)] +#[case::i64_lower(vec![i64::from(i32::MIN) - 1, 0], PType::I64)] +#[case::i64_upper(vec![0i64, i64::from(i32::MAX) + 1], PType::I64)] +#[case::positive_signed(vec![0i64, 255], PType::I16)] +fn test_signed_encoding(#[case] values: Vec, #[case] storage: PType) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let input = PrimitiveArray::from_iter(values); + let encoded = NarrowArray::encode(input.clone(), &mut ctx)?; + let stored_ptype = encoded.as_opt::().map_or_else( + || encoded.dtype().as_ptype(), + |array| array.values().dtype().as_ptype(), + ); + + assert_eq!(stored_ptype, storage); + assert_eq!(encoded.dtype(), input.dtype()); + assert_arrays_eq!(encoded, input.into_array(), &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::u8(vec![0u64, 255], PType::U8)] +#[case::u16(vec![256u64, 65535], PType::U16)] +#[case::u32(vec![65536u64, u64::from(u32::MAX)], PType::U32)] +#[case::u64(vec![u64::from(u32::MAX) + 1, u64::MAX], PType::U64)] +fn test_unsigned_encoding(#[case] values: Vec, #[case] storage: PType) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let input = PrimitiveArray::from_iter(values); + let encoded = NarrowArray::encode(input.clone(), &mut ctx)?; + let stored_ptype = encoded.as_opt::().map_or_else( + || encoded.dtype().as_ptype(), + |array| array.values().dtype().as_ptype(), + ); + + assert_eq!(stored_ptype, storage); + assert_eq!(encoded.dtype(), input.dtype()); + assert_arrays_eq!(encoded, input.into_array(), &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::empty(PrimitiveArray::new(Buffer::::empty(), Validity::NonNullable))] +#[case::all_null(PrimitiveArray::new(buffer![i64::MIN, i64::MAX], Validity::AllInvalid))] +#[case::null_payloads(PrimitiveArray::new( + buffer![i64::MIN, 42i64, i64::MAX], + Validity::from_iter([false, true, false]), +))] +fn test_null_payloads_and_empty(#[case] input: PrimitiveArray) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let encoded = NarrowArray::encode(input.clone(), &mut ctx)?; + + assert_eq!( + encoded.as_::().values().dtype().as_ptype(), + PType::I8 + ); + assert_arrays_eq!(encoded, input.into_array(), &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::equal_width(PType::I8, PType::I8)] +#[case::wider_storage(PType::I16, PType::I8)] +#[case::signedness(PType::I8, PType::U64)] +#[case::float_storage(PType::F32, PType::I64)] +#[case::float_logical(PType::I8, PType::F64)] +fn test_reject_invalid_dtypes(#[case] storage: PType, #[case] logical: PType) { + let child = crate::Canonical::empty(&storage.into()).into_array(); + assert!(NarrowArray::try_new(child, logical.into()).is_err()); +} + +#[test] +fn test_validate_slots() { + assert!( + NarrowArray::try_new( + buffer![1i8].into_array(), + DType::Primitive(PType::I64, Nullability::Nullable), + ) + .is_err() + ); + for slots in [ + vec![], + vec![None], + vec![Some(buffer![1i8].into_array()), None], + ] { + assert!( + NarrowArray::try_from_parts(ArrayParts::new( + Narrow, + PType::I64.into(), + 1, + EmptyArrayData, + slots.into() + ),) + .is_err() + ); + } + assert!( + NarrowArray::try_from_parts(ArrayParts::new( + Narrow, + PType::I64.into(), + 2, + EmptyArrayData, + vec![Some(buffer![1i8].into_array())].into() + ),) + .is_err() + ); +} + +#[test] +fn test_flatten_and_scalar_dtype() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let child = PrimitiveArray::from_option_iter([Some(-128i8), None, Some(127)]).into_array(); + let inner = NarrowArray::try_new(child, DType::Primitive(PType::I16, Nullability::Nullable))?; + let outer = NarrowArray::try_new( + inner.into_array(), + DType::Primitive(PType::I64, Nullability::Nullable), + )?; + + assert_eq!(outer.values().dtype().as_ptype(), PType::I8); + assert_eq!( + outer.as_ref().execute_scalar(0, &mut ctx)?, + Scalar::primitive(-128i64, Nullability::Nullable) + ); + assert_eq!( + outer.as_ref().execute_scalar(1, &mut ctx)?, + Scalar::null(outer.dtype().clone()) + ); + + Ok(()) +} + +#[test] +fn test_selection_and_fill_null() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let child = + PrimitiveArray::from_option_iter([Some(-12i8), None, Some(34), Some(56)]).into_array(); + let array = NarrowArray::try_new(child, DType::Primitive(PType::I64, Nullability::Nullable))? + .into_array(); + let selected = array + .slice(1..4)? + .take(buffer![2u32, 0, 1].into_array())? + .filter(Mask::from_iter([true, false, true]))? + .optimize()?; + + assert!(selected.is::()); + assert_eq!( + selected.as_::().values().dtype().as_ptype(), + PType::I8 + ); + assert_arrays_eq!( + selected, + buffer![56i64, 34] + .into_array() + .cast(array.dtype().clone())?, + &mut ctx + ); + + let masked = array + .clone() + .mask(BoolArray::from_iter([true, true, false, true]).into_array())? + .optimize()?; + assert!(masked.is::()); + assert_arrays_eq!( + masked, + PrimitiveArray::from_option_iter([Some(-12i64), None, None, Some(56)]), + &mut ctx + ); + + let filled = array.fill_null(Scalar::from(100i64))?; + assert!(filled.is::()); + assert_arrays_eq!(filled, buffer![-12i64, 100, 34, 56].into_array(), &mut ctx); + assert_arrays_eq!( + array.fill_null(Scalar::from(1000i64))?, + buffer![-12i64, 1000, 34, 56].into_array(), + &mut ctx + ); + + Ok(()) +} + +#[test] +fn test_cast_preserves_checked_nullability() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let child = PrimitiveArray::from_option_iter([Some(1i8), None]).into_array(); + let array = NarrowArray::try_new(child, DType::Primitive(PType::I32, Nullability::Nullable))? + .into_array(); + let widened = array.cast(DType::Primitive(PType::I64, Nullability::Nullable))?; + assert!(widened.is::()); + assert!( + array + .cast(PType::I64.into())? + .execute::(&mut ctx) + .is_err() + ); + assert_arrays_eq!( + array.cast(DType::Primitive(PType::I8, Nullability::Nullable))?, + PrimitiveArray::from_option_iter([Some(1i8), None]), + &mut ctx + ); + + Ok(()) +} + +#[rstest] +#[case::nonnullable(Nullability::NonNullable)] +#[case::nullable(Nullability::Nullable)] +fn test_fill_null_preserves_result_dtype(#[case] nullability: Nullability) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let child = PrimitiveArray::from_option_iter([Some(1i8), None]).into_array(); + let dtype = DType::Primitive(PType::I64, Nullability::Nullable); + let array = NarrowArray::try_new(child, dtype)?.into_array(); + let filled = array.fill_null(Scalar::primitive(2i64, nullability))?; + let expected_dtype = DType::Primitive(PType::I64, nullability); + + assert_eq!(filled.dtype(), &expected_dtype); + assert!(filled.is::()); + assert_arrays_eq!( + filled, + buffer![1i64, 2].into_array().cast(expected_dtype)?, + &mut ctx + ); + + Ok(()) +} + +#[rstest] +#[case::eq(Operator::Eq)] +#[case::ne(Operator::NotEq)] +#[case::lt(Operator::Lt)] +#[case::le(Operator::Lte)] +#[case::gt(Operator::Gt)] +#[case::ge(Operator::Gte)] +fn test_comparison(#[case] operator: Operator) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let input = + PrimitiveArray::from_option_iter([Some(-128i64), None, Some(0), Some(127)]).into_array(); + let array = NarrowArray::encode(input.clone().execute::(&mut ctx)?, &mut ctx)?; + for value in [-129i64, -128, 0, 127, 128] { + let constant = ConstantArray::new(Scalar::from(value), input.len()).into_array(); + assert_arrays_eq!( + array.binary(constant.clone(), operator)?, + input.binary(constant.clone(), operator)?, + &mut ctx + ); + assert_arrays_eq!( + constant.binary(array.clone(), operator)?, + constant.binary(input.clone(), operator)?, + &mut ctx + ); + } + + let rhs = NarrowArray::try_new( + buffer![0i16, 300, -300, 127].into_array(), + PType::I64.into(), + )? + .into_array(); + let expected_rhs = rhs + .clone() + .execute::(&mut ctx)? + .into_array(); + assert_arrays_eq!( + array.binary(rhs, operator)?, + input.binary(expected_rhs, operator)?, + &mut ctx + ); + + Ok(()) +} + +#[test] +fn test_aggregate_partial_states() -> VortexResult<()> { + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let batches = [ + vec![Some(-128i64), None, Some(42)], + vec![Some(127i64), None], + vec![None, None], + ]; + let fns = [ + MinMax.bind(NumericalAggregateOpts::default()), + Sum.bind(NumericalAggregateOpts::default()), + SumV2.bind(NumericalAggregateOpts::default()), + IsConstant.bind(EmptyOptions), + IsSorted.bind(IsSortedOptions { strict: false }), + ]; + for aggregate in fns { + assert!( + session + .aggregate_fns() + .find_aggregate_kernel(Narrow.id(), aggregate.id()) + .is_some() + ); + + let dtype = DType::Primitive(PType::I64, Nullability::Nullable); + let mut expected = aggregate.accumulator(&dtype)?; + let mut actual = aggregate.accumulator(&dtype)?; + for values in &batches { + let input = PrimitiveArray::from_option_iter(values.clone()); + let encoded = NarrowArray::encode(input.clone(), &mut ctx)?; + expected.accumulate(&input.into_array(), &mut ctx)?; + actual.accumulate(&encoded, &mut ctx)?; + assert_eq!(actual.partial_scalar()?, expected.partial_scalar()?); + } + assert_eq!(actual.final_scalar()?, expected.final_scalar()?); + } + + Ok(()) +} + +#[test] +fn test_unsupported_aggregate_keeps_logical_width() -> VortexResult<()> { + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let array = + NarrowArray::try_new(buffer![-1i8, 0, 1].into_array(), PType::I64.into())?.into_array(); + let aggregate = UncompressedSizeInBytes.bind(EmptyOptions); + let kernel = session + .aggregate_fns() + .find_aggregate_kernel(Narrow.id(), aggregate.id()) + .vortex_expect("Narrow registers a fallback aggregate kernel"); + assert!(kernel.aggregate(&aggregate, &array, &mut ctx)?.is_none()); + assert_eq!( + uncompressed_size_in_bytes(&array, &mut ctx)?, + 3 * size_of::() + ); + + Ok(()) +} + +#[test] +fn test_encoded_child_serde() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let values = DictArray::try_new( + buffer![2u8, 0, 1, 2].into_array(), + buffer![-1i8, 0, 1].into_array(), + )? + .into_array(); + let array = NarrowArray::try_new(values, PType::I64.into())?.into_array(); + let decoded = roundtrip(&array)?; + assert!(decoded.is::()); + assert_eq!( + decoded.as_::().values().dtype().as_ptype(), + PType::I8 + ); + assert_arrays_eq!(decoded, array, &mut ctx); + assert_eq!( + decoded + .execute::(&mut ctx)? + .as_slice::(), + &[1, -1, 0, 1] + ); + + Ok(()) +} + +#[test] +fn test_arithmetic_uses_logical_width() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = + NarrowArray::try_new(buffer![120i8, -120].into_array(), PType::I64.into())?.into_array(); + let rhs = ConstantArray::new(Scalar::from(20i64), 2).into_array(); + + assert_arrays_eq!( + array.binary(rhs, Operator::Add)?, + buffer![140i64, -100].into_array(), + &mut ctx + ); + + Ok(()) +} + +#[rstest] +#[case::missing(vec![])] +#[case::extra(vec![PType::I8 as u8, 0])] +#[case::unknown(vec![255])] +#[case::float(vec![PType::F32 as u8])] +#[case::unsigned(vec![PType::U8 as u8])] +#[case::equal_width(vec![PType::I64 as u8])] +fn test_reject_malformed_metadata(#[case] metadata: Vec) { + let session = array_session(); + let values = [buffer![1i8].into_array()]; + assert!( + Narrow + .deserialize(&PType::I64.into(), 1, &metadata, &[], &values, &session) + .is_err() + ); +} diff --git a/vortex-array/src/arrays/narrow/vtable.rs b/vortex-array/src/arrays/narrow/vtable.rs new file mode 100644 index 00000000000..3fd94cafc0a --- /dev/null +++ b/vortex-array/src/arrays/narrow/vtable.rs @@ -0,0 +1,167 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Implement the Narrow storage and execution contract. +//! +//! The values child owns the buffers and validity. Deserialization validates its dtype before +//! construction; canonical execution yields a cast to the logical integer width. + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use super::Narrow; +use super::NarrowArraySlotsExt; +use super::NarrowSlots; +use super::rules::RULES; +use super::validate_dtypes; +use crate::Array; +use crate::ArrayId; +use crate::ArrayParts; +use crate::ArrayRef; +use crate::ArrayView; +use crate::EmptyArrayData; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::VTable; +use crate::array::with_empty_buffers; +use crate::buffer::BufferHandle; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::PType; +use crate::scalar::Scalar; +use crate::serde::ArrayChildren; +use crate::vtable::OperationsVTable; +use crate::vtable::ValidityChild; +use crate::vtable::ValidityVTableFromChild; + +impl VTable for Narrow { + type TypedArrayData = EmptyArrayData; + type OperationsVTable = Self; + type ValidityVTable = ValidityVTableFromChild; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.narrow"); + *ID + } + + fn validate( + &self, + _data: &EmptyArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure_eq!(slots.len(), 1); + let values = slots[0] + .as_ref() + .ok_or_else(|| vortex_err!("Narrow requires a values child"))?; + validate_dtypes(values.dtype(), dtype)?; + vortex_ensure_eq!(values.len(), len); + + Ok(()) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("Narrow buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option { + None + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn serialize( + array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(Some(vec![PType::try_from(array.values().dtype())? as u8])) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure_eq!(buffers.len(), 0); + vortex_ensure_eq!(children.len(), 1); + let [storage_ptype] = metadata else { + return Err(vortex_err!( + "Narrow requires one metadata byte, got {}", + metadata.len() + )); + }; + + let storage_ptype = PType::try_from(i32::from(*storage_ptype)) + .map_err(|err| vortex_err!("Invalid Narrow storage type: {err}"))?; + let storage_dtype = DType::Primitive(storage_ptype, dtype.nullability()); + validate_dtypes(&storage_dtype, dtype)?; + + let values = children.get(0, &storage_dtype, len)?; + Ok(ArrayParts::new( + Self, + dtype.clone(), + len, + EmptyArrayData, + NarrowSlots { values }.into_slots(), + )) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + NarrowSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + // Yield a cast so the child can fuse decoding and widening in its own execution kernel. + Ok(ExecutionResult::done( + array.values().cast(array.dtype().clone())?, + )) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + RULES.evaluate(array, parent, child_idx) + } +} + +impl OperationsVTable for Narrow { + type ProbeState = (); + + fn scalar_at( + array: ArrayView<'_, Self>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + array + .values() + .execute_scalar(index, ctx)? + .cast(array.dtype()) + } +} + +impl ValidityChild for Narrow { + fn validity_child(array: ArrayView<'_, Self>) -> ArrayRef { + array.values().clone() + } +} diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index e9d0d36a55f..eb45c45c869 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -29,6 +29,7 @@ use crate::arrays::List; use crate::arrays::ListView; use crate::arrays::Map; use crate::arrays::Masked; +use crate::arrays::Narrow; use crate::arrays::Null; use crate::arrays::PiecewiseSequence; use crate::arrays::Primitive; @@ -105,6 +106,7 @@ impl Default for ArraySession { this.register(Dict); this.register(List); this.register(Masked); + this.register(Narrow); this.register(PiecewiseSequence); this.register(VarBin);