You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/08/05 07:04:18 UTC
(datafusion) 01/01: perf: preallocate RowsGroupColumn buffers in take_n (#24070)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24070-a0124a4f0f7d9614d13c8d1753df5263a2c42b51
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 2dd1a14232b8021cb8979b0438367f1fde7c97e7
Author: Saad Tajwar <59...@users.noreply.github.com>
AuthorDate: Wed Aug 5 03:03:52 2026 -0400
perf: preallocate RowsGroupColumn buffers in take_n (#24070)
## Which issue does this PR close?
- Closes #23994
## Rationale for this change
Currently `RowsGroupColumn`'s `take_n` method knows the remaining number
of rows and bytes after taking the first n group rows, but when invoking
`empty_rows` passes in 0 for the row capacity and data capacity
parameters instead of the known remainders, potentially causing
unnecessary copying and allocating
## What changes are included in this PR?
Using the known remaining number of rows and bytes in the invocation to
`empty_rows`
## Are these changes tested?
Yes
## Are there any user-facing changes?
No
---
.../group_values/multi_group_by/row_backed.rs | 114 ++++++++++++++++++++-
1 file changed, 112 insertions(+), 2 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 31735559cd..1445a81f21 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<Utf8>` 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<Option<Option<&str>>>) -> ArrayRef {
+ let child = StringArray::from(
+ rows.iter()
+ .map(|row| row.and_then(|inner| inner))
+ .collect::<Vec<_>>(),
+ );
+ let outer_nulls = arrow::buffer::NullBuffer::from(
+ rows.iter().map(|row| row.is_some()).collect::<Vec<_>>(),
+ );
+ 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<Int32>`.
#[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<Utf8>` 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::<FixedSizeListArray>()
+ .unwrap();
+ assert_eq!(emitted.len(), 3);
+ assert_eq!(
+ emitted
+ .value(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .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::<StringArray>()
+ .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::<FixedSizeListArray>().unwrap();
+ assert_eq!(rest.len(), 4);
+
+ let value_at = |idx: usize| {
+ rest.value(idx)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .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<a: Int32>` too — proves the column is type-generic.
#[test]
fn struct_roundtrip() {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org