From 080216ebfd475f0438dfbcc1604a75b8172d1e71 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:24:09 -0400 Subject: [PATCH 1/3] Specialize VarBin PiecewiseSequence takes Signed-off-by: Daniel King --- .../src/arrays/varbin/compute/take.rs | 392 ++++++++++++++++++ .../src/arrays/varbinview/compute/take.rs | 119 ++++++ 2 files changed, 511 insertions(+) diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 6a6454d0603..b956359ef0b 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -1,26 +1,33 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use itertools::Itertools as _; use vortex_buffer::BitBufferMut; use vortex_buffer::BufferMut; use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; use crate::array::ArrayView; +use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; use crate::arrays::VarBinArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::ConstantOrArray; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; use crate::arrays::varbin::VarBinArrayExt; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::PType; +use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_unsigned_integer_ptype; use crate::validity::Validity; @@ -43,6 +50,12 @@ impl TakeExecute for VarBin { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? + { + return Ok(Some(taken)); + } + // TODO(joe): Be lazy with execute let offsets = array.offsets().clone().execute::(ctx)?; let data = array.bytes(); @@ -113,6 +126,54 @@ impl TakeExecute for VarBin { } } +fn take_contiguous_ranges( + array: ArrayView<'_, VarBin>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { + return Ok(None); + }; + let offsets = array.offsets().clone().execute::(ctx)?; + let out_offset_ptype = taken_offset_ptype(offsets.ptype()); + let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); + let bytes = array.bytes(); + let data = bytes.as_slice(); + let dtype = array.dtype().clone(); + let output_len = indices_ref.len(); + + let result = match &lengths { + ConstantOrArray::Constant(length) => gather_slices_constant_dispatch( + &starts, + *length, + &offsets, + data, + output_len, + out_offset_ptype, + )?, + ConstantOrArray::Array(lengths) => gather_slices_dispatch( + &starts, + lengths, + &offsets, + data, + output_len, + out_offset_ptype, + )?, + }; + + let validity = array.validity()?.take(indices_ref)?; + + // SAFETY: output offsets are built from valid input offsets, start at zero, are monotonically + // non-decreasing, and the copied data buffer has exactly the referenced byte length. + unsafe { + Ok(Some( + VarBinArray::new_unchecked(result.offsets, result.data.freeze(), dtype, validity) + .into_array(), + )) + } +} + fn take( dtype: DType, offsets: &[Offset], @@ -183,6 +244,337 @@ fn take( } } +struct GatheredPiecewiseVarBin { + offsets: ArrayRef, + data: ByteBufferMut, +} + +fn gather_slices_constant_dispatch( + starts: &PrimitiveArray, + length: usize, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + gather_slices_constant_start_dispatch::( + starts, + length, + offsets, + data, + output_len, + out_offset_ptype, + ) + }) +} + +fn gather_slices_constant_start_dispatch( + starts: &PrimitiveArray, + length: usize, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, +{ + match offsets.ptype() { + PType::U8 => gather_slices_constant_length::( + offsets.as_slice::(), + data, + starts.as_slice::(), + length, + output_len, + out_offset_ptype, + ), + PType::U16 => gather_slices_constant_length::( + offsets.as_slice::(), + data, + starts.as_slice::(), + length, + output_len, + out_offset_ptype, + ), + PType::U32 => gather_slices_constant_length::( + offsets.as_slice::(), + data, + starts.as_slice::(), + length, + output_len, + out_offset_ptype, + ), + PType::U64 => gather_slices_constant_length::( + offsets.as_slice::(), + data, + starts.as_slice::(), + length, + output_len, + out_offset_ptype, + ), + _ => unreachable!("offsets were reinterpreted to an unsigned integer ptype"), + } +} + +fn gather_slices_dispatch( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + gather_slices_start_dispatch::( + starts, + lengths, + offsets, + data, + output_len, + out_offset_ptype, + ) + }) +} + +fn gather_slices_start_dispatch( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, +{ + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + gather_slices_start_length_dispatch::( + starts, + lengths, + offsets, + data, + output_len, + out_offset_ptype, + ) + }) +} + +fn gather_slices_start_length_dispatch( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, +{ + match offsets.ptype() { + PType::U8 => gather_slices::( + offsets.as_slice::(), + data, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + PType::U16 => gather_slices::( + offsets.as_slice::(), + data, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + PType::U32 => gather_slices::( + offsets.as_slice::(), + data, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + PType::U64 => gather_slices::( + offsets.as_slice::(), + data, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + _ => unreachable!("offsets were reinterpreted to an unsigned integer ptype"), + } +} + +fn gather_slices_constant_length( + offsets: &[Offset], + data: &[u8], + starts: &[S], + length: usize, + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, + Offset: IntegerPType, + NewOffset: IntegerPType, +{ + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); + + let mut new_offsets = BufferMut::::with_capacity(output_len + 1); + new_offsets.push(NewOffset::zero()); + let mut output_bytes = 0usize; + + for &start in starts { + let start = start.as_(); + if length == 0 { + continue; + } + + let offset_range = &offsets[start..][..=length]; + let byte_start = offset_range[0].as_(); + let byte_end = offset_range[length].as_(); + vortex_ensure!( + byte_start <= byte_end && byte_end <= data.len(), + "VarBin offsets range {byte_start}..{byte_end} exceeds data length {}", + data.len() + ); + + for &offset in &offset_range[1..] { + let offset = offset.as_(); + let relative = offset.checked_sub(byte_start).ok_or_else(|| { + vortex_err!("VarBin offsets are not monotonic at offset {offset}") + })?; + let output_offset = output_bytes.checked_add(relative).ok_or_else(|| { + vortex_err!("PiecewiseSequence VarBin output byte length overflow") + })?; + new_offsets.push(new_offset_value::(output_offset)?); + } + + output_bytes = output_bytes + .checked_add(byte_end - byte_start) + .ok_or_else(|| vortex_err!("PiecewiseSequence VarBin output byte length overflow"))?; + } + + let mut new_data = ByteBufferMut::with_capacity(output_bytes); + for &start in starts { + let start = start.as_(); + if length == 0 { + continue; + } + + let offset_range = &offsets[start..][..=length]; + let byte_start = offset_range[0].as_(); + let byte_end = offset_range[length].as_(); + new_data.extend_from_slice(&data[byte_start..][..byte_end - byte_start]); + } + + let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) + .reinterpret_cast(out_offset_ptype) + .into_array(); + Ok(GatheredPiecewiseVarBin { + offsets, + data: new_data, + }) +} + +fn gather_slices( + offsets: &[Offset], + data: &[u8], + starts: &[S], + lengths: &[L], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, + Offset: IntegerPType, + NewOffset: IntegerPType, +{ + let mut new_offsets = BufferMut::::with_capacity(output_len + 1); + new_offsets.push(NewOffset::zero()); + let mut output_bytes = 0usize; + + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + if length == 0 { + continue; + } + + let offset_range = &offsets[start..][..=length]; + let byte_start = offset_range[0].as_(); + let byte_end = offset_range[length].as_(); + vortex_ensure!( + byte_start <= byte_end && byte_end <= data.len(), + "VarBin offsets range {byte_start}..{byte_end} exceeds data length {}", + data.len() + ); + + for &offset in &offset_range[1..] { + let offset = offset.as_(); + let relative = offset.checked_sub(byte_start).ok_or_else(|| { + vortex_err!("VarBin offsets are not monotonic at offset {offset}") + })?; + let output_offset = output_bytes.checked_add(relative).ok_or_else(|| { + vortex_err!("PiecewiseSequence VarBin output byte length overflow") + })?; + new_offsets.push(new_offset_value::(output_offset)?); + } + + output_bytes = output_bytes + .checked_add(byte_end - byte_start) + .ok_or_else(|| vortex_err!("PiecewiseSequence VarBin output byte length overflow"))?; + } + vortex_ensure!( + new_offsets.len() == output_len + 1, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + new_offsets.len() - 1 + ); + + let mut new_data = ByteBufferMut::with_capacity(output_bytes); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + if length == 0 { + continue; + } + + let offset_range = &offsets[start..][..=length]; + let byte_start = offset_range[0].as_(); + let byte_end = offset_range[length].as_(); + new_data.extend_from_slice(&data[byte_start..][..byte_end - byte_start]); + } + + let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) + .reinterpret_cast(out_offset_ptype) + .into_array(); + Ok(GatheredPiecewiseVarBin { + offsets, + data: new_data, + }) +} + +fn new_offset_value(value: usize) -> VortexResult { + T::from(value).ok_or_else(|| { + vortex_err!( + "PiecewiseSequence VarBin offset value {value} does not fit in {}", + T::PTYPE + ) + }) +} + fn take_nullable( dtype: DType, offsets: &[Offset], diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index ce10abda3d0..40e05e08737 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -4,23 +4,32 @@ use std::iter; use std::sync::Arc; +use itertools::Itertools as _; use num_traits::AsPrimitive; use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_mask::AllOr; use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; use crate::array::ArrayView; +use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::ConstantOrArray; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::varbinview::BinaryView; use crate::buffer::BufferHandle; +use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_integer_ptype; +use crate::match_each_unsigned_integer_ptype; impl TakeExecute for VarBinView { /// Take involves creating a new array that references the old array, just with the given set of views. @@ -29,6 +38,12 @@ impl TakeExecute for VarBinView { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? + { + return Ok(Some(taken)); + } + let validity = array.validity()?.take(indices)?; let indices = indices.clone().execute::(ctx)?; @@ -57,6 +72,58 @@ impl TakeExecute for VarBinView { } } +fn take_contiguous_ranges( + array: ArrayView<'_, VarBinView>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { + return Ok(None); + }; + let source = array.views(); + let output_len = indices_ref.len(); + let views = match &lengths { + ConstantOrArray::Constant(length) => { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + gather_view_slices_constant_length( + source, + starts.as_slice::(), + *length, + output_len, + )? + }) + } + ConstantOrArray::Array(lengths) => { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + gather_view_slices( + source, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + )? + }) + }) + } + }; + let validity = array.validity()?.take(indices_ref)?; + + // SAFETY: ranges were validated against the source views, and copied views still reference the + // same backing data buffers. + unsafe { + Ok(Some( + VarBinViewArray::new_handle_unchecked( + BufferHandle::new_host(views.into_byte_buffer()), + Arc::clone(array.data_buffers()), + array.dtype().clone(), + validity, + ) + .into_array(), + )) + } +} + fn take_views>( views_ref: &[BinaryView], indices: &[I], @@ -85,6 +152,58 @@ fn take_views>( } } +fn gather_view_slices_constant_length( + source: &[BinaryView], + starts: &[S], + length: usize, + output_len: usize, +) -> VortexResult> +where + S: UnsignedPType, +{ + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); + + let mut views = BufferMut::::with_capacity(output_len); + for &start in starts { + let start = start.as_(); + views.extend_from_slice(&source[start..][..length]); + } + + Ok(views.freeze()) +} + +fn gather_view_slices( + source: &[BinaryView], + starts: &[S], + lengths: &[L], + output_len: usize, +) -> VortexResult> +where + S: UnsignedPType, + L: UnsignedPType, +{ + let mut views = BufferMut::::with_capacity(output_len); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + views.extend_from_slice(&source[start..][..length]); + } + + vortex_ensure!( + views.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + views.len() + ); + Ok(views.freeze()) +} + #[cfg(test)] mod tests { use rstest::rstest; From fe7dd28c6930de60fe20897a1da0e151ad8633bb Mon Sep 17 00:00:00 2001 From: Daniel King Date: Mon, 20 Jul 2026 12:34:22 -0400 Subject: [PATCH 2/3] Use Columnar lengths in VarBin takes Signed-off-by: Daniel King --- .../src/arrays/varbin/compute/take.rs | 36 ++++++++++++------- .../src/arrays/varbinview/compute/take.rs | 20 ++++++++--- 2 files changed, 39 insertions(+), 17 deletions(-) diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index b956359ef0b..41e6dd2a936 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -7,12 +7,15 @@ use vortex_buffer::BufferMut; use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_mask::Mask; use crate::ArrayRef; +use crate::Canonical; +use crate::Columnar; use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::PiecewiseSequence; @@ -20,7 +23,7 @@ use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; use crate::arrays::VarBinArray; use crate::arrays::dict::TakeExecute; -use crate::arrays::piecewise_sequence::ConstantOrArray; +use crate::arrays::piecewise_sequence::constant_unsigned_usize; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; use crate::arrays::varbin::VarBinArrayExt; @@ -143,23 +146,32 @@ fn take_contiguous_ranges( let dtype = array.dtype().clone(); let output_len = indices_ref.len(); - let result = match &lengths { - ConstantOrArray::Constant(length) => gather_slices_constant_dispatch( - &starts, - *length, - &offsets, - data, - output_len, - out_offset_ptype, - )?, - ConstantOrArray::Array(lengths) => gather_slices_dispatch( + let result = match lengths { + Columnar::Constant(lengths) => { + let length = constant_unsigned_usize(&lengths)?; + gather_slices_constant_dispatch( + &starts, + length, + &offsets, + data, + output_len, + out_offset_ptype, + )? + } + Columnar::Canonical(Canonical::Primitive(lengths)) => gather_slices_dispatch( &starts, - lengths, + &lengths, &offsets, data, output_len, out_offset_ptype, )?, + Columnar::Canonical(lengths) => { + vortex_bail!( + "PiecewiseSequenceArray lengths must be primitive or constant, got {}", + lengths.dtype() + ) + } }; let validity = array.validity()?.take(indices_ref)?; diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index 40e05e08737..9dbde23df1e 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -9,12 +9,15 @@ use num_traits::AsPrimitive; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_mask::AllOr; use vortex_mask::Mask; use crate::ArrayRef; +use crate::Canonical; +use crate::Columnar; use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::PiecewiseSequence; @@ -22,7 +25,7 @@ use crate::arrays::PrimitiveArray; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::dict::TakeExecute; -use crate::arrays::piecewise_sequence::ConstantOrArray; +use crate::arrays::piecewise_sequence::constant_unsigned_usize; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::varbinview::BinaryView; use crate::buffer::BufferHandle; @@ -83,18 +86,19 @@ fn take_contiguous_ranges( }; let source = array.views(); let output_len = indices_ref.len(); - let views = match &lengths { - ConstantOrArray::Constant(length) => { + let views = match lengths { + Columnar::Constant(lengths) => { + let length = constant_unsigned_usize(&lengths)?; match_each_unsigned_integer_ptype!(starts.ptype(), |S| { gather_view_slices_constant_length( source, starts.as_slice::(), - *length, + length, output_len, )? }) } - ConstantOrArray::Array(lengths) => { + Columnar::Canonical(Canonical::Primitive(lengths)) => { match_each_unsigned_integer_ptype!(starts.ptype(), |S| { match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { gather_view_slices( @@ -106,6 +110,12 @@ fn take_contiguous_ranges( }) }) } + Columnar::Canonical(lengths) => { + vortex_bail!( + "PiecewiseSequenceArray lengths must be primitive or constant, got {}", + lengths.dtype() + ) + } }; let validity = array.validity()?.take(indices_ref)?; From 5d31540246f28d18d13cac102733066e32473d5e Mon Sep 17 00:00:00 2001 From: Daniel King Date: Mon, 20 Jul 2026 12:51:54 -0400 Subject: [PATCH 3/3] Rely on PiecewiseSequence validation in VarBin takes Signed-off-by: Daniel King --- .../src/arrays/varbin/compute/take.rs | 25 ++++++++----------- .../src/arrays/varbinview/compute/take.rs | 13 +++------- 2 files changed, 13 insertions(+), 25 deletions(-) diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 41e6dd2a936..f69ed43e8fa 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -7,14 +7,12 @@ use vortex_buffer::BufferMut; use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_mask::Mask; use crate::ArrayRef; -use crate::Canonical; use crate::Columnar; use crate::IntoArray; use crate::array::ArrayView; @@ -148,7 +146,7 @@ fn take_contiguous_ranges( let result = match lengths { Columnar::Constant(lengths) => { - let length = constant_unsigned_usize(&lengths)?; + let length = constant_unsigned_usize(&lengths); gather_slices_constant_dispatch( &starts, length, @@ -158,19 +156,16 @@ fn take_contiguous_ranges( out_offset_ptype, )? } - Columnar::Canonical(Canonical::Primitive(lengths)) => gather_slices_dispatch( - &starts, - &lengths, - &offsets, - data, - output_len, - out_offset_ptype, - )?, Columnar::Canonical(lengths) => { - vortex_bail!( - "PiecewiseSequenceArray lengths must be primitive or constant, got {}", - lengths.dtype() - ) + let lengths = lengths.into_primitive(); + gather_slices_dispatch( + &starts, + &lengths, + &offsets, + data, + output_len, + out_offset_ptype, + )? } }; diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index 9dbde23df1e..5af1c1b032c 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -9,14 +9,12 @@ use num_traits::AsPrimitive; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_mask::AllOr; use vortex_mask::Mask; use crate::ArrayRef; -use crate::Canonical; use crate::Columnar; use crate::IntoArray; use crate::array::ArrayView; @@ -88,7 +86,7 @@ fn take_contiguous_ranges( let output_len = indices_ref.len(); let views = match lengths { Columnar::Constant(lengths) => { - let length = constant_unsigned_usize(&lengths)?; + let length = constant_unsigned_usize(&lengths); match_each_unsigned_integer_ptype!(starts.ptype(), |S| { gather_view_slices_constant_length( source, @@ -98,7 +96,8 @@ fn take_contiguous_ranges( )? }) } - Columnar::Canonical(Canonical::Primitive(lengths)) => { + Columnar::Canonical(lengths) => { + let lengths = lengths.into_primitive(); match_each_unsigned_integer_ptype!(starts.ptype(), |S| { match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { gather_view_slices( @@ -110,12 +109,6 @@ fn take_contiguous_ranges( }) }) } - Columnar::Canonical(lengths) => { - vortex_bail!( - "PiecewiseSequenceArray lengths must be primitive or constant, got {}", - lengths.dtype() - ) - } }; let validity = array.validity()?.take(indices_ref)?;