1- use arrow_buffer:: ArrowNativeType ;
21use vortex_array:: accessor:: ArrayAccessor ;
32use vortex_array:: arrays:: { BoolArray , ListArray , PrimitiveArray , StructArray , VarBinViewArray } ;
43use vortex_array:: validity:: Validity ;
54use vortex_array:: variants:: { PrimitiveArrayTrait , StructArrayTrait } ;
65use vortex_array:: { Array , ArrayRef , ToCanonical } ;
7- use vortex_dtype:: { DType , NativePType , match_each_native_ptype} ;
6+ use vortex_dtype:: {
7+ DType , NativePType , Nullability , match_each_integer_ptype, match_each_native_ptype,
8+ } ;
89use vortex_error:: VortexResult ;
910
1011pub fn slice_canonical_array (
@@ -61,13 +62,14 @@ pub fn slice_canonical_array(
6162 let offsets =
6263 slice_canonical_array ( list_array. offsets ( ) , start, stop + 1 ) ?. to_primitive ( ) ?;
6364
64- let elements = slice_canonical_array (
65- list_array. elements ( ) ,
66- usize:: try_from ( offsets. get_as_cast :: < u64 > ( 0 ) ) ?,
67- usize:: try_from ( offsets. get_as_cast :: < u64 > ( offsets. len ( ) - 1 ) ) ?,
68- ) ?;
69- let offsets = match_each_native_ptype ! ( offsets. ptype( ) , |$P | {
70- shift_offsets:: <$P >( offsets)
65+ let ( start, end) = match_each_integer_ptype ! ( offsets. ptype( ) , |$P | {
66+ let offset_slice = offsets. as_slice:: <$P >( ) ;
67+ ( usize :: try_from( offset_slice[ 0 ] ) ?, usize :: try_from( offset_slice[ offsets. len( ) - 1 ] ) ?)
68+ } ) ;
69+
70+ let elements = slice_canonical_array ( list_array. elements ( ) , start, end) ?;
71+ let offsets = match_each_integer_ptype ! ( offsets. ptype( ) , |$P | {
72+ shift_offsets( offsets. as_slice:: <$P >( ) )
7173 } )
7274 . into_array ( ) ;
7375 ListArray :: try_new ( elements, offsets, validity) . map ( |a| a. into_array ( ) )
@@ -76,11 +78,16 @@ pub fn slice_canonical_array(
7678 }
7779}
7880
79- fn shift_offsets < O : NativePType + ArrowNativeType > ( offsets : PrimitiveArray ) -> PrimitiveArray {
81+ fn shift_offsets < O : NativePType > ( offsets : & [ O ] ) -> PrimitiveArray {
8082 if offsets. is_empty ( ) {
81- return offsets ;
83+ return PrimitiveArray :: empty :: < O > ( Nullability :: NonNullable ) ;
8284 }
83- let offsets: Vec < O > = offsets. as_slice ( ) . to_vec ( ) ;
8485 let start = offsets[ 0 ] ;
85- PrimitiveArray :: from_iter ( offsets. into_iter ( ) . map ( |o| o - start) . collect :: < Vec < _ > > ( ) )
86+ PrimitiveArray :: from_iter (
87+ offsets
88+ . iter ( )
89+ . copied ( )
90+ . map ( |o| o - start)
91+ . collect :: < Vec < _ > > ( ) ,
92+ )
8693}
0 commit comments