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/09/03 06:13:32 UTC
(datafusion) branch main updated: fix: clear memory after emit all in count distinct (#24888)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 7e5561dff9 fix: clear memory after emit all in count distinct (#24888)
7e5561dff9 is described below
commit 7e5561dff96a8103fc8a28a0723deb6b252a6b49
Author: Raz Luvaton <16...@users.noreply.github.com>
AuthorDate: Thu Sep 3 06:03:33 2026 +0000
fix: clear memory after emit all in count distinct (#24888)
## Which issue does this PR close?
N/A
## Rationale for this change
after emit all, the memory is still being held in count distinct causing
OOM issues
## What changes are included in this PR?
release memory in emit all and added tests
## What is the testing strategy for this PR?
integration test
## Are there any user-facing changes?
no
----
Founded while running:
- #24881
---
datafusion/core/tests/memory_limit/mod.rs | 84 +++++++++++++++++++++-
.../src/aggregate/count_distinct/groups.rs | 9 ++-
2 files changed, 90 insertions(+), 3 deletions(-)
diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs
index 15b224d200..39be33587f 100644
--- a/datafusion/core/tests/memory_limit/mod.rs
+++ b/datafusion/core/tests/memory_limit/mod.rs
@@ -26,11 +26,14 @@ mod nlj_spill_unmatched;
mod repartition_mem_limit;
mod union_nullable_spill;
mod view_spill_compaction;
-use arrow::array::{ArrayRef, DictionaryArray, Int32Array, RecordBatch, StringViewArray};
+use arrow::array::{
+ ArrayRef, DictionaryArray, Int32Array, Int64Array, RecordBatch, StringViewArray,
+};
use arrow::compute::SortOptions;
use arrow::datatypes::{Int32Type, SchemaRef};
use arrow_schema::{DataType, Field, Schema};
use datafusion::assert_batches_eq;
+use datafusion::assert_batches_sorted_eq;
use datafusion::config::SpillCompression;
use datafusion::datasource::memory::MemorySourceConfig;
use datafusion::datasource::source::DataSourceExec;
@@ -125,6 +128,85 @@ async fn group_by_hash() {
.await
}
+/// `count(distinct)` over integers under a memory limit.
+///
+/// The integer distinct-count groups accumulator reports the capacity of its
+/// buffers in `size()`. After an aggregate stream emits all groups, either to
+/// emit partial state early or to spill, it resizes its reservation to the
+/// table's reported size and expects it to have shrunk. If the accumulator
+/// keeps its capacity, that resize is a grow against an exhausted pool and the
+/// query fails although everything was already written out.
+const COUNT_DISTINCT_ROWS: usize = 200_000;
+const COUNT_DISTINCT_GROUPS: i64 = 64;
+const COUNT_DISTINCT_BATCH_ROWS: usize = 8_192;
+
+/// Far below the distinct sets (200k values, several megabytes across the
+/// partial and final tables), far above the fixed cost of the stages. With the
+/// accumulator releasing its buffers the query passes from 2 MB upwards;
+/// without, it fails up to 4 MB with "Decreasing allocation after spilling
+/// should succeed" in the final stage or a failed emit in the partial stage.
+const COUNT_DISTINCT_MEMORY_LIMIT: usize = 4 * 1024 * 1024;
+
+/// `g` has 64 groups, `v` is unique, so every group holds 3125 distinct values.
+fn count_distinct_table() -> MemTable {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("g", DataType::Int64, false),
+ Field::new("v", DataType::Int64, false),
+ ]));
+ let batches = (0..COUNT_DISTINCT_ROWS)
+ .step_by(COUNT_DISTINCT_BATCH_ROWS)
+ .map(|start| {
+ let rows =
+ start..(start + COUNT_DISTINCT_BATCH_ROWS).min(COUNT_DISTINCT_ROWS);
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int64Array::from_iter_values(
+ rows.clone().map(|row| row as i64 % COUNT_DISTINCT_GROUPS),
+ )),
+ Arc::new(Int64Array::from_iter_values(rows.map(|row| row as i64))),
+ ],
+ )
+ .unwrap()
+ })
+ .collect();
+ MemTable::try_new(schema, vec![batches]).unwrap()
+}
+
+/// Four partial stages emit their state early and four hash-partitioned final
+/// stages spill; every one of them must see the accumulator memory drop after
+/// emitting all groups.
+#[tokio::test]
+async fn count_distinct_releases_memory_after_emitting_all() {
+ let runtime = RuntimeEnvBuilder::new()
+ .with_memory_limit(COUNT_DISTINCT_MEMORY_LIMIT, 1.0)
+ .with_disk_manager_builder(DiskManagerBuilder::default())
+ .build_arc()
+ .unwrap();
+ let config = SessionConfig::new().with_target_partitions(4);
+ let ctx = SessionContext::new_with_config_rt(config, runtime);
+ ctx.register_table("t", Arc::new(count_distinct_table()))
+ .unwrap();
+
+ let batches = ctx
+ .sql("select count(distinct v) as d, count(*) as n from t group by g")
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap_or_else(|error| panic!("query failed under the memory limit: {error}"));
+
+ let per_group = (COUNT_DISTINCT_ROWS as i64 / COUNT_DISTINCT_GROUPS).to_string();
+ let row = format!("| {per_group} | {per_group} |");
+ let mut expected = vec!["+------+------+", "| d | n |", "+------+------+"];
+ expected.extend(std::iter::repeat_n(
+ row.as_str(),
+ COUNT_DISTINCT_GROUPS as usize,
+ ));
+ expected.push("+------+------+");
+ assert_batches_sorted_eq!(expected, &batches);
+}
+
#[tokio::test]
async fn join_by_key_multiple_partitions() {
let config = SessionConfig::new().with_target_partitions(2);
diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs
index 6e3e3b91a7..38f6d26200 100644
--- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs
+++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs
@@ -89,7 +89,10 @@ where
match emit_to {
EmitTo::All => {
- self.seen.clear();
+ // Release the capacity, not just the entries: `size()` reports
+ // capacity, and the aggregate streams rely on it dropping after
+ // emitting everything.
+ self.seen = HashSet::default();
}
EmitTo::First(n) => {
let mut remaining = HashSet::default();
@@ -145,7 +148,9 @@ where
all_values[pos] = value;
cursors[group_idx] += 1;
}
- self.counts.clear();
+ // Release the capacity, see `evaluate`.
+ self.seen = HashSet::default();
+ self.counts = Vec::new();
} else {
let mut remaining = HashSet::default();
for (group_idx, value) in self.seen.drain() {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org