Jefffrey commented on code in PR #10046:
URL: https://github.com/apache/arrow-rs/pull/10046#discussion_r3456978481
##########
arrow-select/src/interleave.rs:
##########
@@ -1941,4 +2020,145 @@ mod tests {
&[3]
);
}
+
+ #[test]
+ fn test_interleave_fixed_size_list() {
+ // a: [[1, 2], [3, 4], [5, 6]]
+ let field = Arc::new(Field::new("item", DataType::Int32, false));
+ let a = FixedSizeListArray::new(
+ field.clone(),
+ 2,
+ Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6])),
+ None,
+ );
+ // b: [[7, 8], [9, 10]]
+ let b = FixedSizeListArray::new(
+ field.clone(),
+ 2,
+ Arc::new(Int32Array::from(vec![7, 8, 9, 10])),
+ None,
+ );
+
+ let result = interleave(&[&a, &b], &[(0, 2), (1, 0), (0, 0), (1, 1),
(0, 1)]).unwrap();
+ let result = result.as_fixed_size_list();
+ assert_eq!(result.len(), 5);
+ assert_eq!(result.value_length(), 2);
+
+ let values = result.values().as_primitive::<Int32Type>();
+ // [[5,6], [7,8], [1,2], [9,10], [3,4]]
+ assert_eq!(values.values(), &[5, 6, 7, 8, 1, 2, 9, 10, 3, 4]);
+ }
+
+ #[test]
+ fn test_interleave_fixed_size_list_with_nulls() {
Review Comment:
do we have a test for fixedsizelist with non-primitive child?
##########
arrow-select/src/interleave.rs:
##########
@@ -513,6 +517,81 @@ fn interleave_list<O: OffsetSizeTrait>(
Ok(Arc::new(list_array))
}
+fn interleave_fixed_size_list(
+ values: &[&dyn Array],
+ indices: &[(usize, usize)],
+ field: &FieldRef,
+ size: i32,
+) -> Result<ArrayRef, ArrowError> {
+ let interleaved = Interleave::<'_, FixedSizeListArray>::new(values,
indices);
+ let capacity = indices.len() * size as usize;
+
+ macro_rules! fsl_primitive_helper {
+ ($t:ty) => {
+ interleave_list_like_primitive_child::<FixedSizeListArray, $t>(
+ &interleaved,
+ indices,
+ capacity,
+ field.data_type(),
+ )
+ };
+ }
+
+ let interleaved_values = downcast_primitive! {
+ field.data_type() => (fsl_primitive_helper),
+ _ => {
+ interleave_list_like_child(&interleaved, indices, capacity)?
+ }
+ };
+
+ let array = FixedSizeListArray::new(field.clone(), size,
interleaved_values, interleaved.nulls);
+ Ok(Arc::new(array))
+}
+
+fn interleave_map(
+ values: &[&dyn Array],
+ indices: &[(usize, usize)],
+ field: &FieldRef,
+ ordered: bool,
+) -> Result<ArrayRef, ArrowError> {
+ let interleaved = Interleave::<'_, MapArray>::new(values, indices);
+
+ let mut capacity = 0usize;
+ let mut offsets = Vec::with_capacity(indices.len() + 1);
+ offsets.push(0i32);
+ for &(array, row) in indices {
+ let o = interleaved.arrays[array].value_offsets();
+ let element_len = (o[row + 1] - o[row]) as usize;
+ capacity += element_len;
+ offsets
+ .push(i32::try_from(capacity).map_err(|_|
ArrowError::OffsetOverflowError(capacity))?);
+ }
+
+ let mut child_indices = Vec::with_capacity(capacity);
+ for &(array, row) in indices {
+ let o = interleaved.arrays[array].value_offsets();
+ let start = o[row] as usize;
+ let end = o[row + 1] as usize;
+ child_indices.extend((start..end).map(|i| (array, i)));
+ }
+
+ let entries_arrays: Vec<&dyn Array> = interleaved
+ .arrays
+ .iter()
+ .map(|m| m.entries() as &dyn Array)
+ .collect();
+ let interleaved_entries = interleave(&entries_arrays, &child_indices)?;
+
+ let offsets = OffsetBuffer::new(offsets.into());
+ let entries = interleaved_entries
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .unwrap()
+ .clone();
Review Comment:
```suggestion
let entries = interleaved_entries.as_struct().clone();
```
##########
arrow-select/src/interleave.rs:
##########
@@ -513,6 +517,81 @@ fn interleave_list<O: OffsetSizeTrait>(
Ok(Arc::new(list_array))
}
+fn interleave_fixed_size_list(
+ values: &[&dyn Array],
+ indices: &[(usize, usize)],
+ field: &FieldRef,
+ size: i32,
+) -> Result<ArrayRef, ArrowError> {
+ let interleaved = Interleave::<'_, FixedSizeListArray>::new(values,
indices);
+ let capacity = indices.len() * size as usize;
+
+ macro_rules! fsl_primitive_helper {
+ ($t:ty) => {
+ interleave_list_like_primitive_child::<FixedSizeListArray, $t>(
+ &interleaved,
+ indices,
+ capacity,
+ field.data_type(),
+ )
+ };
+ }
+
+ let interleaved_values = downcast_primitive! {
+ field.data_type() => (fsl_primitive_helper),
+ _ => {
+ interleave_list_like_child(&interleaved, indices, capacity)?
+ }
+ };
+
+ let array = FixedSizeListArray::new(field.clone(), size,
interleaved_values, interleaved.nulls);
+ Ok(Arc::new(array))
+}
+
+fn interleave_map(
+ values: &[&dyn Array],
+ indices: &[(usize, usize)],
+ field: &FieldRef,
+ ordered: bool,
+) -> Result<ArrayRef, ArrowError> {
+ let interleaved = Interleave::<'_, MapArray>::new(values, indices);
+
+ let mut capacity = 0usize;
+ let mut offsets = Vec::with_capacity(indices.len() + 1);
Review Comment:
could consider using
[`OffsetBufferBuilder`](https://docs.rs/arrow/latest/arrow/array/struct.OffsetBufferBuilder.html)
though i suppose it doesn't have a way to handle overflow without panicking
🤔
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
