rich7420 commented on code in PR #6498:
URL: https://github.com/apache/datafusion-comet/pull/6498#discussion_r4277563935
##########
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:
got it, thanks for the advice
--
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]
