diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs index 31735559cdb42..1445a81f2189b 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs @@ -302,7 +302,11 @@ impl GroupColumn for RowsGroupColumn { // Shift the remaining rows to the front by rebuilding the buffer. // TODO: mirror the arrow-rs efficiency TODO in `GroupValuesRows::emit`. - let mut remaining = self.row_converter.empty_rows(0, 0); + let remaining_rows = self.group_values.num_rows() - n; + let remaining_bytes = self.group_values.lengths().skip(n).sum(); + let mut remaining = self + .row_converter + .empty_rows(remaining_rows, remaining_bytes); for row in self.group_values.iter().skip(n) { remaining.push(row); } @@ -316,7 +320,9 @@ impl GroupColumn for RowsGroupColumn { mod tests { use super::*; - use arrow::array::{Array, ArrayRef, FixedSizeListArray, Int32Array, StructArray}; + use arrow::array::{ + Array, ArrayRef, FixedSizeListArray, Int32Array, StringArray, StructArray, + }; use arrow::datatypes::{DataType, Field, Int32Type}; use std::sync::Arc; @@ -326,6 +332,28 @@ mod tests { )) } + /// Build a `FixedSizeList` with `list_len == 1`. Each entry is one + /// row holding a single (optionally null) string, and an outer `None` + /// marks a null list. Variable-length string payloads give retained rows + /// distinct encoded lengths, which is what `take_n`'s byte preallocation + /// depends on. + fn fsl_utf8(rows: Vec>>) -> ArrayRef { + let child = StringArray::from( + rows.iter() + .map(|row| row.and_then(|inner| inner)) + .collect::>(), + ); + let outer_nulls = arrow::buffer::NullBuffer::from( + rows.iter().map(|row| row.is_some()).collect::>(), + ); + Arc::new(FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Utf8, true)), + 1, + Arc::new(child), + Some(outer_nulls), + )) + } + /// The generic column must agree with a per-row reference for equality, /// including inner-null and outer-null rows, on a `FixedSizeList`. #[test] @@ -425,6 +453,88 @@ mod tests { assert_eq!(g0, 20); } + /// `take_n` preallocates the retained-row buffer from the known retained + /// row count and byte size + /// + /// To exercise the byte-sum path directly, the retained rows are + /// `FixedSizeList` values with deliberately unequal payload + /// lengths plus an inner-null. Here we assert every emitted and + /// every shifted-down value is byte-for-byte unchanged. + #[test] + fn take_n_preallocated_rebuild_preserves_variable_length_rows() { + let dt = DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Utf8, true)), + 1, + ); + let mut col = RowsGroupColumn::try_new(dt).unwrap(); + + // Rows 0-2 are emitted; rows 3-6 are retained and shifted to the + // front. The retained rows intentionally have different encoded + // lengths so `lengths().skip(3).sum()` is not a simple row_count * k. + let input = fsl_utf8(vec![ + Some(Some("emit_a")), // 0: emitted + Some(None), // 1: emitted (inner-null) + None, // 2: emitted (outer-null) + Some(Some("")), // 3: retained, empty payload + Some(Some("xyz")), // 4: retained, short payload + Some(None), // 5: retained, inner-null + Some(Some("a_much_longer_payload_string")), // 6: retained, long payload + ]); + col.vectorized_append(&input, &[0, 1, 2, 3, 4, 5, 6]) + .unwrap(); + assert_eq!(col.len(), 7); + + // Emit the first three rows; four rows should remain. + let emitted = col.take_n(3); + let emitted = emitted + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(emitted.len(), 3); + assert_eq!( + emitted + .value(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + "emit_a" + ); + // Row 1 was an inner-null; row 2 was an outer-null. + assert!( + emitted + .value(1) + .as_any() + .downcast_ref::() + .unwrap() + .is_null(0) + ); + assert!(emitted.is_null(2)); + + assert_eq!(col.len(), 4); + + // The four retained rows must survive the rebuild intact, in order: + // "", "xyz", inner-null, "a_much_longer_payload_string". + let rest = Box::new(col).build(); + let rest = rest.as_any().downcast_ref::().unwrap(); + assert_eq!(rest.len(), 4); + + let value_at = |idx: usize| { + rest.value(idx) + .as_any() + .downcast_ref::() + .unwrap() + .clone() + }; + assert_eq!(value_at(0).value(0), ""); + assert_eq!(value_at(1).value(0), "xyz"); + assert!( + value_at(2).is_null(0), + "retained inner-null row must be preserved" + ); + assert_eq!(value_at(3).value(0), "a_much_longer_payload_string"); + } + /// Works for `Struct` too — proves the column is type-generic. #[test] fn struct_roundtrip() {