From 951fac225d07bd641ae6829421b07095a5f0c701 Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Mon, 3 Aug 2026 18:49:15 -0400 Subject: [PATCH 1/3] perf: preallocate rowsgroupcolumn buffers in take_n method --- .../aggregates/group_values/multi_group_by/row_backed.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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..9d2bc9c99d36e 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); } From 016e86caf2747b1ed8fa070af578d84cf4fa8f58 Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Mon, 3 Aug 2026 18:50:24 -0400 Subject: [PATCH 2/3] added tests --- .../group_values/multi_group_by/row_backed.rs | 78 +++++++++++++++++++ 1 file changed, 78 insertions(+) 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 9d2bc9c99d36e..a8a8b56acf6e1 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 @@ -429,6 +429,84 @@ mod tests { assert_eq!(g0, 20); } + /// `take_n` preallocates the retained-row buffer from the known row + /// count and byte size - this tests that rebuild path with multiple emitted and multiple retained rows, + /// asserting that every emitted value and every shifted-down value is byte-for-byte unchanged. + #[test] + fn take_n_preallocated_rebuild_preserves_all_rows() { + let dt = DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Int32, true)), + 1, + ); + let mut col = RowsGroupColumn::try_new(dt).unwrap(); + + // Include an outer-null and an inner-null row to make sure the + // preallocated byte capacity covers variable-length encodings. + let input = fsl_i32( + vec![ + Some(vec![Some(10)]), + Some(vec![Some(20)]), + None, + Some(vec![None]), + Some(vec![Some(50)]), + ], + 1, + ); + col.vectorized_append(&input, &[0, 1, 2, 3, 4]).unwrap(); + assert_eq!(col.len(), 5); + + // Emit the first three rows; two rows should remain and shift to front. + 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), + 10 + ); + assert_eq!( + emitted + .value(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 20 + ); + // Third emitted row was an outer-null. + assert!(emitted.is_null(2)); + + assert_eq!(col.len(), 2); + + // Remaining rows (inner-null, then 50) must survive the rebuild intact. + let rest = Box::new(col).build(); + let rest = rest.as_any().downcast_ref::().unwrap(); + assert_eq!(rest.len(), 2); + + let r0 = rest + .value(0) + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + assert!(r0.is_null(0), "retained inner-null row must be preserved"); + + let r1 = rest + .value(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(r1, 50); + } + /// Works for `Struct` too — proves the column is type-generic. #[test] fn struct_roundtrip() { From 9a26f49658e2389ec94c16ff1b9996304a73a39c Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Tue, 4 Aug 2026 11:39:52 -0400 Subject: [PATCH 3/3] tests updated --- .../group_values/multi_group_by/row_backed.rs | 120 +++++++++++------- 1 file changed, 74 insertions(+), 46 deletions(-) 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 a8a8b56acf6e1..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 @@ -320,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; @@ -330,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] @@ -429,33 +453,38 @@ mod tests { assert_eq!(g0, 20); } - /// `take_n` preallocates the retained-row buffer from the known row - /// count and byte size - this tests that rebuild path with multiple emitted and multiple retained rows, - /// asserting that every emitted value and every shifted-down value is byte-for-byte unchanged. + /// `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_all_rows() { + fn take_n_preallocated_rebuild_preserves_variable_length_rows() { let dt = DataType::FixedSizeList( - Arc::new(Field::new("item", DataType::Int32, true)), + Arc::new(Field::new("item", DataType::Utf8, true)), 1, ); let mut col = RowsGroupColumn::try_new(dt).unwrap(); - // Include an outer-null and an inner-null row to make sure the - // preallocated byte capacity covers variable-length encodings. - let input = fsl_i32( - vec![ - Some(vec![Some(10)]), - Some(vec![Some(20)]), - None, - Some(vec![None]), - Some(vec![Some(50)]), - ], - 1, - ); - col.vectorized_append(&input, &[0, 1, 2, 3, 4]).unwrap(); - assert_eq!(col.len(), 5); + // 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; two rows should remain and shift to front. + // Emit the first three rows; four rows should remain. let emitted = col.take_n(3); let emitted = emitted .as_any() @@ -466,45 +495,44 @@ mod tests { emitted .value(0) .as_any() - .downcast_ref::() + .downcast_ref::() .unwrap() .value(0), - 10 + "emit_a" ); - assert_eq!( + // Row 1 was an inner-null; row 2 was an outer-null. + assert!( emitted .value(1) .as_any() - .downcast_ref::() + .downcast_ref::() .unwrap() - .value(0), - 20 + .is_null(0) ); - // Third emitted row was an outer-null. assert!(emitted.is_null(2)); - assert_eq!(col.len(), 2); + assert_eq!(col.len(), 4); - // Remaining rows (inner-null, then 50) must survive the rebuild intact. + // 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(), 2); - - let r0 = rest - .value(0) - .as_any() - .downcast_ref::() - .unwrap() - .clone(); - assert!(r0.is_null(0), "retained inner-null row must be preserved"); + assert_eq!(rest.len(), 4); - let r1 = rest - .value(1) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - assert_eq!(r1, 50); + 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.