@@ -11,6 +11,7 @@ use vortex_error::vortex_err;
1111use crate :: ArrayRef ;
1212use crate :: IntoArray ;
1313use crate :: array:: ArrayView ;
14+ use crate :: arrays:: ConstantArray ;
1415use crate :: arrays:: List ;
1516use crate :: arrays:: ListArray ;
1617use crate :: arrays:: PiecewiseSequence ;
@@ -19,7 +20,7 @@ use crate::arrays::Primitive;
1920use crate :: arrays:: PrimitiveArray ;
2021use crate :: arrays:: dict:: TakeExecute ;
2122use crate :: arrays:: list:: ListArrayExt ;
22- use crate :: arrays:: piecewise_sequence:: execute_index_arrays ;
23+ use crate :: arrays:: piecewise_sequence:: execute_unit_multiplier_index_arrays ;
2324use crate :: arrays:: piecewise_sequence:: validate_index_ranges;
2425use crate :: arrays:: primitive:: PrimitiveArrayExt ;
2526use crate :: builders:: ArrayBuilder ;
@@ -156,7 +157,9 @@ fn take_piecewise_sequence(
156157 return Ok ( None ) ;
157158 }
158159
159- let ( starts, lengths) = execute_index_arrays ( indices, ctx) ?;
160+ let Some ( ( starts, lengths) ) = execute_unit_multiplier_index_arrays ( indices, ctx) ? else {
161+ return Ok ( None ) ;
162+ } ;
160163 let offsets = array. offsets ( ) . clone ( ) . execute :: < PrimitiveArray > ( ctx) ?;
161164 let offsets = offsets. reinterpret_cast ( offsets. ptype ( ) . to_unsigned ( ) ) ;
162165
@@ -305,12 +308,14 @@ where
305308 debug_assert_eq ! ( output_elements, total_elements) ;
306309
307310 let offsets = PrimitiveArray :: new ( new_offsets. freeze ( ) , Validity :: NonNullable ) . into_array ( ) ;
311+ let multipliers = ConstantArray :: new ( 1u64 , element_starts. len ( ) ) . into_array ( ) ;
308312 // SAFETY: element ranges are derived from validated source list offsets, and total_elements is
309- // the sum of the gathered element range lengths.
313+ // the sum of the gathered element range lengths. Multiplier 1 preserves contiguous ranges.
310314 let element_indices = unsafe {
311315 PiecewiseSequenceArray :: new_unchecked (
312316 element_starts. into_array ( ) ,
313317 element_lengths. into_array ( ) ,
318+ multipliers,
314319 total_elements,
315320 )
316321 } ;
@@ -416,6 +421,7 @@ mod test {
416421 use crate :: VortexSessionExecute ;
417422 use crate :: array_session;
418423 use crate :: arrays:: BoolArray ;
424+ use crate :: arrays:: ConstantArray ;
419425 use crate :: arrays:: ListArray ;
420426 use crate :: arrays:: ListViewArray ;
421427 use crate :: arrays:: PiecewiseSequenceArray ;
@@ -618,6 +624,7 @@ mod test {
618624 let idx = PiecewiseSequenceArray :: try_new (
619625 buffer ! [ 1u64 , 0 ] . into_array ( ) ,
620626 buffer ! [ 2u64 , 1 ] . into_array ( ) ,
627+ ConstantArray :: new ( 1u64 , 2 ) . into_array ( ) ,
621628 3 ,
622629 )
623630 . unwrap ( )
0 commit comments