From 2f132a9ab9d17b9e3d110b4ea173a4e3abdacd1e Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 27 May 2026 16:26:06 +0100 Subject: [PATCH] fix Signed-off-by: Joe Isaacs --- .../fastlanes/src/rle/array/rle_decompress.rs | 54 +++++------ encodings/sparse/src/canonical.rs | 11 ++- vortex-array/src/arrays/decimal/array.rs | 11 ++- .../arrays/fixed_size_list/compute/take.rs | 10 +- vortex-array/src/arrays/list/compute/take.rs | 29 ++++-- vortex-array/src/arrays/listview/array.rs | 44 ++++++--- vortex-array/src/arrays/listview/rebuild.rs | 91 +++++++++++++++++-- .../src/arrays/varbin/compute/take.rs | 88 +++++++++--------- vortex-array/src/patches.rs | 16 +++- .../src/scalar_fn/fns/list_contains/mod.rs | 11 ++- 10 files changed, 241 insertions(+), 124 deletions(-) diff --git a/encodings/fastlanes/src/rle/array/rle_decompress.rs b/encodings/fastlanes/src/rle/array/rle_decompress.rs index 7240f3ff71f..51766398070 100644 --- a/encodings/fastlanes/src/rle/array/rle_decompress.rs +++ b/encodings/fastlanes/src/rle/array/rle_decompress.rs @@ -19,35 +19,44 @@ use crate::RLEArray; use crate::rle::RLEArrayExt; /// Decompresses an RLE array back into a primitive array. -#[expect( - clippy::cognitive_complexity, - reason = "complexity is from nested match_each_* macros" -)] pub fn rle_decompress(array: &RLEArray, ctx: &mut ExecutionCtx) -> VortexResult { + // The per-chunk value-index offsets are tiny (one entry per 1024-element chunk), so cast them + // to `u64` once here instead of monomorphizing the whole decode loop over the offset width. + let values_idx_offsets = array + .values_idx_offsets() + .clone() + .execute::(ctx)?; + let values_idx_offsets: Vec = + match_each_unsigned_integer_ptype!(values_idx_offsets.ptype(), |O| { + values_idx_offsets + .as_slice::() + .iter() + .map(|&o| o.as_()) + .collect() + }); + match_each_native_ptype!(array.values().dtype().as_ptype(), |V| { - match_each_unsigned_integer_ptype!(array.values_idx_offsets().dtype().as_ptype(), |O| { - // RLE indices are always u16 (or u8 if downcasted). - match array.indices().dtype().as_ptype() { - PType::U8 => rle_decode_typed::(array, ctx), - PType::U16 => rle_decode_typed::(array, ctx), - _ => vortex_panic!( - "Unsupported index type for RLE decoding: {}", - array.indices().dtype().as_ptype() - ), - } - }) + // RLE indices are always u16 (or u8 if downcasted). + match array.indices().dtype().as_ptype() { + PType::U8 => rle_decode_typed::(array, &values_idx_offsets, ctx), + PType::U16 => rle_decode_typed::(array, &values_idx_offsets, ctx), + _ => vortex_panic!( + "Unsupported index type for RLE decoding: {}", + array.indices().dtype().as_ptype() + ), + } }) } /// Decompresses an `RLEArray` into to a primitive array of unsigned integers. -fn rle_decode_typed( +fn rle_decode_typed( array: &RLEArray, + values_idx_offsets: &[u64], ctx: &mut ExecutionCtx, ) -> VortexResult where V: NativePType + RLE + Clone + Copy, I: NativePType + Into, - O: NativePType + AsPrimitive, { let values = array.values().clone().execute::(ctx)?; let values = values.as_slice::(); @@ -64,23 +73,16 @@ where let mut buffer = BufferMut::::with_capacity(num_chunks * FL_CHUNK_SIZE); let (out_buf, _) = buffer.spare_capacity_mut().as_chunks_mut::(); - let values_idx_offsets = array - .values_idx_offsets() - .clone() - .execute::(ctx)?; - let values_idx_offsets = values_idx_offsets.as_slice::(); - for (chunk_idx, (chunk_indices, chunk_out)) in indices_sl.iter().zip(out_buf.iter_mut()).enumerate() { // Offsets in `values_idx_offsets` are absolute and need to be shifted // by the offset of the first chunk, respective of the current slice, // to make them relative. - let value_idx_offset = - (values_idx_offsets[chunk_idx].as_() - values_idx_offsets[0].as_()) as usize; + let value_idx_offset = (values_idx_offsets[chunk_idx] - values_idx_offsets[0]) as usize; let next_value_idx_offset = if chunk_idx + 1 < num_chunks { - (values_idx_offsets[chunk_idx + 1].as_() - values_idx_offsets[0].as_()) as usize + (values_idx_offsets[chunk_idx + 1] - values_idx_offsets[0]) as usize } else { values.len() }; diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index 91052a63681..ac8c072e7b1 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -22,6 +22,7 @@ use vortex_array::arrays::VarBinView; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::arrays::varbinview::build_views::BinaryView; use vortex_array::buffer::BufferHandle; @@ -41,6 +42,7 @@ use vortex_array::dtype::StructFields; use vortex_array::match_each_decimal_value_type; use vortex_array::match_each_integer_ptype; use vortex_array::match_each_native_ptype; +use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::match_smallest_offset_type; use vortex_array::patches::Patches; use vortex_array::scalar::DecimalScalar; @@ -163,10 +165,6 @@ pub(super) fn execute_sparse(parts: SparseParts, ctx: &mut ExecutionCtx) -> Vort }) } -#[expect( - clippy::cognitive_complexity, - reason = "complexity is from nested match_smallest_offset_type macro" -)] fn execute_sparse_lists( resolved: &Patches, fill_value: &Scalar, @@ -175,14 +173,17 @@ fn execute_sparse_lists( nullability: Nullability, ctx: &mut ExecutionCtx, ) -> VortexResult { + // Patch indices are non-negative; reinterpret to unsigned so this dispatches over 4 widths + // instead of 8. `O` is already unsigned (from `match_smallest_offset_type`). let indices = resolved.indices().as_::().into_owned(); + let indices = indices.reinterpret_cast(indices.ptype().to_unsigned()); let values = resolved.values().as_::().into_owned(); let fill_list = fill_value.as_list(); let n_filled = len - resolved.num_patches(); let total_canonical_values = values.elements().len() + fill_list.len() * n_filled; - Ok(match_each_integer_ptype!(indices.ptype(), |I| { + Ok(match_each_unsigned_integer_ptype!(indices.ptype(), |I| { match_smallest_offset_type!(total_canonical_values, |O| { execute_sparse_lists_inner::( indices.as_slice(), diff --git a/vortex-array/src/arrays/decimal/array.rs b/vortex-array/src/arrays/decimal/array.rs index 68f929be319..1f4d6c0c277 100644 --- a/vortex-array/src/arrays/decimal/array.rs +++ b/vortex-array/src/arrays/decimal/array.rs @@ -28,6 +28,7 @@ use crate::array::validity_to_child; use crate::arrays::Decimal; use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; +use crate::arrays::primitive::PrimitiveArrayExt; use crate::buffer::BufferHandle; use crate::dtype::BigCast; use crate::dtype::DType; @@ -37,7 +38,7 @@ use crate::dtype::IntegerPType; use crate::dtype::NativeDecimalType; use crate::dtype::Nullability; use crate::match_each_decimal_value_type; -use crate::match_each_integer_ptype; +use crate::match_each_unsigned_integer_ptype; use crate::patches::Patches; use crate::validity::Validity; @@ -556,8 +557,12 @@ impl Array { assert_eq!(self.decimal_dtype(), patch_values.decimal_dtype()); let data = self.into_data(); - let data = match_each_integer_ptype!(patch_indices.ptype(), |I| { - let patch_indices = patch_indices.as_slice::(); + // Patch indices are non-negative; reinterpret to unsigned so this dispatches over 4 widths + // instead of 8 (the decimal value-type dimensions are unaffected). + let patch_indices_unsigned = + patch_indices.reinterpret_cast(patch_indices.ptype().to_unsigned()); + let data = match_each_unsigned_integer_ptype!(patch_indices_unsigned.ptype(), |I| { + let patch_indices = patch_indices_unsigned.as_slice::(); match_each_decimal_value_type!(patch_values.values_type(), |PatchDVT| { let patch_values = patch_values.buffer::(); match_each_decimal_value_type!(data.values_type(), |ValuesDVT| { diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 3f2cf451c00..8df3bae0b0c 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -19,9 +19,10 @@ use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; use crate::arrays::dict::TakeExecute; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::primitive::PrimitiveArrayExt; use crate::dtype::IntegerPType; use crate::executor::ExecutionCtx; -use crate::match_each_integer_ptype; +use crate::match_each_unsigned_integer_ptype; use crate::match_smallest_offset_type; use crate::validity::Validity; @@ -31,14 +32,15 @@ use crate::validity::Validity; /// that elements start at offset 0 and be perfectly packed without gaps. We expand list indices /// into element indices and push them down to the child elements array. impl TakeExecute for FixedSizeList { - #[expect(clippy::cognitive_complexity)] fn take( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { let max_element_idx = array.elements().len(); - match_each_integer_ptype!(indices.dtype().as_ptype(), |I| { + // Indices are non-negative; dispatch over the 4 unsigned widths (the executed array is + // reinterpreted to unsigned in `take_with_indices`). `E` is already unsigned. + match_each_unsigned_integer_ptype!(indices.dtype().as_ptype().to_unsigned(), |I| { match_smallest_offset_type!(max_element_idx, |E| { take_with_indices::(array, indices, ctx) }) @@ -56,6 +58,8 @@ fn take_with_indices( let list_size = array.list_size() as usize; let indices_array = indices.clone().execute::(ctx)?; + // Reinterpret to unsigned so `as_slice::` (with unsigned `I`) matches; values are unchanged. + let indices_array = indices_array.reinterpret_cast(indices_array.ptype().to_unsigned()); // Make sure to handle degenerate case where lists have size 0 (these can take fast paths). if list_size == 0 { diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 7cc5ccd78cd..43117b1b834 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -19,7 +19,7 @@ use crate::builders::PrimitiveBuilder; use crate::dtype::IntegerPType; use crate::dtype::Nullability; use crate::executor::ExecutionCtx; -use crate::match_each_integer_ptype; +use crate::match_each_unsigned_integer_ptype; use crate::match_smallest_offset_type; // TODO(connor)[ListView]: Re-revert to the version where we simply convert to a `ListView` and call @@ -38,16 +38,22 @@ impl TakeExecute for List { ctx: &mut ExecutionCtx, ) -> VortexResult> { let indices = indices.clone().execute::(ctx)?; + let indices = indices.reinterpret_cast(indices.ptype().to_unsigned()); + let offsets = array.offsets().clone().execute::(ctx)?; + let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); // This is an over-approximation of the total number of elements in the resulting array. let total_approx = array.elements().len().saturating_mul(indices.len()); - match_each_integer_ptype!(array.offsets().dtype().as_ptype(), |O| { - match_each_integer_ptype!(indices.ptype(), |I| { + match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + match_each_unsigned_integer_ptype!(indices.ptype(), |I| { match_smallest_offset_type!(total_approx, |OutputOffsetType| { - { - let indices = indices.as_view(); - _take::(array, indices, ctx).map(Some) - } + _take::( + array, + offsets.as_view(), + indices.as_view(), + ctx, + ) + .map(Some) }) }) }) @@ -56,6 +62,7 @@ impl TakeExecute for List { fn _take( array: ArrayView<'_, List>, + offsets_array: ArrayView<'_, Primitive>, indices_array: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> VortexResult { @@ -68,10 +75,9 @@ fn _take( .execute_mask(indices_array.as_ref().len(), ctx)?; if !indices_validity.all_true() || !data_validity.all_true() { - return _take_nullable::(array, indices_array, ctx); + return _take_nullable::(array, offsets_array, indices_array, ctx); } - let offsets_array = array.offsets().clone().execute::(ctx)?; let offsets: &[O] = offsets_array.as_slice(); let indices: &[I] = indices_array.as_slice(); @@ -121,12 +127,15 @@ fn _take( .into_array()) } +// Kept out-of-line: as a single-callsite generic helper it would otherwise be inlined into every +// monomorphization of `_take`, duplicating the entire nullable path across all specializations. +#[inline(never)] fn _take_nullable( array: ArrayView<'_, List>, + offsets_array: ArrayView<'_, Primitive>, indices_array: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> VortexResult { - let offsets_array = array.offsets().clone().execute::(ctx)?; let offsets: &[O] = offsets_array.as_slice(); let indices: &[I] = indices_array.as_slice(); let data_validity = array diff --git a/vortex-array/src/arrays/listview/array.rs b/vortex-array/src/arrays/listview/array.rs index 3e7aedc5f5e..64e59d1687f 100644 --- a/vortex-array/src/arrays/listview/array.rs +++ b/vortex-array/src/arrays/listview/array.rs @@ -31,10 +31,12 @@ use crate::arrays::ListView; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::bool; +use crate::arrays::primitive::PrimitiveArrayExt; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::expr::stats::Stat; use crate::match_each_integer_ptype; +use crate::match_each_unsigned_integer_ptype; use crate::validity::Validity; /// The `elements` data array, where each list scalar is a _slice_ of the `elements` array, and @@ -241,10 +243,6 @@ impl ListViewData { sizes.len() ); - // Check that the size type can fit within the offset type to prevent overflows. - let size_ptype = sizes.dtype().as_ptype(); - let offset_ptype = offsets.dtype().as_ptype(); - // If a validity array is present, it must be the same length as the `ListViewArray`. if let Some(validity_len) = validity.maybe_len() { vortex_ensure!( @@ -260,10 +258,16 @@ impl ListViewData { let offsets_primitive = offsets.to_primitive(); #[expect(deprecated)] let sizes_primitive = sizes.to_primitive(); + // Offsets and sizes are non-negative; reinterpret to unsigned to dispatch over 4 widths + // each (4x4 instead of 8x8). This is a read-only validation, so result types are moot. + let offsets_primitive = + offsets_primitive.reinterpret_cast(offsets_primitive.ptype().to_unsigned()); + let sizes_primitive = + sizes_primitive.reinterpret_cast(sizes_primitive.ptype().to_unsigned()); // Validate the `offsets` and `sizes` arrays. - match_each_integer_ptype!(offset_ptype, |O| { - match_each_integer_ptype!(size_ptype, |S| { + match_each_unsigned_integer_ptype!(offsets_primitive.ptype(), |O| { + match_each_unsigned_integer_ptype!(sizes_primitive.ptype(), |S| { let offsets_slice = offsets_primitive.as_slice::(); let sizes_slice = sizes_primitive.as_slice::(); @@ -445,8 +449,13 @@ pub trait ListViewArrayExt: TypedArrayRef { let mut buf = BitBufferMut::new_unset(len); - match_each_integer_ptype!(offsets_primitive.ptype(), |O| { - match_each_integer_ptype!(sizes_primitive.ptype(), |S| { + // Offsets/sizes are non-negative; reinterpret to unsigned (4x4 instead of 8x8). + let offsets_primitive = + offsets_primitive.reinterpret_cast(offsets_primitive.ptype().to_unsigned()); + let sizes_primitive = + sizes_primitive.reinterpret_cast(sizes_primitive.ptype().to_unsigned()); + match_each_unsigned_integer_ptype!(offsets_primitive.ptype(), |O| { + match_each_unsigned_integer_ptype!(sizes_primitive.ptype(), |S| { fill_referenced_mask::( &mut buf, offsets_primitive.as_slice::(), @@ -721,11 +730,16 @@ fn validate_zctl( let sizes_dtype = sizes_primitive.dtype(); let len = offsets_primitive.len(); + // Offsets/sizes are non-negative; reinterpret to unsigned (4x4 instead of 8x8). + let offsets_unsigned = + offsets_primitive.reinterpret_cast(offsets_dtype.as_ptype().to_unsigned()); + let sizes_unsigned = sizes_primitive.reinterpret_cast(sizes_dtype.as_ptype().to_unsigned()); + // Check that offset + size values are monotonic (no overlaps) - match_each_integer_ptype!(offsets_dtype.as_ptype(), |O| { - match_each_integer_ptype!(sizes_dtype.as_ptype(), |S| { - let offsets_slice = offsets_primitive.as_slice::(); - let sizes_slice = sizes_primitive.as_slice::(); + match_each_unsigned_integer_ptype!(offsets_unsigned.ptype(), |O| { + match_each_unsigned_integer_ptype!(sizes_unsigned.ptype(), |S| { + let offsets_slice = offsets_unsigned.as_slice::(); + let sizes_slice = sizes_unsigned.as_slice::(); validate_monotonic_ends(offsets_slice, sizes_slice, len)?; }) @@ -756,9 +770,9 @@ fn validate_zctl( } } - match_each_integer_ptype!(offsets_primitive.ptype(), |O| { - match_each_integer_ptype!(sizes_primitive.ptype(), |S| { - count_references::(&mut element_references, offsets_primitive, sizes_primitive); + match_each_unsigned_integer_ptype!(offsets_unsigned.ptype(), |O| { + match_each_unsigned_integer_ptype!(sizes_unsigned.ptype(), |S| { + count_references::(&mut element_references, offsets_unsigned, sizes_unsigned); }) }); diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 0082cad427a..21dbcd70cee 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -14,7 +14,9 @@ use crate::VortexSessionExecute; use crate::aggregate_fn::fns::min_max::min_max; use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; +use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::primitive::PrimitiveArrayExt; use crate::builders::builder_with_capacity; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; @@ -22,8 +24,25 @@ use crate::dtype::IntegerPType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::match_each_integer_ptype; +use crate::match_each_unsigned_integer_ptype; use crate::scalar::Scalar; use crate::scalar_fn::fns::operators::Operator; +use crate::validity::Validity; + +/// The widened offset type used for a rebuilt, zero-copy-to-list array. +/// +/// Offsets are widened to at least 32 bits (Arrow only permits 32/64-bit list offsets) while +/// preserving signedness, since a signed result keeps the array zero-copyable to Arrow's +/// `ListArray`. +fn rebuilt_offset_ptype(offsets_ptype: PType) -> PType { + match offsets_ptype { + PType::U8 | PType::U16 | PType::U32 => PType::U32, + PType::U64 => PType::U64, + PType::I8 | PType::I16 | PType::I32 => PType::I32, + PType::I64 => PType::I64, + _ => unreachable!("invalid offsets PType"), + } +} /// Density threshold to decide whether to rebuild a sparse `ListViewArray`. /// @@ -106,16 +125,12 @@ impl ListViewArray { // `ListArray`), we rebuild the offsets as 32-bit or 64-bit integer types. // TODO(connor)[ListView]: This is true for `sizes` as well, we should do this conversion // for sizes as well. - match_each_integer_ptype!(sizes_ptype, |S| { - match offsets_ptype { + match_each_unsigned_integer_ptype!(sizes_ptype.to_unsigned(), |S| { + match offsets_ptype.to_unsigned() { PType::U8 => self.naive_rebuild::(), PType::U16 => self.naive_rebuild::(), PType::U32 => self.naive_rebuild::(), PType::U64 => self.naive_rebuild::(), - PType::I8 => self.naive_rebuild::(), - PType::I16 => self.naive_rebuild::(), - PType::I32 => self.naive_rebuild::(), - PType::I64 => self.naive_rebuild::(), _ => unreachable!("invalid offsets PType"), } }) @@ -129,6 +144,8 @@ impl ListViewArray { ) -> VortexResult { #[expect(deprecated)] let sizes_canonical = self.sizes().to_primitive(); + let sizes_canonical = + sizes_canonical.reinterpret_cast(sizes_canonical.ptype().to_unsigned()); let total: u64 = sizes_canonical .as_slice::() .iter() @@ -161,11 +178,18 @@ impl ListViewArray { fn rebuild_with_take( &self, ) -> VortexResult { + let new_offset_ptype = rebuilt_offset_ptype(self.offsets().dtype().as_ptype()); + let size_ptype = self.sizes().dtype().as_ptype(); + #[expect(deprecated)] let offsets_canonical = self.offsets().to_primitive(); + let offsets_canonical = + offsets_canonical.reinterpret_cast(offsets_canonical.ptype().to_unsigned()); let offsets_slice = offsets_canonical.as_slice::(); #[expect(deprecated)] let sizes_canonical = self.sizes().to_primitive(); + let sizes_canonical = + sizes_canonical.reinterpret_cast(sizes_canonical.ptype().to_unsigned()); let sizes_slice = sizes_canonical.as_slice::(); let len = offsets_slice.len(); @@ -194,8 +218,13 @@ impl ListViewArray { } let elements = self.elements().take(take_indices.into_array())?; - let offsets = new_offsets.into_array(); - let sizes = new_sizes.into_array(); + // Built unsigned; reinterpret back to the signed-preserving result types. + let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) + .reinterpret_cast(new_offset_ptype) + .into_array(); + let sizes = PrimitiveArray::new(new_sizes.freeze(), Validity::NonNullable) + .reinterpret_cast(size_ptype) + .into_array(); // SAFETY: same invariants as `rebuild_list_by_list` — offsets are sequential and // non-overlapping, all (offset, size) pairs reference valid elements, and the validity @@ -216,11 +245,18 @@ impl ListViewArray { .as_list_element_opt() .vortex_expect("somehow had a canonical list that was not a list"); + let new_offset_ptype = rebuilt_offset_ptype(self.offsets().dtype().as_ptype()); + let size_ptype = self.sizes().dtype().as_ptype(); + #[expect(deprecated)] let offsets_canonical = self.offsets().to_primitive(); + let offsets_canonical = + offsets_canonical.reinterpret_cast(offsets_canonical.ptype().to_unsigned()); let offsets_slice = offsets_canonical.as_slice::(); #[expect(deprecated)] let sizes_canonical = self.sizes().to_primitive(); + let sizes_canonical = + sizes_canonical.reinterpret_cast(sizes_canonical.ptype().to_unsigned()); let sizes_slice = sizes_canonical.as_slice::(); let len = offsets_slice.len(); @@ -268,8 +304,13 @@ impl ListViewArray { n_elements += num_traits::cast(size).vortex_expect("Cast failed"); } - let offsets = new_offsets.into_array(); - let sizes = new_sizes.into_array(); + // Built unsigned; reinterpret back to the signed-preserving result types. + let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) + .reinterpret_cast(new_offset_ptype) + .into_array(); + let sizes = PrimitiveArray::new(new_sizes.freeze(), Validity::NonNullable) + .reinterpret_cast(size_ptype) + .into_array(); let elements = new_elements_builder.finish(); debug_assert_eq!( @@ -625,6 +666,36 @@ mod tests { Ok(()) } + /// Rebuild with signed offsets/sizes: exercises the unsigned-reinterpret read path and asserts + /// the result offset/size dtypes preserve signedness (widened to >=32-bit for offsets). + #[test] + fn test_rebuild_preserves_signed_offset_and_size_types() -> VortexResult<()> { + use crate::dtype::PType; + + // Overlapping lists force an actual rebuild rather than the zero-copy fast path. + let elements = PrimitiveArray::from_iter(vec![1i32, 2, 3]).into_array(); + let offsets = PrimitiveArray::from_iter(vec![0i32, 1]).into_array(); + let sizes = PrimitiveArray::from_iter(vec![3i16, 2]).into_array(); + + let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable); + let rebuilt = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList)?; + + // Values: [1,2,3] and [2,3]. + assert_arrays_eq!( + rebuilt.list_elements_at(0)?, + PrimitiveArray::from_iter([1i32, 2, 3]) + ); + assert_arrays_eq!( + rebuilt.list_elements_at(1)?, + PrimitiveArray::from_iter([2i32, 3]) + ); + + // Signed input -> signed result (offsets widened to i32, sizes kept i16). + assert_eq!(rebuilt.offsets().dtype().as_ptype(), PType::I32); + assert_eq!(rebuilt.sizes().dtype().as_ptype(), PType::I16); + Ok(()) + } + // ── should_use_take heuristic tests ──────────────────────────────────── #[test] diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 7c799bfa54b..cb4b4031cd4 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -16,13 +16,27 @@ use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; use crate::arrays::VarBinArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::primitive::PrimitiveArrayExt; use crate::arrays::varbin::VarBinArrayExt; use crate::dtype::DType; use crate::dtype::IntegerPType; +use crate::dtype::PType; use crate::executor::ExecutionCtx; -use crate::match_each_integer_ptype; +use crate::match_each_unsigned_integer_ptype; use crate::validity::Validity; +/// The widened offset type used for a taken `VarBinArray`: offsets are widened to at least 32 bits +/// (to avoid overflow) while preserving signedness, so a signed result stays Arrow-compatible. +fn taken_offset_ptype(offsets_ptype: PType) -> PType { + match offsets_ptype { + PType::U8 | PType::U16 | PType::U32 => PType::U32, + PType::U64 => PType::U64, + PType::I8 | PType::I16 | PType::I32 => PType::I32, + PType::I64 => PType::I64, + _ => unreachable!("invalid PType for offsets"), + } +} + impl TakeExecute for VarBin { fn take( array: ArrayView<'_, VarBin>, @@ -45,9 +59,15 @@ impl TakeExecute for VarBin { .validity()? .execute_mask(indices.as_ref().len(), ctx)?; - let array = match_each_integer_ptype!(indices.ptype(), |I| { - // On take, offsets get widened to either 32- or 64-bit based on the original type, - // to avoid overflow issues. + // Offsets and indices are non-negative; read them through their unsigned reinterpretations + // so we only monomorphize over the 4 unsigned widths each (4x4 instead of 8x8). On take, + // offsets get widened to either 32- or 64-bit (to avoid overflow); the built output offsets + // are reinterpreted back to `out_offset_ptype` to preserve the result's offset signedness. + let out_offset_ptype = taken_offset_ptype(offsets.ptype()); + let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); + let indices = indices.reinterpret_cast(indices.ptype().to_unsigned()); + + let array = match_each_unsigned_integer_ptype!(indices.ptype(), |I| { match offsets.ptype() { PType::U8 => take::( dtype, @@ -56,6 +76,7 @@ impl TakeExecute for VarBin { indices.as_slice::(), array_validity, indices_validity, + out_offset_ptype, ), PType::U16 => take::( dtype, @@ -64,6 +85,7 @@ impl TakeExecute for VarBin { indices.as_slice::(), array_validity, indices_validity, + out_offset_ptype, ), PType::U32 => take::( dtype, @@ -72,6 +94,7 @@ impl TakeExecute for VarBin { indices.as_slice::(), array_validity, indices_validity, + out_offset_ptype, ), PType::U64 => take::( dtype, @@ -80,38 +103,7 @@ impl TakeExecute for VarBin { indices.as_slice::(), array_validity, indices_validity, - ), - PType::I8 => take::( - dtype, - offsets.as_slice::(), - data.as_slice(), - indices.as_slice::(), - array_validity, - indices_validity, - ), - PType::I16 => take::( - dtype, - offsets.as_slice::(), - data.as_slice(), - indices.as_slice::(), - array_validity, - indices_validity, - ), - PType::I32 => take::( - dtype, - offsets.as_slice::(), - data.as_slice(), - indices.as_slice::(), - array_validity, - indices_validity, - ), - PType::I64 => take::( - dtype, - offsets.as_slice::(), - data.as_slice(), - indices.as_slice::(), - array_validity, - indices_validity, + out_offset_ptype, ), _ => unreachable!("invalid PType for offsets"), } @@ -128,6 +120,7 @@ fn take( indices: &[Index], validity_mask: Mask, indices_validity_mask: Mask, + out_offset_ptype: PType, ) -> VortexResult { if !validity_mask.all_true() || !indices_validity_mask.all_true() { return Ok(take_nullable::( @@ -137,6 +130,7 @@ fn take( indices, validity_mask, indices_validity_mask, + out_offset_ptype, )); } @@ -172,11 +166,16 @@ fn take( let array_validity = Validity::from(dtype.nullability()); + // Built unsigned; reinterpret back to the signedness-preserving result offset type. + let new_offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) + .reinterpret_cast(out_offset_ptype) + .into_array(); + // Safety: // All variants of VarBinArray are satisfied here. unsafe { Ok(VarBinArray::new_unchecked( - PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable).into_array(), + new_offsets, new_data.freeze(), dtype, array_validity, @@ -191,6 +190,7 @@ fn take_nullable VarBinArray { let mut new_offsets = BufferMut::::with_capacity(indices.len() + 1); new_offsets.push(NewOffset::zero()); @@ -239,16 +239,14 @@ fn take_nullable VortexResult> { let take_indices_validity = take_indices.validity()?; + // Take indices are non-negative; reinterpret to unsigned so the `TakeT` dimension + // dispatches over 4 widths instead of 8. + let take_indices_unsigned = + take_indices.reinterpret_cast(take_indices.ptype().to_unsigned()); let patch_indices = self.indices.clone().execute::(ctx)?; let chunk_offsets = self .chunk_offsets() @@ -856,8 +861,8 @@ impl Patches { let (values_indices, new_indices): (BufferMut, BufferMut) = match_each_unsigned_integer_ptype!(patch_indices.ptype(), |PatchT| { let patch_indices_slice = patch_indices.as_slice::(); - match_each_integer_ptype!(take_indices.ptype(), |TakeT| { - let take_slice = take_indices.as_slice::(); + match_each_unsigned_integer_ptype!(take_indices_unsigned.ptype(), |TakeT| { + let take_slice = take_indices_unsigned.as_slice::(); if let Some(chunk_offsets) = chunk_offsets { match_each_unsigned_integer_ptype!(chunk_offsets.ptype(), |OffsetT| { @@ -930,18 +935,21 @@ impl Patches { ) -> VortexResult> { let indices = self.indices.clone().execute::(ctx)?; let new_length = take_indices.len(); + // Take indices are non-negative; reinterpret to unsigned (4 widths instead of 8). + let take_indices_unsigned = + take_indices.reinterpret_cast(take_indices.ptype().to_unsigned()); let min_index = self.min_index()?; let max_index = self.max_index()?; let Some((new_sparse_indices, value_indices)) = match_each_unsigned_integer_ptype!(indices.ptype(), |Indices| { - match_each_integer_ptype!(take_indices.ptype(), |TakeIndices| { + match_each_unsigned_integer_ptype!(take_indices_unsigned.ptype(), |TakeIndices| { let take_validity = take_indices .validity()? .execute_mask(take_indices.len(), ctx)?; let take_nullability = take_indices.validity()?.nullability(); - let take_slice = take_indices.as_slice::(); + let take_slice = take_indices_unsigned.as_slice::(); take_map::<_, TakeIndices>( indices.as_slice::(), take_slice, diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs index e91936ca6dc..e964af75a32 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -26,6 +26,7 @@ use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::primitive::PrimitiveArrayExt; use crate::arrays::scalar_fn::ScalarFnFactoryExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; @@ -39,6 +40,7 @@ use crate::expr::lit; use crate::expr::lt; use crate::expr::or; use crate::match_each_integer_ptype; +use crate::match_each_unsigned_integer_ptype; use crate::scalar::ListScalar; use crate::scalar::Scalar; use crate::scalar_fn::Arity; @@ -316,16 +318,19 @@ fn list_contains_scalar( }; } - // Get the offsets and sizes as primitive arrays. + // Get the offsets and sizes as primitive arrays. They are non-negative, so reinterpret to + // unsigned and dispatch over the 4 unsigned widths each (4x4 instead of 8x8). let offsets = list_array .offsets() .clone() .execute::(ctx)?; + let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); let sizes = list_array.sizes().clone().execute::(ctx)?; + let sizes = sizes.reinterpret_cast(sizes.ptype().to_unsigned()); // Process based on the offset and size types. - let list_matches = match_each_integer_ptype!(offsets.ptype(), |O| { - match_each_integer_ptype!(sizes.ptype(), |S| { + let list_matches = match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + match_each_unsigned_integer_ptype!(sizes.ptype(), |S| { process_matches::(matches, list_array.len(), offsets, sizes) }) });