Skip to content

test: skip count distinct spill memory test under forced hash collisions #24918

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rluvaton merged 1 commit into from
Sep 3, 2026
Merged

test: skip count distinct spill memory test under forced hash collisions #24918

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
173 changes: 93 additions & 80 deletions datafusion/core/tests/memory_limit/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,14 +26,11 @@ mod nlj_spill_unmatched;
mod repartition_mem_limit;
mod union_nullable_spill;
mod view_spill_compaction;
use arrow::array::{
ArrayRef, DictionaryArray, Int32Array, Int64Array, RecordBatch, StringViewArray,
};
use arrow::array::{ArrayRef, DictionaryArray, Int32Array, 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;
Expand Down Expand Up @@ -128,83 +125,99 @@ 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();
/// With `force_hash_collisions` every key hashes alike, so the hash
/// repartitioning sends all groups to a single final stage, whose table then
/// does not fit the limit however well memory is released. The limit is sized
/// for the real distribution across four final stages, so this test is skipped
/// under that feature.
#[cfg(not(feature = "force_hash_collisions"))]
mod count_distinct_spill {
use super::*;
use arrow::array::Int64Array;
use datafusion::assert_batches_sorted_eq;

/// `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()
}

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);
/// 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]
Expand Down