From f85700c45c9e6151161a9ce488dbd748ee102c73 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 2 Oct 2026 13:05:47 -0400 Subject: [PATCH 1/9] feat(array): add Narrow integer encoding Signed-off-by: Connor Tsui --- encodings/fastlanes/tests/narrow.rs | 58 +++ vortex-array/src/arrays/mod.rs | 5 + vortex-array/src/arrays/narrow/aggregates.rs | 80 ++++ vortex-array/src/arrays/narrow/compare.rs | 117 ++++++ vortex-array/src/arrays/narrow/encoding.rs | 57 +++ vortex-array/src/arrays/narrow/mod.rs | 97 +++++ vortex-array/src/arrays/narrow/rules.rs | 97 +++++ vortex-array/src/arrays/narrow/tests.rs | 384 +++++++++++++++++++ vortex-array/src/arrays/narrow/vtable.rs | 170 ++++++++ vortex-array/src/session/mod.rs | 2 + 10 files changed, 1067 insertions(+) create mode 100644 encodings/fastlanes/tests/narrow.rs create mode 100644 vortex-array/src/arrays/narrow/aggregates.rs create mode 100644 vortex-array/src/arrays/narrow/compare.rs create mode 100644 vortex-array/src/arrays/narrow/encoding.rs create mode 100644 vortex-array/src/arrays/narrow/mod.rs create mode 100644 vortex-array/src/arrays/narrow/rules.rs create mode 100644 vortex-array/src/arrays/narrow/tests.rs create mode 100644 vortex-array/src/arrays/narrow/vtable.rs diff --git a/encodings/fastlanes/tests/narrow.rs b/encodings/fastlanes/tests/narrow.rs new file mode 100644 index 00000000000..d5f224ad8f4 --- /dev/null +++ b/encodings/fastlanes/tests/narrow.rs @@ -0,0 +1,58 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::LazyLock; + +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::fns::sum::sum; +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)?, + buffer![true, true, false, false].into_array(), + &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/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..598359f83be --- /dev/null +++ b/vortex-array/src/arrays/narrow/aggregates.rs @@ -0,0 +1,80 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use super::Narrow; +use super::NarrowArraySlotsExt; +use crate::ArrayRef; +use crate::ArrayVTable; +use crate::ExecutionCtx; +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::AggregateFnSessionExt; +use crate::scalar::Scalar; + +pub(super) fn initialize(session: &VortexSession) { + for id in [ + 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(), + ] { + session.aggregate_fns().register_aggregate_kernel( + Narrow.id(), + Some(id), + &NarrowAggregateKernel, + ); + } +} + +#[derive(Debug)] +struct NarrowAggregateKernel; + +impl DynAggregateKernel for NarrowAggregateKernel { + fn aggregate( + &self, + aggregate_fn: &AggregateFnRef, + batch: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + 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..f502f3c6c25 --- /dev/null +++ b/vortex-array/src/arrays/narrow/compare.rs @@ -0,0 +1,117 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +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..df0c6ca0577 --- /dev/null +++ b/vortex-array/src/arrays/narrow/encoding.rs @@ -0,0 +1,57 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +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::arrays::primitive::PrimitiveArrayExt; +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..93189ce5281 --- /dev/null +++ b/vortex-array/src/arrays/narrow/mod.rs @@ -0,0 +1,97 @@ +// 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`]. + +mod encoding; + +mod vtable; + +mod rules; + +mod compare; + +mod aggregates; + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_session::VortexSession; + +use crate::Array; +use crate::ArrayParts; +use crate::ArrayRef; +use crate::EmptyArrayData; +use crate::array_slots; +use crate::arrays::PrimitiveArray; +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) + .with_slots(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!( + storage_ptype.is_signed_int() == logical_ptype.is_signed_int(), + "Narrow requires matching signedness, got {storage} and {logical}" + ); + vortex_ensure!( + storage_ptype.byte_width() < logical_ptype.byte_width(), + "Narrow requires storage narrower than {logical}, got {storage}" + ); + vortex_ensure!( + storage.nullability() == logical.nullability(), + "Narrow requires matching nullability, got {storage} and {logical}" + ); + + Ok(()) +} + +pub(crate) fn initialize(session: &VortexSession) { + compare::initialize(session); + aggregates::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..dde3964c1a2 --- /dev/null +++ b/vortex-array/src/arrays/narrow/rules.rs @@ -0,0 +1,97 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +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().as_nonnullable(); + 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..01871c80411 --- /dev/null +++ b/vortex-array/src/arrays/narrow/tests.rs @@ -0,0 +1,384 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_buffer::Buffer; +use vortex_buffer::ByteBufferMut; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; +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::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::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; + +static SESSION: LazyLock = LazyLock::new(crate::array_session); + +fn roundtrip(array: &ArrayRef) -> VortexResult { + 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 = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_iter(values); + let encoded = NarrowArray::encode(input.clone(), &mut ctx)?; + let values = encoded.as_opt::().map_or(&encoded, |array| array.values()); + + assert_eq!(values.dtype().as_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 = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_iter(values); + let encoded = NarrowArray::encode(input.clone(), &mut ctx)?; + let values = encoded.as_opt::().map_or(&encoded, |array| array.values()); + + assert_eq!(values.dtype().as_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 = 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) + .with_slots(slots.into()), + ) + .is_err() + ); + } + assert!( + NarrowArray::try_from_parts( + ArrayParts::new(Narrow, PType::I64.into(), 2, EmptyArrayData) + .with_slots(vec![Some(buffer![1i8].into_array())].into()), + ) + .is_err() + ); +} + +#[test] +fn test_flatten_and_scalar_dtype() -> VortexResult<()> { + let mut ctx = 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 = 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(buffer![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 = 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::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 = 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 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(crate::aggregate_fn::EmptyOptions), + IsSorted.bind(IsSortedOptions::default()), + ]; + for aggregate in fns { + 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_encoded_child_serde() -> VortexResult<()> { + let mut ctx = 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.clone(), 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 = 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 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..ec435e518af --- /dev/null +++ b/vortex-array/src/arrays/narrow/vtable.rs @@ -0,0 +1,170 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use 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!( + slots.len() == 1, + "Narrow requires one child, got {}", + slots.len() + ); + let values = slots[0] + .as_ref() + .ok_or_else(|| vortex_err!("Narrow requires a values child"))?; + validate_dtypes(values.dtype(), dtype)?; + vortex_ensure!( + values.len() == len, + "Narrow requires child length {len}, got {}", + values.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!( + buffers.is_empty(), + "Narrow requires no buffers, got {}", + buffers.len() + ); + vortex_ensure!( + children.len() == 1, + "Narrow requires one child, got {}", + children.len() + ); + 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) + .with_slots(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); From f5387a0bc665ad774b9c31175f3d1e5941b603b9 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 2 Oct 2026 14:24:26 -0400 Subject: [PATCH 2/9] fix(array): preserve Narrow fill result nullability A valid nullable fill scalar still requires a nullable result dtype. Preserve that tag when casting the fill scalar into the stored child dtype. Signed-off-by: Connor Tsui --- vortex-array/src/arrays/narrow/encoding.rs | 1 - vortex-array/src/arrays/narrow/rules.rs | 5 ++++- vortex-array/src/arrays/narrow/tests.rs | 21 +++++++++++++++++++++ 3 files changed, 25 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/arrays/narrow/encoding.rs b/vortex-array/src/arrays/narrow/encoding.rs index df0c6ca0577..a818977f7da 100644 --- a/vortex-array/src/arrays/narrow/encoding.rs +++ b/vortex-array/src/arrays/narrow/encoding.rs @@ -11,7 +11,6 @@ use crate::IntoArray; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min_max::min_max; use crate::arrays::PrimitiveArray; -use crate::arrays::primitive::PrimitiveArrayExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::PType; diff --git a/vortex-array/src/arrays/narrow/rules.rs b/vortex-array/src/arrays/narrow/rules.rs index dde3964c1a2..a48a0a64b5e 100644 --- a/vortex-array/src/arrays/narrow/rules.rs +++ b/vortex-array/src/arrays/narrow/rules.rs @@ -86,7 +86,10 @@ impl CastReduce for Narrow { impl FillNullReduce for Narrow { fn fill_null(array: ArrayView<'_, Self>, fill_value: &Scalar) -> VortexResult> { - let storage_dtype = array.values().dtype().as_nonnullable(); + 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); diff --git a/vortex-array/src/arrays/narrow/tests.rs b/vortex-array/src/arrays/narrow/tests.rs index 01871c80411..bc1e653cdbd 100644 --- a/vortex-array/src/arrays/narrow/tests.rs +++ b/vortex-array/src/arrays/narrow/tests.rs @@ -254,6 +254,27 @@ fn test_cast_preserves_checked_nullability() -> VortexResult<()> { Ok(()) } +#[rstest] +#[case(Nullability::NonNullable)] +#[case(Nullability::Nullable)] +fn test_fill_null_preserves_result_dtype(#[case] nullability: Nullability) -> VortexResult<()> { + let mut ctx = 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)] From 3db36321215fa5cc50ed61020c69987ec7eb190f Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 2 Oct 2026 14:53:01 -0400 Subject: [PATCH 3/9] fix(array): compile Narrow regression tests Signed-off-by: Connor Tsui --- encodings/fastlanes/tests/narrow.rs | 3 ++- vortex-array/src/arrays/narrow/mod.rs | 4 ++-- vortex-array/src/arrays/narrow/tests.rs | 17 +++++++++++------ 3 files changed, 15 insertions(+), 9 deletions(-) diff --git a/encodings/fastlanes/tests/narrow.rs b/encodings/fastlanes/tests/narrow.rs index d5f224ad8f4..04e1f69e1b0 100644 --- a/encodings/fastlanes/tests/narrow.rs +++ b/encodings/fastlanes/tests/narrow.rs @@ -6,6 +6,7 @@ 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; @@ -45,7 +46,7 @@ fn test_narrow_bitpacked_child() -> VortexResult<()> { ); assert_arrays_eq!( array.binary(ConstantArray::new(Scalar::from(3u64), 4).into_array(), Operator::Lt)?, - buffer![true, true, false, false].into_array(), + BoolArray::from_iter([true, true, false, false]), &mut ctx ); assert_eq!(sum(&array, &mut ctx)?, Scalar::from(11u64).into_nullable()); diff --git a/vortex-array/src/arrays/narrow/mod.rs b/vortex-array/src/arrays/narrow/mod.rs index 93189ce5281..9bcc5103f4a 100644 --- a/vortex-array/src/arrays/narrow/mod.rs +++ b/vortex-array/src/arrays/narrow/mod.rs @@ -5,7 +5,8 @@ //! //! [`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`]. +//! Canonical execution returns a widening cast to +//! [`PrimitiveArray`](crate::arrays::PrimitiveArray). mod encoding; @@ -26,7 +27,6 @@ use crate::ArrayParts; use crate::ArrayRef; use crate::EmptyArrayData; use crate::array_slots; -use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::PType; diff --git a/vortex-array/src/arrays/narrow/tests.rs b/vortex-array/src/arrays/narrow/tests.rs index bc1e653cdbd..17643cbc041 100644 --- a/vortex-array/src/arrays/narrow/tests.rs +++ b/vortex-array/src/arrays/narrow/tests.rs @@ -30,6 +30,7 @@ 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::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; @@ -76,9 +77,11 @@ fn test_signed_encoding(#[case] values: Vec, #[case] storage: PType) -> Vor let mut ctx = SESSION.create_execution_ctx(); let input = PrimitiveArray::from_iter(values); let encoded = NarrowArray::encode(input.clone(), &mut ctx)?; - let values = encoded.as_opt::().map_or(&encoded, |array| array.values()); + let stored_ptype = encoded.as_opt::().map_or(encoded.dtype().as_ptype(), |array| { + array.values().dtype().as_ptype() + }); - assert_eq!(values.dtype().as_ptype(), storage); + assert_eq!(stored_ptype, storage); assert_eq!(encoded.dtype(), input.dtype()); assert_arrays_eq!(encoded, input.into_array(), &mut ctx); @@ -94,9 +97,11 @@ fn test_unsigned_encoding(#[case] values: Vec, #[case] storage: PType) -> V let mut ctx = SESSION.create_execution_ctx(); let input = PrimitiveArray::from_iter(values); let encoded = NarrowArray::encode(input.clone(), &mut ctx)?; - let values = encoded.as_opt::().map_or(&encoded, |array| array.values()); + let stored_ptype = encoded.as_opt::().map_or(encoded.dtype().as_ptype(), |array| { + array.values().dtype().as_ptype() + }); - assert_eq!(values.dtype().as_ptype(), storage); + assert_eq!(stored_ptype, storage); assert_eq!(encoded.dtype(), input.dtype()); assert_arrays_eq!(encoded, input.into_array(), &mut ctx); @@ -209,7 +214,7 @@ fn test_selection_and_fill_null() -> VortexResult<()> { let masked = array .clone() - .mask(buffer![true, true, false, true].into_array())? + .mask(BoolArray::from_iter([true, true, false, true]).into_array())? .optimize()?; assert!(masked.is::()); assert_arrays_eq!( @@ -332,7 +337,7 @@ fn test_aggregate_partial_states() -> VortexResult<()> { Sum.bind(NumericalAggregateOpts::default()), SumV2.bind(NumericalAggregateOpts::default()), IsConstant.bind(crate::aggregate_fn::EmptyOptions), - IsSorted.bind(IsSortedOptions::default()), + IsSorted.bind(IsSortedOptions { strict: false }), ]; for aggregate in fns { let dtype = DType::Primitive(PType::I64, Nullability::Nullable); From 0c196b385145ed42530127453b58e5382f4fbe5c Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 2 Oct 2026 15:00:03 -0400 Subject: [PATCH 4/9] style(array): clarify Narrow contracts and test fixtures Signed-off-by: Connor Tsui --- encodings/fastlanes/tests/narrow.rs | 7 +- vortex-array/src/arrays/narrow/aggregates.rs | 5 + vortex-array/src/arrays/narrow/compare.rs | 5 + vortex-array/src/arrays/narrow/encoding.rs | 15 ++- vortex-array/src/arrays/narrow/mod.rs | 10 +- vortex-array/src/arrays/narrow/rules.rs | 10 +- vortex-array/src/arrays/narrow/tests.rs | 115 +++++++++++-------- vortex-array/src/arrays/narrow/vtable.rs | 10 +- vortex-array/src/lib.rs | 4 + 9 files changed, 120 insertions(+), 61 deletions(-) diff --git a/encodings/fastlanes/tests/narrow.rs b/encodings/fastlanes/tests/narrow.rs index 04e1f69e1b0..777b9614459 100644 --- a/encodings/fastlanes/tests/narrow.rs +++ b/encodings/fastlanes/tests/narrow.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +#![cfg(test)] + use std::sync::LazyLock; use vortex_array::IntoArray; @@ -45,7 +47,10 @@ fn test_narrow_bitpacked_child() -> VortexResult<()> { &mut ctx ); assert_arrays_eq!( - array.binary(ConstantArray::new(Scalar::from(3u64), 4).into_array(), Operator::Lt)?, + array.binary( + ConstantArray::new(Scalar::from(3u64), 4).into_array(), + Operator::Lt + )?, BoolArray::from_iter([true, true, false, false]), &mut ctx ); diff --git a/vortex-array/src/arrays/narrow/aggregates.rs b/vortex-array/src/arrays/narrow/aggregates.rs index 598359f83be..20e0bf34816 100644 --- a/vortex-array/src/arrays/narrow/aggregates.rs +++ b/vortex-array/src/arrays/narrow/aggregates.rs @@ -1,6 +1,11 @@ // 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 vortex_error::VortexResult; use vortex_session::VortexSession; diff --git a/vortex-array/src/arrays/narrow/compare.rs b/vortex-array/src/arrays/narrow/compare.rs index f502f3c6c25..8cd34b642d4 100644 --- a/vortex-array/src/arrays/narrow/compare.rs +++ b/vortex-array/src/arrays/narrow/compare.rs @@ -1,6 +1,11 @@ // 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; diff --git a/vortex-array/src/arrays/narrow/encoding.rs b/vortex-array/src/arrays/narrow/encoding.rs index a818977f7da..223b7de33dc 100644 --- a/vortex-array/src/arrays/narrow/encoding.rs +++ b/vortex-array/src/arrays/narrow/encoding.rs @@ -1,6 +1,11 @@ // 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; @@ -23,7 +28,10 @@ impl NarrowArray { /// 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}"); + vortex_ensure!( + ptype.is_int(), + "Narrow requires integer values, got {ptype}" + ); if ptype.byte_width() == 1 { return Ok(array.into_array()); } @@ -47,7 +55,10 @@ impl NarrowArray { continue; } - let values = array.as_ref().cast(storage_dtype)?.execute::(ctx)?; + let values = array + .as_ref() + .cast(storage_dtype)? + .execute::(ctx)?; return Ok(Self::try_new(values.into_array(), array.dtype().clone())?.into_array()); } diff --git a/vortex-array/src/arrays/narrow/mod.rs b/vortex-array/src/arrays/narrow/mod.rs index 9bcc5103f4a..d6412fcd165 100644 --- a/vortex-array/src/arrays/narrow/mod.rs +++ b/vortex-array/src/arrays/narrow/mod.rs @@ -8,15 +8,11 @@ //! Canonical execution returns a widening cast to //! [`PrimitiveArray`](crate::arrays::PrimitiveArray). +mod aggregates; +mod compare; mod encoding; - -mod vtable; - mod rules; - -mod compare; - -mod aggregates; +mod vtable; use vortex_error::VortexResult; use vortex_error::vortex_ensure; diff --git a/vortex-array/src/arrays/narrow/rules.rs b/vortex-array/src/arrays/narrow/rules.rs index a48a0a64b5e..c3002658454 100644 --- a/vortex-array/src/arrays/narrow/rules.rs +++ b/vortex-array/src/arrays/narrow/rules.rs @@ -1,6 +1,11 @@ // 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; @@ -85,7 +90,10 @@ impl CastReduce for Narrow { } impl FillNullReduce for Narrow { - fn fill_null(array: ArrayView<'_, Self>, fill_value: &Scalar) -> VortexResult> { + fn fill_null( + array: ArrayView<'_, Self>, + fill_value: &Scalar, + ) -> VortexResult> { let storage_dtype = array .values() .dtype() diff --git a/vortex-array/src/arrays/narrow/tests.rs b/vortex-array/src/arrays/narrow/tests.rs index 17643cbc041..52617155272 100644 --- a/vortex-array/src/arrays/narrow/tests.rs +++ b/vortex-array/src/arrays/narrow/tests.rs @@ -1,15 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::sync::LazyLock; - use rstest::rstest; use vortex_buffer::Buffer; use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_mask::Mask; -use vortex_session::VortexSession; use vortex_session::registry::ReadContext; use super::Narrow; @@ -20,6 +17,7 @@ use crate::ArrayParts; use crate::ArrayRef; use crate::EmptyArrayData; use crate::IntoArray; +use crate::TEST_SESSION; use crate::VTable; use crate::VortexSessionExecute; use crate::aggregate_fn::AggregateFnVTableExt; @@ -46,11 +44,9 @@ use crate::serde::SerializeOptions; use crate::serde::SerializedArray; use crate::validity::Validity; -static SESSION: LazyLock = LazyLock::new(crate::array_session); - fn roundtrip(array: &ArrayRef) -> VortexResult { let array_ctx = ArrayContext::empty(); - let buffers = array.serialize(&array_ctx, &SESSION, &SerializeOptions::default())?; + let buffers = array.serialize(&array_ctx, &TEST_SESSION, &SerializeOptions::default())?; let mut bytes = ByteBufferMut::empty(); for buffer in buffers { bytes.extend_from_slice(buffer.as_ref()); @@ -60,7 +56,7 @@ fn roundtrip(array: &ArrayRef) -> VortexResult { array.dtype(), array.len(), &ReadContext::new(array_ctx.to_ids()), - &SESSION, + &TEST_SESSION, ) } @@ -74,12 +70,13 @@ fn roundtrip(array: &ArrayRef) -> VortexResult { #[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 = SESSION.create_execution_ctx(); + let mut ctx = TEST_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(encoded.dtype().as_ptype(), |array| { - array.values().dtype().as_ptype() - }); + 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()); @@ -94,12 +91,13 @@ fn test_signed_encoding(#[case] values: Vec, #[case] storage: PType) -> Vor #[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 = SESSION.create_execution_ctx(); + let mut ctx = TEST_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(encoded.dtype().as_ptype(), |array| { - array.values().dtype().as_ptype() - }); + 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()); @@ -116,10 +114,13 @@ fn test_unsigned_encoding(#[case] values: Vec, #[case] storage: PType) -> V Validity::from_iter([false, true, false]), ))] fn test_null_payloads_and_empty(#[case] input: PrimitiveArray) -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); + let mut ctx = TEST_SESSION.create_execution_ctx(); let encoded = NarrowArray::encode(input.clone(), &mut ctx)?; - assert_eq!(encoded.as_::().values().dtype().as_ptype(), PType::I8); + assert_eq!( + encoded.as_::().values().dtype().as_ptype(), + PType::I8 + ); assert_arrays_eq!(encoded, input.into_array(), &mut ctx); Ok(()) @@ -169,10 +170,9 @@ fn test_validate_slots() { #[test] fn test_flatten_and_scalar_dtype() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); + let mut ctx = TEST_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 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), @@ -193,11 +193,11 @@ fn test_flatten_and_scalar_dtype() -> VortexResult<()> { #[test] fn test_selection_and_fill_null() -> VortexResult<()> { - let mut ctx = 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 mut ctx = TEST_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())? @@ -205,10 +205,15 @@ fn test_selection_and_fill_null() -> VortexResult<()> { .optimize()?; assert!(selected.is::()); - assert_eq!(selected.as_::().values().dtype().as_ptype(), PType::I8); + assert_eq!( + selected.as_::().values().dtype().as_ptype(), + PType::I8 + ); assert_arrays_eq!( selected, - buffer![56i64, 34].into_array().cast(array.dtype().clone())?, + buffer![56i64, 34] + .into_array() + .cast(array.dtype().clone())?, &mut ctx ); @@ -237,11 +242,10 @@ fn test_selection_and_fill_null() -> VortexResult<()> { #[test] fn test_cast_preserves_checked_nullability() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); + let mut ctx = TEST_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 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!( @@ -260,10 +264,10 @@ fn test_cast_preserves_checked_nullability() -> VortexResult<()> { } #[rstest] -#[case(Nullability::NonNullable)] -#[case(Nullability::Nullable)] +#[case::nonnullable(Nullability::NonNullable)] +#[case::nullable(Nullability::Nullable)] fn test_fill_null_preserves_result_dtype(#[case] nullability: Nullability) -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); + let mut ctx = TEST_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(); @@ -277,6 +281,7 @@ fn test_fill_null_preserves_result_dtype(#[case] nullability: Nullability) -> Vo buffer![1i64, 2].into_array().cast(expected_dtype)?, &mut ctx ); + Ok(()) } @@ -288,13 +293,10 @@ fn test_fill_null_preserves_result_dtype(#[case] nullability: Nullability) -> Vo #[case::gt(Operator::Gt)] #[case::ge(Operator::Gte)] fn test_comparison(#[case] operator: Operator) -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); + let mut ctx = TEST_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, - )?; + 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!( @@ -314,7 +316,10 @@ fn test_comparison(#[case] operator: Operator) -> VortexResult<()> { PType::I64.into(), )? .into_array(); - let expected_rhs = rhs.clone().execute::(&mut ctx)?.into_array(); + let expected_rhs = rhs + .clone() + .execute::(&mut ctx)? + .into_array(); assert_arrays_eq!( array.binary(rhs, operator)?, input.binary(expected_rhs, operator)?, @@ -326,7 +331,7 @@ fn test_comparison(#[case] operator: Operator) -> VortexResult<()> { #[test] fn test_aggregate_partial_states() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); + let mut ctx = TEST_SESSION.create_execution_ctx(); let batches = [ vec![Some(-128i64), None, Some(42)], vec![Some(127i64), None], @@ -358,7 +363,7 @@ fn test_aggregate_partial_states() -> VortexResult<()> { #[test] fn test_encoded_child_serde() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); + let mut ctx = TEST_SESSION.create_execution_ctx(); let values = DictArray::try_new( buffer![2u8, 0, 1, 2].into_array(), buffer![-1i8, 0, 1].into_array(), @@ -367,10 +372,15 @@ fn test_encoded_child_serde() -> VortexResult<()> { 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.clone(), array, &mut ctx); assert_eq!( - decoded.execute::(&mut ctx)?.as_slice::(), + 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] ); @@ -379,9 +389,9 @@ fn test_encoded_child_serde() -> VortexResult<()> { #[test] fn test_arithmetic_uses_logical_width() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let array = NarrowArray::try_new(buffer![120i8, -120].into_array(), PType::I64.into())? - .into_array(); + let mut ctx = TEST_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!( @@ -404,7 +414,14 @@ fn test_reject_malformed_metadata(#[case] metadata: Vec) { let values = [buffer![1i8].into_array()]; assert!( Narrow - .deserialize(&PType::I64.into(), 1, &metadata, &[], &values, &SESSION) + .deserialize( + &PType::I64.into(), + 1, + &metadata, + &[], + &values, + &TEST_SESSION + ) .is_err() ); } diff --git a/vortex-array/src/arrays/narrow/vtable.rs b/vortex-array/src/arrays/narrow/vtable.rs index ec435e518af..26338dfd137 100644 --- a/vortex-array/src/arrays/narrow/vtable.rs +++ b/vortex-array/src/arrays/narrow/vtable.rs @@ -1,6 +1,11 @@ // 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; use vortex_error::vortex_err; @@ -159,7 +164,10 @@ impl OperationsVTable for Narrow { index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - array.values().execute_scalar(index, ctx)?.cast(array.dtype()) + array + .values() + .execute_scalar(index, ctx)? + .cast(array.dtype()) } } diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index fc64f01de52..5b5090fa5a3 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -157,6 +157,10 @@ pub fn initialize(session: &VortexSession) { } } +/// Shared default session for unit tests. +#[cfg(test)] +pub(crate) static TEST_SESSION: LazyLock = LazyLock::new(array_session); + /// Builds a fresh [`VortexSession`] registered with all of vortex-array's built-in session /// variables: arrays, dtypes, scalar functions, stats, optimizer kernels, aggregate functions, /// and memory. From d1e48317294eca1a2ae904d3a550bccf321975e0 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 2 Oct 2026 16:40:32 -0400 Subject: [PATCH 5/9] fix(array): register Narrow aggregates in the aggregate session Signed-off-by: Connor Tsui --- vortex-array/src/aggregate_fn/session.rs | 3 ++ vortex-array/src/arrays/narrow/aggregates.rs | 29 ++++++++++------- vortex-array/src/arrays/narrow/mod.rs | 2 +- vortex-array/src/arrays/narrow/tests.rs | 33 +++++++++++++++++++- 4 files changed, 53 insertions(+), 14 deletions(-) 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/narrow/aggregates.rs b/vortex-array/src/arrays/narrow/aggregates.rs index 20e0bf34816..1e64fd95cba 100644 --- a/vortex-array/src/arrays/narrow/aggregates.rs +++ b/vortex-array/src/arrays/narrow/aggregates.rs @@ -6,14 +6,16 @@ //! 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 vortex_session::VortexSession; 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; @@ -30,11 +32,16 @@ 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::AggregateFnSessionExt; +use crate::aggregate_fn::session::AggregateFnSession; use crate::scalar::Scalar; -pub(super) fn initialize(session: &VortexSession) { - for id in [ +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(), @@ -48,14 +55,8 @@ pub(super) fn initialize(session: &VortexSession) { NullCount.id(), AllNull.id(), AllNonNull.id(), - ] { - session.aggregate_fns().register_aggregate_kernel( - Narrow.id(), - Some(id), - &NarrowAggregateKernel, - ); - } -} + ] +}); #[derive(Debug)] struct NarrowAggregateKernel; @@ -67,6 +68,10 @@ impl DynAggregateKernel for NarrowAggregateKernel { 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); }; diff --git a/vortex-array/src/arrays/narrow/mod.rs b/vortex-array/src/arrays/narrow/mod.rs index d6412fcd165..8ac873dedc8 100644 --- a/vortex-array/src/arrays/narrow/mod.rs +++ b/vortex-array/src/arrays/narrow/mod.rs @@ -9,6 +9,7 @@ //! [`PrimitiveArray`](crate::arrays::PrimitiveArray). mod aggregates; +pub(crate) use aggregates::register_aggregate_kernels; mod compare; mod encoding; mod rules; @@ -86,7 +87,6 @@ pub(super) fn validate_dtypes(storage: &DType, logical: &DType) -> VortexResult< pub(crate) fn initialize(session: &VortexSession) { compare::initialize(session); - aggregates::initialize(session); } #[cfg(test)] diff --git a/vortex-array/src/arrays/narrow/tests.rs b/vortex-array/src/arrays/narrow/tests.rs index 52617155272..15971062e91 100644 --- a/vortex-array/src/arrays/narrow/tests.rs +++ b/vortex-array/src/arrays/narrow/tests.rs @@ -5,6 +5,7 @@ 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; @@ -21,6 +22,7 @@ use crate::TEST_SESSION; 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; @@ -28,6 +30,9 @@ 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::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::DictArray; @@ -341,10 +346,17 @@ fn test_aggregate_partial_states() -> VortexResult<()> { MinMax.bind(NumericalAggregateOpts::default()), Sum.bind(NumericalAggregateOpts::default()), SumV2.bind(NumericalAggregateOpts::default()), - IsConstant.bind(crate::aggregate_fn::EmptyOptions), + IsConstant.bind(EmptyOptions), IsSorted.bind(IsSortedOptions { strict: false }), ]; for aggregate in fns { + assert!( + TEST_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)?; @@ -361,6 +373,25 @@ fn test_aggregate_partial_states() -> VortexResult<()> { Ok(()) } +#[test] +fn test_unsupported_aggregate_keeps_logical_width() -> VortexResult<()> { + let mut ctx = TEST_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 = TEST_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 = TEST_SESSION.create_execution_ctx(); From c287456d0deee0c79228ebffd61f493ab5ecdc26 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 5 Oct 2026 11:51:29 +0100 Subject: [PATCH 6/9] test(array): create Narrow sessions per test Signed-off-by: Connor Tsui --- vortex-array/src/arrays/narrow/tests.rs | 47 ++++++++++++------------- vortex-array/src/lib.rs | 4 --- 2 files changed, 22 insertions(+), 29 deletions(-) diff --git a/vortex-array/src/arrays/narrow/tests.rs b/vortex-array/src/arrays/narrow/tests.rs index 15971062e91..64c8fe10910 100644 --- a/vortex-array/src/arrays/narrow/tests.rs +++ b/vortex-array/src/arrays/narrow/tests.rs @@ -18,7 +18,6 @@ use crate::ArrayParts; use crate::ArrayRef; use crate::EmptyArrayData; use crate::IntoArray; -use crate::TEST_SESSION; use crate::VTable; use crate::VortexSessionExecute; use crate::aggregate_fn::AggregateFnVTableExt; @@ -33,6 +32,7 @@ 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; @@ -50,8 +50,9 @@ 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, &TEST_SESSION, &SerializeOptions::default())?; + 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()); @@ -61,7 +62,7 @@ fn roundtrip(array: &ArrayRef) -> VortexResult { array.dtype(), array.len(), &ReadContext::new(array_ctx.to_ids()), - &TEST_SESSION, + &session, ) } @@ -75,7 +76,7 @@ fn roundtrip(array: &ArrayRef) -> VortexResult { #[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 = TEST_SESSION.create_execution_ctx(); + 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( @@ -96,7 +97,7 @@ fn test_signed_encoding(#[case] values: Vec, #[case] storage: PType) -> Vor #[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 = TEST_SESSION.create_execution_ctx(); + 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( @@ -119,7 +120,7 @@ fn test_unsigned_encoding(#[case] values: Vec, #[case] storage: PType) -> V Validity::from_iter([false, true, false]), ))] fn test_null_payloads_and_empty(#[case] input: PrimitiveArray) -> VortexResult<()> { - let mut ctx = TEST_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let encoded = NarrowArray::encode(input.clone(), &mut ctx)?; assert_eq!( @@ -175,7 +176,7 @@ fn test_validate_slots() { #[test] fn test_flatten_and_scalar_dtype() -> VortexResult<()> { - let mut ctx = TEST_SESSION.create_execution_ctx(); + 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( @@ -198,7 +199,7 @@ fn test_flatten_and_scalar_dtype() -> VortexResult<()> { #[test] fn test_selection_and_fill_null() -> VortexResult<()> { - let mut ctx = TEST_SESSION.create_execution_ctx(); + 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))? @@ -247,7 +248,7 @@ fn test_selection_and_fill_null() -> VortexResult<()> { #[test] fn test_cast_preserves_checked_nullability() -> VortexResult<()> { - let mut ctx = TEST_SESSION.create_execution_ctx(); + 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(); @@ -272,7 +273,7 @@ fn test_cast_preserves_checked_nullability() -> VortexResult<()> { #[case::nonnullable(Nullability::NonNullable)] #[case::nullable(Nullability::Nullable)] fn test_fill_null_preserves_result_dtype(#[case] nullability: Nullability) -> VortexResult<()> { - let mut ctx = TEST_SESSION.create_execution_ctx(); + 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(); @@ -298,7 +299,7 @@ fn test_fill_null_preserves_result_dtype(#[case] nullability: Nullability) -> Vo #[case::gt(Operator::Gt)] #[case::ge(Operator::Gte)] fn test_comparison(#[case] operator: Operator) -> VortexResult<()> { - let mut ctx = TEST_SESSION.create_execution_ctx(); + 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)?; @@ -336,7 +337,8 @@ fn test_comparison(#[case] operator: Operator) -> VortexResult<()> { #[test] fn test_aggregate_partial_states() -> VortexResult<()> { - let mut ctx = TEST_SESSION.create_execution_ctx(); + let session = array_session(); + let mut ctx = session.create_execution_ctx(); let batches = [ vec![Some(-128i64), None, Some(42)], vec![Some(127i64), None], @@ -351,7 +353,7 @@ fn test_aggregate_partial_states() -> VortexResult<()> { ]; for aggregate in fns { assert!( - TEST_SESSION + session .aggregate_fns() .find_aggregate_kernel(Narrow.id(), aggregate.id()) .is_some() @@ -375,11 +377,12 @@ fn test_aggregate_partial_states() -> VortexResult<()> { #[test] fn test_unsupported_aggregate_keeps_logical_width() -> VortexResult<()> { - let mut ctx = TEST_SESSION.create_execution_ctx(); + 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 = TEST_SESSION + let kernel = session .aggregate_fns() .find_aggregate_kernel(Narrow.id(), aggregate.id()) .vortex_expect("Narrow registers a fallback aggregate kernel"); @@ -394,7 +397,7 @@ fn test_unsupported_aggregate_keeps_logical_width() -> VortexResult<()> { #[test] fn test_encoded_child_serde() -> VortexResult<()> { - let mut ctx = TEST_SESSION.create_execution_ctx(); + 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(), @@ -420,7 +423,7 @@ fn test_encoded_child_serde() -> VortexResult<()> { #[test] fn test_arithmetic_uses_logical_width() -> VortexResult<()> { - let mut ctx = TEST_SESSION.create_execution_ctx(); + 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(); @@ -442,17 +445,11 @@ fn test_arithmetic_uses_logical_width() -> VortexResult<()> { #[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, - &TEST_SESSION - ) + .deserialize(&PType::I64.into(), 1, &metadata, &[], &values, &session) .is_err() ); } diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index 5b5090fa5a3..fc64f01de52 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -157,10 +157,6 @@ pub fn initialize(session: &VortexSession) { } } -/// Shared default session for unit tests. -#[cfg(test)] -pub(crate) static TEST_SESSION: LazyLock = LazyLock::new(array_session); - /// Builds a fresh [`VortexSession`] registered with all of vortex-array's built-in session /// variables: arrays, dtypes, scalar functions, stats, optimizer kernels, aggregate functions, /// and memory. From 083e95c53448b9f244f5f1872d0164a2c7aa90d4 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 5 Oct 2026 16:21:48 +0100 Subject: [PATCH 7/9] style(array): use equality checks for Narrow invariants Signed-off-by: Connor Tsui --- vortex-array/src/arrays/narrow/mod.rs | 11 +++------- vortex-array/src/arrays/narrow/vtable.rs | 26 +++++------------------- 2 files changed, 8 insertions(+), 29 deletions(-) diff --git a/vortex-array/src/arrays/narrow/mod.rs b/vortex-array/src/arrays/narrow/mod.rs index 8ac873dedc8..22df1a7d065 100644 --- a/vortex-array/src/arrays/narrow/mod.rs +++ b/vortex-array/src/arrays/narrow/mod.rs @@ -17,6 +17,7 @@ mod vtable; use vortex_error::VortexResult; use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; use vortex_session::VortexSession; use crate::Array; @@ -69,18 +70,12 @@ pub(super) fn validate_dtypes(storage: &DType, logical: &DType) -> VortexResult< storage_ptype.is_int() && logical_ptype.is_int(), "Narrow requires integer dtypes, got {storage} and {logical}" ); - vortex_ensure!( - storage_ptype.is_signed_int() == logical_ptype.is_signed_int(), - "Narrow requires matching signedness, 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!( - storage.nullability() == logical.nullability(), - "Narrow requires matching nullability, got {storage} and {logical}" - ); + vortex_ensure_eq!(storage.nullability(), logical.nullability()); Ok(()) } diff --git a/vortex-array/src/arrays/narrow/vtable.rs b/vortex-array/src/arrays/narrow/vtable.rs index 26338dfd137..be2b572cc78 100644 --- a/vortex-array/src/arrays/narrow/vtable.rs +++ b/vortex-array/src/arrays/narrow/vtable.rs @@ -7,7 +7,7 @@ //! construction; canonical execution yields a cast to the logical integer width. use vortex_error::VortexResult; -use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; @@ -55,20 +55,12 @@ impl VTable for Narrow { len: usize, slots: &[Option], ) -> VortexResult<()> { - vortex_ensure!( - slots.len() == 1, - "Narrow requires one child, got {}", - slots.len() - ); + 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!( - values.len() == len, - "Narrow requires child length {len}, got {}", - values.len() - ); + vortex_ensure_eq!(values.len(), len); Ok(()) } @@ -109,16 +101,8 @@ impl VTable for Narrow { children: &dyn ArrayChildren, _session: &VortexSession, ) -> VortexResult> { - vortex_ensure!( - buffers.is_empty(), - "Narrow requires no buffers, got {}", - buffers.len() - ); - vortex_ensure!( - children.len() == 1, - "Narrow requires one child, got {}", - children.len() - ); + 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 {}", From cc903d480cc675e3c41b0f952a743cb21d5984ae Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 6 Oct 2026 14:34:47 +0100 Subject: [PATCH 8/9] fix(array): initialize Narrow parts with child slots Signed-off-by: Connor Tsui --- vortex-array/src/arrays/narrow/mod.rs | 11 +++++++---- vortex-array/src/arrays/narrow/tests.rs | 22 ++++++++++++++-------- vortex-array/src/arrays/narrow/vtable.rs | 9 +++++++-- 3 files changed, 28 insertions(+), 14 deletions(-) diff --git a/vortex-array/src/arrays/narrow/mod.rs b/vortex-array/src/arrays/narrow/mod.rs index 22df1a7d065..eea555854fb 100644 --- a/vortex-array/src/arrays/narrow/mod.rs +++ b/vortex-array/src/arrays/narrow/mod.rs @@ -56,10 +56,13 @@ impl NarrowArray { } let len = values.len(); - Self::try_from_parts( - ArrayParts::new(Narrow, dtype, len, EmptyArrayData) - .with_slots(NarrowSlots { values }.into_slots()), - ) + Self::try_from_parts(ArrayParts::new( + Narrow, + dtype, + len, + EmptyArrayData, + NarrowSlots { values }.into_slots(), + )) } } diff --git a/vortex-array/src/arrays/narrow/tests.rs b/vortex-array/src/arrays/narrow/tests.rs index 64c8fe10910..b6e7da65519 100644 --- a/vortex-array/src/arrays/narrow/tests.rs +++ b/vortex-array/src/arrays/narrow/tests.rs @@ -158,18 +158,24 @@ fn test_validate_slots() { vec![Some(buffer![1i8].into_array()), None], ] { assert!( - NarrowArray::try_from_parts( - ArrayParts::new(Narrow, PType::I64.into(), 1, EmptyArrayData) - .with_slots(slots.into()), - ) + 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) - .with_slots(vec![Some(buffer![1i8].into_array())].into()), - ) + NarrowArray::try_from_parts(ArrayParts::new( + Narrow, + PType::I64.into(), + 2, + EmptyArrayData, + vec![Some(buffer![1i8].into_array())].into() + ),) .is_err() ); } diff --git a/vortex-array/src/arrays/narrow/vtable.rs b/vortex-array/src/arrays/narrow/vtable.rs index be2b572cc78..3fd94cafc0a 100644 --- a/vortex-array/src/arrays/narrow/vtable.rs +++ b/vortex-array/src/arrays/narrow/vtable.rs @@ -116,8 +116,13 @@ impl VTable for Narrow { validate_dtypes(&storage_dtype, dtype)?; let values = children.get(0, &storage_dtype, len)?; - Ok(ArrayParts::new(Self, dtype.clone(), len, EmptyArrayData) - .with_slots(NarrowSlots { values }.into_slots())) + Ok(ArrayParts::new( + Self, + dtype.clone(), + len, + EmptyArrayData, + NarrowSlots { values }.into_slots(), + )) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { From 44780e80bbe3579afadeb0f3b3b83f2c99a246c4 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 6 Oct 2026 15:07:04 +0100 Subject: [PATCH 9/9] style(array): clarify Narrow module boundaries Signed-off-by: Connor Tsui --- encodings/fastlanes/tests/narrow.rs | 5 +++++ vortex-array/src/arrays/narrow/encoding.rs | 2 ++ vortex-array/src/arrays/narrow/mod.rs | 1 + 3 files changed, 8 insertions(+) diff --git a/encodings/fastlanes/tests/narrow.rs b/encodings/fastlanes/tests/narrow.rs index 777b9614459..c99da2dbf81 100644 --- a/encodings/fastlanes/tests/narrow.rs +++ b/encodings/fastlanes/tests/narrow.rs @@ -1,6 +1,11 @@ // 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; diff --git a/vortex-array/src/arrays/narrow/encoding.rs b/vortex-array/src/arrays/narrow/encoding.rs index 223b7de33dc..26307c514c4 100644 --- a/vortex-array/src/arrays/narrow/encoding.rs +++ b/vortex-array/src/arrays/narrow/encoding.rs @@ -32,6 +32,7 @@ impl NarrowArray { ptype.is_int(), "Narrow requires integer values, got {ptype}" ); + if ptype.byte_width() == 1 { return Ok(array.into_array()); } @@ -59,6 +60,7 @@ impl NarrowArray { .as_ref() .cast(storage_dtype)? .execute::(ctx)?; + return Ok(Self::try_new(values.into_array(), array.dtype().clone())?.into_array()); } diff --git a/vortex-array/src/arrays/narrow/mod.rs b/vortex-array/src/arrays/narrow/mod.rs index eea555854fb..980530db737 100644 --- a/vortex-array/src/arrays/narrow/mod.rs +++ b/vortex-array/src/arrays/narrow/mod.rs @@ -10,6 +10,7 @@ mod aggregates; pub(crate) use aggregates::register_aggregate_kernels; + mod compare; mod encoding; mod rules;