@@ -10,7 +10,7 @@ use vortex_buffer::Buffer;
1010use vortex_dtype:: {
1111 DType , DecimalDType , NativeDecimalType , NativePType , Nullability , match_each_native_ptype,
1212} ;
13- use vortex_error:: VortexResult ;
13+ use vortex_error:: { VortexExpect , VortexResult } ;
1414use vortex_scalar:: match_each_decimal_value_type;
1515
1616pub fn take_canonical_array_non_nullable_indices (
@@ -31,11 +31,7 @@ pub fn take_canonical_array(
3131 array : & dyn Array ,
3232 indices : & [ Option < usize > ] ,
3333) -> VortexResult < ArrayRef > {
34- let nullable = if indices. contains ( & None ) {
35- Nullability :: Nullable
36- } else {
37- Nullability :: NonNullable
38- } ;
34+ let nullable: Nullability = indices. contains ( & None ) . into ( ) ;
3935
4036 let validity = if array. dtype ( ) . is_nullable ( ) || nullable == Nullability :: Nullable {
4137 let validity_idx = array. validity_mask ( ) . to_bit_buffer ( ) ;
@@ -116,10 +112,18 @@ pub fn take_canonical_array(
116112 . map ( |a| a. into_array ( ) )
117113 }
118114 DType :: List ( ..) | DType :: FixedSizeList ( ..) => {
119- let mut builder = builder_with_capacity ( array. dtype ( ) , indices_slice_non_opt. len ( ) ) ;
115+ let mut builder = builder_with_capacity (
116+ & array. dtype ( ) . union_nullability ( nullable) ,
117+ indices_slice_non_opt. len ( ) ,
118+ ) ;
120119 for idx in indices {
121120 if let Some ( idx) = idx {
122- builder. append_scalar ( & array. scalar_at ( * idx) ) ?;
121+ builder. append_scalar (
122+ & array
123+ . scalar_at ( * idx)
124+ . cast ( & array. dtype ( ) . union_nullability ( nullable) )
125+ . vortex_expect ( "cannot cast scalar nullability" ) ,
126+ ) ?;
123127 } else {
124128 builder. append_null ( )
125129 }
0 commit comments