sunchao commented on code in PR #6107:
URL: https://github.com/apache/datafusion-comet/pull/6107#discussion_r4433047749
##########
native/operators/src/explode.rs:
##########
@@ -1439,14 +1733,675 @@ mod tests {
#[tokio::test]
async fn single_row_exceeding_batch_size_is_sliced() {
- // One row cannot be chunked on the input side, so the oversized build
is sliced.
let batches = explode(vec![list_batch(&[Some(25)])], 10,
NullHandling::Preserve)
.await
.unwrap();
assert_eq!(sizes(&batches), vec![10, 10, 5]);
assert_eq!(values(&batches), seq(25));
}
+ #[tokio::test]
+ async fn ordinary_carried_strings_keep_multi_row_chunks() {
+ let lists = list_batch(&vec![Some(300); 8192]);
+ let payload = "x".repeat(1024);
+ let strings =
Arc::new(StringArray::from_iter_values(std::iter::repeat_n(
+ payload.as_str(),
+ 8192,
+ ))) as ArrayRef;
+ let schema = Arc::new(Schema::new(vec![
+ lists.schema().field(0).clone(),
+ Field::new("payload", DataType::Utf8, false),
+ ]));
+ let input = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::clone(lists.column(0)), strings],
+ )
+ .unwrap();
+ let source = MemorySourceConfig::try_new_exec(&[vec![input]], schema,
None).unwrap();
+ let output_schema = Arc::new(Schema::new(vec![
+ Field::new("l", DataType::Int32, true),
+ Field::new("payload", DataType::Utf8, false),
+ ]));
+ let exec = ExplodeExec::new(
+ source,
+ vec![ListUnnest {
+ index_in_input_schema: 0,
+ depth: 1,
+ }],
+ vec![],
+ output_schema,
+ UnnestOptions::new().with_null_handling(NullHandling::Drop),
+ )
+ .unwrap();
+ let context = Arc::new(
+
TaskContext::default().with_session_config(SessionConfig::new().with_batch_size(8192)),
+ );
+ let mut stream = exec.execute(0, context).unwrap();
+ let first = stream.next().await.unwrap().unwrap();
+ assert_eq!(first.num_rows(), 27 * 300);
+ assert_eq!(
+ first.column(1).as_string::<i32>().value_data().len(),
+ 8100 * 1024
+ );
+ }
+
+ #[tokio::test]
+ async fn skewed_carried_bytes_use_each_rows_repetition() {
+ use arrow::array::BinaryArray;
+
+ let wide = "w".repeat(16 * 1024 * 1024);
+ for wide_row in [0, 16, 32] {
+ let mut lens = vec![Some(300); 33];
+ lens[wide_row] = Some(1);
+ let lists = list_batch(&lens);
+ let strings = StringArray::from_iter_values((0..33).map(|row| {
+ if row == wide_row {
+ wide.as_str()
+ } else {
+ "x"
+ }
+ }));
+ let binary = Arc::new(BinaryArray::new(
+ strings.offsets().clone(),
+ strings.values().clone(),
+ None,
+ )) as ArrayRef;
+ let strings = Arc::new(strings) as ArrayRef;
+ let structure = Arc::new(StructArray::from(vec![
+ (
+ Arc::new(Field::new("s", DataType::Utf8, false)),
+ Arc::clone(&strings),
+ ),
+ (
+ Arc::new(Field::new("b", DataType::Binary, false)),
+ Arc::clone(&binary),
+ ),
+ ])) as ArrayRef;
+ let indices = Int64Array::from_iter_values(
+ lens.iter()
+ .enumerate()
+ .flat_map(|(row, len)| std::iter::repeat_n(row as i64,
len.unwrap())),
+ );
+ for carried in [strings, binary, structure] {
+ let input = RecordBatch::try_from_iter(vec![
+ ("l", Arc::clone(lists.column(0))),
+ ("payload", Arc::clone(&carried)),
+ ])
+ .unwrap();
+ let batches = explode(vec![input], 8192, NullHandling::Drop)
+ .await
+ .unwrap();
+ assert_eq!(batches.len(), 2, "wide row at {wide_row}");
+ let schema = batches[0].schema();
+ let expected = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::clone(lists.column(0).as_list::<i32>().values()),
+ kernels::take::take(carried.as_ref(), &indices,
None).unwrap(),
+ ],
+ )
+ .unwrap();
+ assert_eq!(expected.num_rows(), 9601);
+ assert_eq!(
+ arrow::compute::concat_batches(&schema, &batches).unwrap(),
+ expected
+ );
+ }
+ }
+ }
+
+ #[tokio::test]
+ async fn sliced_carried_bytes_ignore_unused_backing() {
+ let padding = "p".repeat(64 * 1024 * 1024);
+ let carried = StringArray::from(vec![padding.as_str(), "x"]).slice(1,
1);
+ let lists = list_batch(&[Some(1024)]);
+ let input = RecordBatch::try_from_iter(vec![
+ ("l", Arc::clone(lists.column(0))),
+ ("payload", Arc::new(carried) as ArrayRef),
+ ])
+ .unwrap();
+ let batches = explode(vec![input], 8192, NullHandling::Drop)
+ .await
+ .unwrap();
+ assert_eq!(sizes(&batches), vec![1024]);
+ assert_eq!(values(&batches), seq(1024));
+ assert!(batches[0]
+ .column(1)
+ .as_string::<i32>()
+ .iter()
+ .all(|s| s == Some("x")));
+ }
+
+ #[tokio::test]
+ async fn null_carried_bytes_do_not_spend_the_offset_budget() {
+ let wide = "w".repeat(2 * 1024 * 1024);
+ let strings = StringArray::from(vec![wide.as_str(), "x"]);
+ let carried = StringArray::new(
+ strings.offsets().clone(),
+ strings.values().clone(),
+ Some(NullBuffer::from(vec![false, true])),
+ );
+ let lists = list_batch(&[Some(1024), Some(300)]);
+ let input = RecordBatch::try_from_iter(vec![
+ ("l", Arc::clone(lists.column(0))),
+ ("payload", Arc::new(carried) as ArrayRef),
+ ])
+ .unwrap();
+ let batches = explode(vec![input], 8192, NullHandling::Drop)
+ .await
+ .unwrap();
+ assert_eq!(sizes(&batches), vec![1324]);
+ assert_eq!(values(&batches), seq(1324));
+ let carried = batches[0].column(1).as_string::<i32>();
+ assert_eq!(carried.null_count(), 1024);
+ assert_eq!(carried.value_data().len(), 300);
+ }
+
+ #[test]
+ fn null_struct_children_still_constrain_copied_offsets() {
+ let wide = "w".repeat(2 * 1024 * 1024);
+ let strings = StringArray::from(vec![wide.as_str(), "x"]);
+ let repeats = Int64Array::from(vec![1024, 300]);
+ assert!(!repeated_offsets_fit(&strings, &repeats, 1024));
+ let nullable_strings = StringArray::new(
+ strings.offsets().clone(),
+ strings.values().clone(),
+ Some(NullBuffer::from(vec![false, true])),
+ );
+ assert!(repeated_offsets_fit(&nullable_strings, &repeats, 1024));
+ let structure = StructArray::new(
+ vec![Arc::new(Field::new("s", DataType::Utf8, false))].into(),
+ vec![Arc::new(strings)],
+ Some(NullBuffer::from(vec![false, true])),
+ );
+ assert!(!repeated_offsets_fit(&structure, &repeats, 1024));
+ }
+
+ #[test]
+ fn repetition_limit_ignores_fixed_width_buffers() {
+ let batch = RecordBatch::try_from_iter(vec![(
+ "id",
+ Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef,
+ )])
+ .unwrap();
+ assert_eq!(repetition_limit(&batch, &[]), i32::MAX as usize);
+ }
+
+ #[tokio::test]
+ async fn
carried_nested_columns_do_not_preallocate_for_unselected_children() {
+ use arrow::array::MapArray;
+
+ let lists = list_batch(&vec![Some(2); 1024]);
+ let list = lists.column(0).as_list::<i32>();
+ let entries = StructArray::from(vec![
+ (
+ Arc::new(Field::new("key", DataType::Int32, false)),
+ Arc::clone(list.values()),
+ ),
+ (
+ Arc::new(Field::new("value", DataType::Int32, true)),
+ Arc::clone(list.values()),
+ ),
+ ]);
+ let map = Arc::new(MapArray::new(
+ Arc::new(Field::new("entries", entries.data_type().clone(),
false)),
+ list.offsets().clone(),
+ entries,
+ None,
+ false,
+ )) as ArrayRef;
+ let nested_field = Arc::new(Field::new("item",
map.data_type().clone(), true));
+ let carried: Vec<ArrayRef> = vec![
+ Arc::clone(&map),
+ Arc::clone(lists.column(0)),
+ Arc::new(LargeListArray::new(
+ Arc::new(Field::new("item", DataType::Int32, true)),
+ OffsetBuffer::new(list.value_offsets().iter().map(|&o|
i64::from(o)).collect()),
+ Arc::clone(list.values()),
+ None,
+ )),
+ Arc::new(StructArray::from(vec![(
+ Arc::new(Field::new("m", map.data_type().clone(), true)),
+ Arc::clone(&map),
+ )])),
+ Arc::new(FixedSizeListArray::new(
+ Arc::clone(&nested_field),
+ 1,
+ Arc::clone(&map),
+ None,
+ )),
+ Arc::new(ListArray::new(
+ nested_field,
+ OffsetBuffer::new((0..=1024).collect()),
+ map,
+ None,
+ )),
+ ];
+
+ fn assert_small_buffers(array: &ArrayRef) {
+ match array.data_type() {
+ DataType::Map(_, _) => {
+ for column in array.as_map().entries().columns() {
+ assert_small_buffers(column);
+ }
+ }
+ DataType::List(_) =>
assert_small_buffers(array.as_list::<i32>().values()),
+ DataType::LargeList(_) =>
assert_small_buffers(array.as_list::<i64>().values()),
+ DataType::FixedSizeList(_, _) => {
+ assert_small_buffers(array.as_fixed_size_list().values());
+ }
+ DataType::Struct(_) => {
+ for column in array.as_struct().columns() {
+ assert_small_buffers(column);
+ }
+ }
+ _ => assert!(
+ array.get_buffer_memory_size() <= 1024,
+ "eight output rows must not allocate for the 1024-row
source: {} bytes",
+ array.get_buffer_memory_size()
+ ),
+ }
+ }
+
+ for carried in carried {
+ let input = RecordBatch::try_from_iter(vec![
+ ("l", Arc::clone(lists.column(0))),
+ ("carried", Arc::clone(&carried)),
+ ])
+ .unwrap();
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("l", DataType::Int32, true),
+ input.schema().field(1).clone(),
+ ]));
+ let input_schema = input.schema();
+ let source =
+ MemorySourceConfig::try_new_exec(&[vec![input]], input_schema,
None).unwrap();
+ let exec = ExplodeExec::new(
+ source,
+ vec![ListUnnest {
+ index_in_input_schema: 0,
+ depth: 1,
+ }],
+ vec![],
+ Arc::clone(&schema),
+ UnnestOptions::new(),
+ )
+ .unwrap();
+ let context = Arc::new(
+
TaskContext::default().with_session_config(SessionConfig::new().with_batch_size(8)),
+ );
+ let mut stream = exec.execute(0, context).unwrap();
+ for start in [0, 8] {
+ let actual = stream.next().await.unwrap().unwrap();
+ let expected = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int32Array::from_iter_values(start..start +
8)),
+ kernels::take::take(
+ &carried,
+ &Int64Array::from_iter_values(
+ (start / 2..start / 2 + 4).flat_map(|row|
[i64::from(row); 2]),
+ ),
+ None,
+ )
+ .unwrap(),
+ ],
+ )
+ .unwrap();
+ assert_eq!(actual, expected);
+ assert_small_buffers(actual.column(1));
+ }
+ }
+ }
+
+ #[test]
+ fn repetition_limit_uses_selected_nested_offsets() {
+ use arrow::array::NullArray;
+
+ let length = i32::MAX as usize / 2 + 1;
+ let list = ListArray::new(
+ Arc::new(Field::new("item", DataType::Null, true)),
+ OffsetBuffer::new(vec![0, length as i32 - 1, length as
i32].into()),
+ Arc::new(NullArray::new(length)),
+ None,
+ );
+ let batch =
+ RecordBatch::try_from_iter(vec![("carried", Arc::new(list) as
ArrayRef)]).unwrap();
+ assert_eq!(repetition_limit(&batch.slice(1, 1), &[]), i32::MAX as
usize);
+ }
+
+ #[test]
+ fn repetition_limit_follows_nested_variable_width_slices() {
+ use arrow::array::{BinaryArray, MapArray};
+
+ let strings = Arc::new(StringArray::from(vec!["ignored prefix", "a",
"bc"])) as ArrayRef;
+ let field = Arc::new(Field::new("item", DataType::Utf8, false));
+ let list = ListArray::new(
+ Arc::clone(&field),
+ OffsetBuffer::new(vec![0, 1, 3].into()),
+ Arc::clone(&strings),
+ None,
+ );
+ let entries = StructArray::from(vec![
+ (
+ Arc::new(Field::new("key", DataType::Int32, false)),
+ Arc::new(Int32Array::from(vec![0, 1, 2])) as ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("value", DataType::Utf8, false)),
+ Arc::clone(&strings),
+ ),
+ ]);
+ let map = MapArray::new(
+ Arc::new(Field::new("entries", entries.data_type().clone(),
false)),
+ list.offsets().clone(),
+ entries,
+ None,
+ false,
+ );
+ let arrays: Vec<ArrayRef> = vec![
+ strings.slice(1, 2),
+ Arc::new(BinaryArray::from(vec![b"prefix".as_slice(), b"a",
b"bc"]).slice(1, 2)),
+ Arc::new(list.slice(1, 1)),
+ Arc::new(map.slice(1, 1)),
+ Arc::new(
+ StructArray::from(vec![(Arc::clone(&field),
Arc::clone(&strings))]).slice(1, 2),
+ ),
+ Arc::new(
+ LargeListArray::new(
+ Arc::clone(&field),
+ OffsetBuffer::new(vec![0i64, 1, 3].into()),
+ Arc::clone(&strings),
+ None,
+ )
+ .slice(1, 1),
+ ),
+ Arc::new(FixedSizeListArray::new(field, 1, strings, None).slice(1,
2)),
+ ];
+ for array in arrays {
+ let batch = RecordBatch::try_from_iter(vec![("carried",
array)]).unwrap();
+ assert_eq!(repetition_limit(&batch, &[]), i32::MAX as usize / 3);
+ }
+ }
+
+ #[test]
+ fn chunk_offset_bound_counts_all_repeated_rows() {
+ use arrow::array::NullArray;
+
+ let child_length = i32::MAX / 3;
+ let carried = ListArray::new(
+ Arc::new(Field::new("item", DataType::Null, true)),
+ OffsetBuffer::new(vec![0, child_length, 2 * child_length].into()),
+ Arc::new(NullArray::new(2 * child_length as usize)),
+ None,
+ );
+ let lists = list_batch(&[Some(2), Some(2)]);
+ let batch = RecordBatch::try_from_iter(vec![
+ ("l", Arc::clone(lists.column(0))),
+ ("carried", Arc::new(carried) as ArrayRef),
+ ])
+ .unwrap();
+ let list_columns = [ListUnnest {
+ index_in_input_schema: 0,
+ depth: 1,
+ }];
+ let mut pending = PendingInput {
+ batch,
+ row_offset: 0,
+ element_offset: 0,
+ safe_repetitions: Some(1),
+ output_lens: Some(Int64Array::from(vec![2, 2])),
+ };
+ // Each row can repeat twice, but their combined child offsets would
overflow.
+ for _ in 0..2 {
+ let (chunk, lengths) = pending.next_chunk(4,
&list_columns).unwrap();
+ assert_eq!(chunk.num_rows(), 1);
+ assert_eq!(lengths.unwrap().values(), &[2]);
+ }
+ assert_eq!(pending.remaining_rows(), 0);
+ }
+
+ #[tokio::test]
+ async fn splits_one_list_before_repeating_carried_values() {
+ let list = list_batch(&[Some(25)]);
+ let schema = Arc::new(Schema::new(vec![
+ list.schema().field(0).clone(),
+ Field::new("payload", DataType::Utf8, false),
+ ]));
+ let payload = "x".repeat(1024);
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::clone(list.column(0)),
+ Arc::new(StringArray::from(vec![payload])),
+ ],
+ )
+ .unwrap();
+ let source = MemorySourceConfig::try_new_exec(&[vec![batch]], schema,
None).unwrap();
+ let output_schema = Arc::new(Schema::new(vec![
+ Field::new("l", DataType::Int32, true),
+ Field::new("payload", DataType::Utf8, false),
+ ]));
+ let exec = ExplodeExec::new(
+ source,
+ vec![ListUnnest {
+ index_in_input_schema: 0,
+ depth: 1,
+ }],
+ vec![],
+ output_schema,
+ UnnestOptions::new().with_null_handling(NullHandling::Drop),
+ )
+ .unwrap();
+ let context = Arc::new(
+
TaskContext::default().with_session_config(SessionConfig::new().with_batch_size(10)),
+ );
+ let mut stream = exec.execute(0, context).unwrap();
+ let mut actual = Vec::new();
+ while let Some(batch) = stream.next().await {
+ let batch = batch.unwrap();
+ let strings = batch.column(1).as_string::<i32>();
+ assert_eq!(
+ strings.value_data().len(),
+ batch.num_rows() * 1024,
+ "carried values must be materialized only for this output
chunk"
+ );
+ actual.extend(batch.column(0).as_primitive::<Int32Type>().iter());
+ }
+ assert_eq!(actual, seq(25));
+ }
+
+ #[tokio::test]
+ async fn wide_carried_string_respects_offset_limit_at_large_batch_size() {
Review Comment:
This test repeats a 1 MiB UTF-8 value 2,049 times and materializes the
output. The first 2,047-row batch alone needs 2,047 MiB for the string payload.
Could the test assert the real 32-bit boundary arithmetically, then drive the
stream with a small injectable byte limit to verify splitting and output
values? That would keep the regression in the default suite without requiring a
multi-gigabyte allocation.
--
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]
