diff --git a/datafusion/physical-plan/src/aggregates/group_values/group_column.rs b/datafusion/physical-plan/src/aggregates/group_values/group_column.rs index 7bd083d36c510..329cc83bd6c0e 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/group_column.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/group_column.rs @@ -1077,18 +1077,40 @@ mod tests { #[test] fn test_byte_view_take_n() { + // `take_n` is really complex, we should consider and test following situations: + // 1. Take all `inlined`s + // 2. Take non-inlined + partial last buffer in `completed` + // 3. Take non-inlined + whole last buffer in `completed` + // 4. Take non-inlined + partial last `in_progress` + // 5. Take non-inlined + while last buffer in ``in_progress` + // 6. Take all views at once + let mut builder = ByteViewGroupValueBuilder::::new().with_max_block_size(60); let input_array = StringViewArray::from(vec![ - Some("this string is quite long"), // in buffer 0 + // Test situation 1 (also test take null together) + None, Some("foo"), Some("bar"), - Some("this string is also quite long"), // buffer 0 + // Test situation 2 (also test take null + inlined) None, - Some("this string is quite long"), // buffer 1 + Some("foo"), + Some("this string is quite long"), + Some("this string is also quite long"), + // Test situation 3 (also test take null + inlined) None, - Some("another string that is is quite long"), // buffer 1 Some("bar"), + Some("this string is quite long"), + // Test situation 4 (also test take null + inlined) + None, + Some("foo"), + Some("another string that is is quite long"), + Some("this string not so long"), + // Test situation 5 (also test take null + inlined + insert again after taking) + None, + Some("bar"), + Some("this string is quite long"), + // Finally, we insert the whole array again for testing situation 6 ]); let input_array: ArrayRef = Arc::new(input_array); for row in 0..input_array.len() { @@ -1097,45 +1119,71 @@ mod tests { // should be 2 completed, one in progress buffer to hold all output assert_eq!(builder.completed.len(), 2); + println!("{}", builder.in_progress.len()); assert!(!builder.in_progress.is_empty()); - let first_4 = builder.take_n(4); + // Take all `inlined`s + let taken_array = builder.take_n(2); println!( "{}", arrow::util::pretty::pretty_format_columns( - "first_4", - &[Arc::clone(&first_4)] + "taken_array", + &[Arc::clone(&taken_array)] ) .unwrap() ); - assert_eq!(&first_4, &input_array.slice(0, 4)); + assert_eq!(&taken_array, &input_array.slice(0, 2)); - // Add some new data after the first n - let input_array = StringViewArray::from(vec![ - Some("short"), - None, - Some("Some new data to add that is long"), // in buffer 0 - Some("short again"), - ]); - let input_array: ArrayRef = Arc::new(input_array); - for row in 0..input_array.len() { - builder.append_val(&input_array, row); - } + // Take non-inlined + partial buffer 0 + let taken_array = builder.take_n(1); + println!( + "{}", + arrow::util::pretty::pretty_format_columns( + "taken_array", + &[Arc::clone(&taken_array)] + ) + .unwrap() + ); + assert_eq!(&taken_array, &input_array.slice(2, 1)); - let result = Box::new(builder).build(); - let expected: ArrayRef = Arc::new(StringViewArray::from(vec![ - // last rows of the original input - None, - Some("this string is quite long"), - None, - Some("another string that is is quite long"), - Some("bar"), - // the subsequent input - Some("short"), - None, - Some("Some new data to add that is long"), // in buffer 0 - Some("short again"), - ])); - assert_eq!(&result, &expected); + // Take non-inlined + remaining partial buffer 0 + let taken_array = builder.take_n(1); + println!( + "{}", + arrow::util::pretty::pretty_format_columns( + "taken_array", + &[Arc::clone(&taken_array)] + ) + .unwrap() + ); + assert_eq!(&taken_array, &input_array.slice(3, 1)); + + // Add some new data after the first n + // let input_array = StringViewArray::from(vec![ + // Some("short"), + // None, + // Some("Some new data to add that is long"), // in buffer 0 + // Some("short again"), + // ]); + // let input_array: ArrayRef = Arc::new(input_array); + // for row in 0..input_array.len() { + // builder.append_val(&input_array, row); + // } + + // let result = Box::new(builder).build(); + // let expected: ArrayRef = Arc::new(StringViewArray::from(vec![ + // // last rows of the original input + // None, + // Some("this string is quite long"), + // None, + // Some("another string that is is quite long"), + // Some("bar"), + // // the subsequent input + // Some("short"), + // None, + // Some("Some new data to add that is long"), // in buffer 0 + // Some("short again"), + // ])); + // assert_eq!(&result, &expected); } }