Skip to content

Commit 730941f

Browse files
authored
feat: Support FixedSizedListArray for length kernel (#4520)
* feat: Support FixedSizedListArray for length kernel * fix clippy * update comment * avoid unsafe * reduce useless trait * remove T
1 parent 47f288c commit 730941f

File tree

1 file changed

+43
-2
lines changed

1 file changed

+43
-2
lines changed

arrow-string/src/length.rs

Lines changed: 43 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,8 @@
1717

1818
//! Defines kernel for length of string arrays and binary arrays
1919
20-
use arrow_array::types::*;
2120
use arrow_array::*;
21+
use arrow_array::{cast::AsArray, types::*};
2222
use arrow_buffer::Buffer;
2323
use arrow_data::ArrayData;
2424
use arrow_schema::{ArrowError, DataType};
@@ -88,6 +88,14 @@ where
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+
9199
fn length_binary<O, T>(array: &dyn Array) -> ArrayRef
92100
where
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.
152160
pub 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> {
215224
mod 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

Comments
 (0)