comphead commented on code in PR #6498:
URL: https://github.com/apache/datafusion-comet/pull/6498#discussion_r4277250644
##########
native/spark-expr/src/array_funcs/list_extract.rs:
##########
@@ -411,6 +525,275 @@ mod test {
use datafusion::physical_expr::expressions::Column;
use datafusion::physical_plan::ColumnarValue;
+ fn skewed_nested_values(selected_len: usize) ->
Vec<arrow::array::ArrayRef> {
+ use arrow::array::{LargeListArray, MapArray, StructArray};
+ use arrow::buffer::{NullBuffer, OffsetBuffer};
+ const ROWS: usize = 8192;
+ let offsets =
+ OffsetBuffer::<i32>::from_lengths((0..ROWS).flat_map(|_|
[selected_len, 128]));
+ let count = *offsets.last().unwrap() as usize;
+ let ints: arrow::array::ArrayRef =
Arc::new(Int32Array::from_iter_values(0..count as i32));
+ let field = Arc::new(Field::new("item", DataType::Int32, true));
+ let lists = Arc::new(ListArray::new(
+ Arc::clone(&field),
+ offsets.clone(),
+ Arc::clone(&ints),
+ None,
+ ));
+ let large_lists = Arc::new(LargeListArray::new(
+ Arc::clone(&field),
+ OffsetBuffer::<i64>::from_lengths((0..ROWS).flat_map(|_|
[selected_len, 128])),
+ Arc::clone(&ints),
+ None,
+ ));
+ let entries = StructArray::new(
+ vec![
+ Arc::new(Field::new("key", DataType::Int32, false)),
+ Arc::new(Field::new("value", DataType::Int32, true)),
+ ]
+ .into(),
+ vec![Arc::clone(&ints), ints],
+ None,
+ );
+ let maps = Arc::new(MapArray::new(
+ Arc::new(Field::new("entries", entries.data_type().clone(),
false)),
+ offsets,
+ entries,
+ None,
+ false,
+ ));
+ let structs = Arc::new(StructArray::new(
+ vec![Arc::new(Field::new(
+ "nested",
+ lists.data_type().clone(),
+ true,
+ ))]
+ .into(),
+ vec![Arc::<ListArray>::clone(&lists)],
+ None,
+ ));
+ let nullable_lists = Arc::new(ListArray::new(
+ field,
+ OffsetBuffer::from_lengths((0..ROWS).flat_map(|_| [selected_len,
128])),
+ Arc::new(Int32Array::from_iter(
+ (0..count as i32).map(|i| (i % 11 != 0).then_some(i)),
+ )),
+ Some((0..2 * ROWS).map(|i| i % 7 != 0).collect::<NullBuffer>()),
+ ));
+ vec![lists, large_lists, maps, structs, nullable_lists]
+ }
+
+ #[test]
+ fn test_list_extract_nested_retained_capacity() -> Result<()> {
+ use arrow::buffer::{NullBuffer, OffsetBuffer};
+ const ROWS: usize = 8192;
+ let error_wrapper = |error: SparkError| DataFusionError::from(error);
+ for selected_len in [0, 1] {
+ for values in skewed_nested_values(selected_len) {
+ for nullable in [false, true] {
+ let outer = ListArray::new(
+ Arc::new(Field::new("item",
values.data_type().clone(), true)),
+ OffsetBuffer::from_lengths(std::iter::repeat_n(2,
ROWS)),
+ Arc::clone(&values),
+ nullable.then(|| (0..ROWS).map(|row| row % 4 !=
0).collect::<NullBuffer>()),
+ );
+ let data = values.to_data();
+ let mut expected = MutableArrayData::new(vec![&data],
true, ROWS);
Review Comment:
The reference gather (`MutableArrayData::new`, then `try_extend` for valid
rows and `try_extend_nulls` for null rows, then `make_array(freeze())`) is
written out in all three new tests, here and at lines 674 and 737. A small
helper that takes the values and a `&[Option<usize>]` of source rows would
remove roughly 30 lines and make the intent of each test easier to see.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
