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/06/05 01:24:23 UTC
(datafusion) branch gh-readonly-queue/main/pr-22558-7b78f0cea0409f60c79a3833781e621344d52ffd created (now c1f0d54bb9)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-22558-7b78f0cea0409f60c79a3833781e621344d52ffd
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at c1f0d54bb9 perf: avoid unnecessary large allocations (#22558)
This branch includes the following new commits:
new c1f0d54bb9 perf: avoid unnecessary large allocations (#22558)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: perf: avoid unnecessary large allocations (#22558)
Posted by gi...@apache.org.
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-22558-7b78f0cea0409f60c79a3833781e621344d52ffd
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit c1f0d54bb9e1054e7c7212037e970b3d3b34885c
Author: Ariel Miculas-Trif <ar...@gmail.com>
AuthorDate: Fri Jun 5 04:24:01 2026 +0300
perf: avoid unnecessary large allocations (#22558)
## Which issue does this PR close?
Related a bit to https://github.com/apache/datafusion/issues/22526
~Needs rebasing once https://github.com/apache/datafusion/pull/22416 is
merged~
## Rationale for this change
split_off does this:
> Returns a newly allocated vector containing the elements in the range
[at, len). After the call, the original vector will be left containing
the elements [0, at) with its previous capacity unchanged.
which is bad when taking a small slice from a large Vec, for two
reasons:
* it will allocate memory for the remaining elements, which are a lot
more than n
* it will return a Vec with a very large capacity compared to its length
split_vec_min_alloc still has some issues:
https://github.com/apache/datafusion/issues/22548 but it uses drain +
collect when n is small, which is better because it only allocates for
the initial n elements and doesn't inflate the capacity
## What changes are included in this PR?
## Are these changes tested?
Yes
## Are there any user-facing changes?
No
---
datafusion/expr-common/src/groups_accumulator.rs | 60 +++++++++++++++++++---
.../group_values/single_group_by/primitive.rs | 53 +++++++++++++++++--
2 files changed, 102 insertions(+), 11 deletions(-)
diff --git a/datafusion/expr-common/src/groups_accumulator.rs b/datafusion/expr-common/src/groups_accumulator.rs
index 9053f7a8ea..da5da384c7 100644
--- a/datafusion/expr-common/src/groups_accumulator.rs
+++ b/datafusion/expr-common/src/groups_accumulator.rs
@@ -18,7 +18,7 @@
//! Vectorized [`GroupsAccumulator`]
use arrow::array::{ArrayRef, BooleanArray};
-use datafusion_common::{Result, not_impl_err};
+use datafusion_common::{Result, not_impl_err, utils::split_vec_min_alloc};
/// Describes how many rows should be emitted during grouping.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -45,13 +45,7 @@ impl EmitTo {
// Take the entire vector, leave new (empty) vector
std::mem::take(v)
}
- Self::First(n) => {
- // get end n+1,.. values into t
- let mut t = v.split_off(*n);
- // leave n+1,.. in v
- std::mem::swap(v, &mut t);
- t
- }
+ Self::First(n) => split_vec_min_alloc(v, *n),
}
}
}
@@ -254,3 +248,53 @@ pub trait GroupsAccumulator: Send + std::any::Any {
/// compute, not `O(num_groups)`
fn size(&self) -> usize;
}
+
+#[cfg(test)]
+mod tests {
+ use super::EmitTo;
+
+ /// When `n` is small relative to `len`, the old `split_off(n) + swap` pattern had
+ /// two allocation problems:
+ ///
+ /// 1. The returned Vec kept the original large backing allocation even though it
+ /// only contains `n` elements (wasted capacity on a short-lived value).
+ /// 2. `split_off` allocated a fresh Vec for the `len - n` remaining elements,
+ /// even though that side is much larger than `n` — the expensive side to
+ /// allocate.
+ ///
+ /// `split_vec_min_alloc` fixes both: when `n * 2 <= len` it uses
+ /// `drain(0..n).collect()`, allocating only `n` elements for the emitted prefix
+ /// and keeping the original large backing in the remaining accumulator.
+ #[test]
+ fn take_needed_first_small_n_allocates_minimally() {
+ let mut v: Vec<i32> = Vec::with_capacity(128);
+ v.extend(0..20i32);
+ let original_capacity = v.capacity(); // 128
+
+ // n=4, n*2=8 <= len=20 -> drain branch in split_vec_min_alloc
+ let emitted = EmitTo::First(4).take_needed(&mut v);
+
+ assert_eq!(emitted, vec![0, 1, 2, 3]);
+ assert_eq!(v, (4..20i32).collect::<Vec<_>>());
+
+ // The emitted prefix must NOT carry the original large allocation.
+ // Old split_off+swap returned a Vec with capacity=128 for only 4 elements.
+ assert!(
+ emitted.capacity() <= 4,
+ "emitted prefix capacity {} should be ~n=4, not the original {}",
+ emitted.capacity(),
+ original_capacity,
+ );
+
+ // The remaining accumulator must retain the original large allocation so
+ // that incoming groups don't immediately force a realloc.
+ // Old split_off+swap left the remaining vec with a small fresh allocation.
+ assert_eq!(
+ v.capacity(),
+ original_capacity,
+ "remaining vec capacity {} should equal original {}",
+ v.capacity(),
+ original_capacity,
+ );
+ }
+}
diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs
index efaf7eba0f..07535cfdaa 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs
@@ -24,6 +24,7 @@ use arrow::array::{
use arrow::datatypes::{DataType, i256};
use datafusion_common::Result;
use datafusion_common::hash_utils::RandomState;
+use datafusion_common::utils::split_vec_min_alloc;
use datafusion_execution::memory_pool::proxy::VecAllocExt;
use datafusion_expr::EmitTo;
use half::f16;
@@ -207,9 +208,7 @@ where
Some(_) => self.null_group.take(),
None => None,
};
- let mut split = self.values.split_off(n);
- std::mem::swap(&mut self.values, &mut split);
- build_primitive(split, null_group)
+ build_primitive(split_vec_min_alloc(&mut self.values, n), null_group)
}
};
@@ -223,3 +222,51 @@ where
self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::types::Int32Type;
+ use arrow::array::{ArrayRef, Int32Array};
+ use arrow::datatypes::DataType;
+ use datafusion_expr::EmitTo;
+ use std::sync::Arc;
+
+ /// Mirror of the `EmitTo::take_needed` regression test, applied to the
+ /// concrete `GroupValuesPrimitive` accumulator.
+ ///
+ /// When `n` is small, the old `split_off(n) + swap` pattern used inside
+ /// `emit(EmitTo::First(n))` left `self.values` with a small fresh allocation
+ /// and returned the emitted prefix carrying the original large backing.
+ ///
+ /// With `split_vec_min_alloc` and `n * 2 <= len`, the drain branch is taken:
+ /// the emitted prefix gets a compact allocation and `self.values` retains the
+ /// original large one.
+ #[test]
+ fn emit_first_small_n_allocates_minimally() -> Result<()> {
+ let mut gv = GroupValuesPrimitive::<Int32Type>::new(DataType::Int32);
+
+ // Intern 20 distinct values; `new()` pre-allocates capacity 128 for `values`.
+ let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..20i32));
+ let mut groups = vec![];
+ gv.intern(&[arr], &mut groups)?;
+ let capacity_before = gv.values.capacity(); // 128
+
+ // n=4, n*2=8 <= len=20 -> drain branch
+ let emitted = gv.emit(EmitTo::First(4))?;
+
+ assert_eq!(emitted[0].len(), 4);
+
+ // `self.values` must retain its original large allocation.
+ // Old split_off+swap left it with a fresh small allocation (~16).
+ assert_eq!(
+ gv.values.capacity(),
+ capacity_before,
+ "self.values capacity {} should equal original {} after small First(n) emit",
+ gv.values.capacity(),
+ capacity_before,
+ );
+
+ Ok(())
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org