1717
1818//! Defines kernel for length of string arrays and binary arrays
1919
20- use arrow_array:: types:: * ;
2120use arrow_array:: * ;
21+ use arrow_array:: { cast:: AsArray , types:: * } ;
2222use arrow_buffer:: Buffer ;
2323use arrow_data:: ArrayData ;
2424use arrow_schema:: { ArrowError , DataType } ;
8888 unary_offsets ! ( array, T :: DATA_TYPE , |x| x)
8989}
9090
91+ fn length_list_fixed_size ( array : & dyn Array , length : i32 ) -> ArrayRef {
92+ let array = array. as_fixed_size_list ( ) ;
93+ let length_list = array. len ( ) ;
94+ let buffer = Buffer :: from_vec ( vec ! [ length; length_list] ) ;
95+ let data = Int32Array :: new ( buffer. into ( ) , array. nulls ( ) . cloned ( ) ) ;
96+ Arc :: new ( data)
97+ }
98+
9199fn length_binary < O , T > ( array : & dyn Array ) -> ArrayRef
92100where
93101 O : OffsetSizeTrait ,
@@ -146,7 +154,7 @@ where
146154/// For list array, length is the number of elements in each list.
147155/// For string array and binary array, length is the number of bytes of each value.
148156///
149- /// * this only accepts ListArray/LargeListArray, StringArray/LargeStringArray and BinaryArray/LargeBinaryArray,
157+ /// * this only accepts ListArray/LargeListArray, StringArray/LargeStringArray, BinaryArray/LargeBinaryArray, and FixedSizeListArray ,
150158/// or DictionaryArray with above Arrays as values
151159/// * length of null is null.
152160pub fn length ( array : & dyn Array ) -> Result < ArrayRef , ArrowError > {
@@ -172,6 +180,7 @@ pub fn length(array: &dyn Array) -> Result<ArrayRef, ArrowError> {
172180 DataType :: LargeUtf8 => Ok ( length_string :: < i64 , Int64Type > ( array) ) ,
173181 DataType :: Binary => Ok ( length_binary :: < i32 , Int32Type > ( array) ) ,
174182 DataType :: LargeBinary => Ok ( length_binary :: < i64 , Int64Type > ( array) ) ,
183+ DataType :: FixedSizeList ( _, len) => Ok ( length_list_fixed_size ( array, * len) ) ,
175184 other => Err ( ArrowError :: ComputeError ( format ! (
176185 "length not supported for {other:?}"
177186 ) ) ) ,
@@ -215,6 +224,8 @@ pub fn bit_length(array: &dyn Array) -> Result<ArrayRef, ArrowError> {
215224mod tests {
216225 use super :: * ;
217226 use arrow_array:: cast:: AsArray ;
227+ use arrow_buffer:: NullBuffer ;
228+ use arrow_schema:: Field ;
218229
219230 fn double_vec < T : Clone > ( v : Vec < T > ) -> Vec < T > {
220231 [ & v[ ..] , & v[ ..] ] . concat ( )
@@ -696,4 +707,34 @@ mod tests {
696707 assert_eq ! ( expected[ i] , actual[ i] , ) ;
697708 }
698709 }
710+
711+ #[ test]
712+ fn test_fixed_size_list_length ( ) {
713+ // Construct a value array
714+ let value_data = ArrayData :: builder ( DataType :: Int32 )
715+ . len ( 9 )
716+ . add_buffer ( Buffer :: from_slice_ref ( [ 0 , 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 ] ) )
717+ . build ( )
718+ . unwrap ( ) ;
719+ let list_data_type = DataType :: FixedSizeList (
720+ Arc :: new ( Field :: new ( "item" , DataType :: Int32 , false ) ) ,
721+ 3 ,
722+ ) ;
723+ let nulls = NullBuffer :: from ( vec ! [ true , false , true ] ) ;
724+ let list_data = ArrayData :: builder ( list_data_type)
725+ . len ( 3 )
726+ . add_child_data ( value_data)
727+ . nulls ( Some ( nulls) )
728+ . build ( )
729+ . unwrap ( ) ;
730+ let list_array = FixedSizeListArray :: from ( list_data) ;
731+
732+ let lengths = length ( & list_array) . unwrap ( ) ;
733+ let lengths = lengths. as_any ( ) . downcast_ref :: < Int32Array > ( ) . unwrap ( ) ;
734+
735+ assert_eq ! ( lengths. len( ) , 3 ) ;
736+ assert_eq ! ( lengths. value( 0 ) , 3 ) ;
737+ assert ! ( lengths. is_null( 1 ) ) ;
738+ assert_eq ! ( lengths. value( 2 ) , 3 ) ;
739+ }
699740}
0 commit comments