From ff64291d1598244d8254fa1bec54ee6ac4f780cd Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 2 Sep 2026 22:50:13 +0300 Subject: [PATCH 1/2] fix: clear memory after emit all --- .../memory_limit/count_distinct_spill.rs | 101 ++++++++++++++++++ datafusion/core/tests/memory_limit/mod.rs | 1 + .../src/aggregate/count_distinct/groups.rs | 9 +- 3 files changed, 109 insertions(+), 2 deletions(-) create mode 100644 datafusion/core/tests/memory_limit/count_distinct_spill.rs diff --git a/datafusion/core/tests/memory_limit/count_distinct_spill.rs b/datafusion/core/tests/memory_limit/count_distinct_spill.rs new file mode 100644 index 0000000000000..d5764bdef7785 --- /dev/null +++ b/datafusion/core/tests/memory_limit/count_distinct_spill.rs @@ -0,0 +1,101 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! `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. + +use std::sync::Arc; + +use arrow::array::{Int64Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema}; +use datafusion::datasource::MemTable; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_common::assert_batches_sorted_eq; +use datafusion_execution::disk_manager::DiskManagerBuilder; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; + +const ROWS: usize = 200_000; +const GROUPS: i64 = 64; +const 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 MEMORY_LIMIT: usize = 4 * 1024 * 1024; + +/// `g` has 64 groups, `v` is unique, so every group holds 3125 distinct values. +fn table() -> MemTable { + let schema = Arc::new(Schema::new(vec![ + Field::new("g", DataType::Int64, false), + Field::new("v", DataType::Int64, false), + ])); + let batches = (0..ROWS) + .step_by(BATCH_ROWS) + .map(|start| { + let rows = start..(start + BATCH_ROWS).min(ROWS); + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from_iter_values( + rows.clone().map(|row| row as i64 % 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(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(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 = (ROWS as i64 / 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(), GROUPS as usize)); + expected.push("+------+------+"); + assert_batches_sorted_eq!(expected, &batches); +} diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 15b224d200bf4..0e914ffd81887 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -20,6 +20,7 @@ use std::num::NonZeroUsize; use std::sync::{Arc, LazyLock}; +mod count_distinct_spill; #[cfg(feature = "extended_tests")] mod memory_limit_validation; mod nlj_spill_unmatched; 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 6e3e3b91a74f7..38f6d26200ec9 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() { From 5bc99b67c71ae08cd4ae4a91c7d2130919647393 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 2 Sep 2026 23:21:23 +0300 Subject: [PATCH 2/2] test: move count distinct spill test into memory_limit tests --- .../memory_limit/count_distinct_spill.rs | 101 ------------------ datafusion/core/tests/memory_limit/mod.rs | 85 ++++++++++++++- 2 files changed, 83 insertions(+), 103 deletions(-) delete mode 100644 datafusion/core/tests/memory_limit/count_distinct_spill.rs diff --git a/datafusion/core/tests/memory_limit/count_distinct_spill.rs b/datafusion/core/tests/memory_limit/count_distinct_spill.rs deleted file mode 100644 index d5764bdef7785..0000000000000 --- a/datafusion/core/tests/memory_limit/count_distinct_spill.rs +++ /dev/null @@ -1,101 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! `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. - -use std::sync::Arc; - -use arrow::array::{Int64Array, RecordBatch}; -use arrow::datatypes::{DataType, Field, Schema}; -use datafusion::datasource::MemTable; -use datafusion::prelude::{SessionConfig, SessionContext}; -use datafusion_common::assert_batches_sorted_eq; -use datafusion_execution::disk_manager::DiskManagerBuilder; -use datafusion_execution::runtime_env::RuntimeEnvBuilder; - -const ROWS: usize = 200_000; -const GROUPS: i64 = 64; -const 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 MEMORY_LIMIT: usize = 4 * 1024 * 1024; - -/// `g` has 64 groups, `v` is unique, so every group holds 3125 distinct values. -fn table() -> MemTable { - let schema = Arc::new(Schema::new(vec![ - Field::new("g", DataType::Int64, false), - Field::new("v", DataType::Int64, false), - ])); - let batches = (0..ROWS) - .step_by(BATCH_ROWS) - .map(|start| { - let rows = start..(start + BATCH_ROWS).min(ROWS); - RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int64Array::from_iter_values( - rows.clone().map(|row| row as i64 % 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(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(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 = (ROWS as i64 / 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(), GROUPS as usize)); - expected.push("+------+------+"); - assert_batches_sorted_eq!(expected, &batches); -} diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 0e914ffd81887..39be33587fb90 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -20,18 +20,20 @@ use std::num::NonZeroUsize; use std::sync::{Arc, LazyLock}; -mod count_distinct_spill; #[cfg(feature = "extended_tests")] mod memory_limit_validation; 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; @@ -126,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);