From 977bdd183dcb491a056dfafce9bc10514d4cf2b0 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 2 Sep 2026 23:28:06 +0300 Subject: [PATCH 01/21] fix: duplicate groups after spilling in legacy hash aggregation with a nested group key When GroupedHashAggregateStream spills and switches to merging the sorted spill files it relies on GroupOrderingFull, which requires group ids to be assigned in first-seen order. It recreated its group values collector to guarantee that only when there was more than one group column, assuming a single column always uses a sequential single-column collector. A single nested column (Struct, Map) has no specialized single-column collector and is handled by GroupValuesColumn through a row-backed column, whose vectorized interning assigns ids out of input order. The ordering then emitted groups that were still in progress and the next batch reopened them as new groups, so the same key came out more than once with its aggregates split between the rows. Always recreate the collector for the merge phase. --- datafusion/core/tests/memory_limit/mod.rs | 121 +++++++++++++++++- .../src/aggregates/grouped_hash_stream.rs | 11 +- 2 files changed, 124 insertions(+), 8 deletions(-) diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 9172e34a87c79..850e07ee1d81d 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -26,9 +26,13 @@ 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, Int64Builder, ListBuilder, + RecordBatch, StringViewArray, StructArray, +}; +use arrow::buffer::NullBuffer; use arrow::compute::SortOptions; -use arrow::datatypes::{Int32Type, SchemaRef}; +use arrow::datatypes::{Fields, Int32Type, SchemaRef}; use arrow_schema::{DataType, Field, Schema}; use datafusion::assert_batches_eq; use datafusion::config::SpillCompression; @@ -43,6 +47,7 @@ use datafusion::physical_plan::{ExecutionPlan, SendableRecordBatchStream}; use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_catalog::Session; use datafusion_catalog::streaming::StreamingTable; +use datafusion_common::test_util::batches_to_sort_string; use datafusion_common::{Result, assert_contains}; use datafusion_execution::TaskContext; use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; @@ -133,7 +138,6 @@ async fn group_by_hash() { #[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. @@ -220,6 +224,117 @@ mod count_distinct_spill { } } +/// `GROUP BY` on a single nested key in the legacy `GroupedHashAggregateStream` +/// under a memory limit. +/// +/// After spilling, the legacy stream re-aggregates the merged spill files with +/// `GroupOrderingFull`, which requires group ids in first-seen order. A single +/// nested key has no specialized single-column group values implementation and +/// is handled by the multi-column one, whose vectorized interning does not +/// guarantee that order. If the stream keeps that implementation for the merge +/// phase, groups are emitted while still in progress and come out duplicated. +const NESTED_KEY_ROWS: usize = 200_000; +const NESTED_KEY_GROUPS: i64 = 16; +const NESTED_KEY_BATCH_ROWS: usize = 8_192; + +/// Small enough that the final stages must spill their `count(distinct)` +/// state, large enough for the merge of the spilled runs. A `FairSpillPool` +/// caps every stage at its share, so the partial stages emit early and the +/// final stages spill repeatedly, which produces the many small merged batches +/// this bug needs. +const NESTED_KEY_MEMORY_LIMIT: usize = 4 * 1024 * 1024; + +fn nested_key_struct_fields() -> Fields { + Fields::from(vec![ + Field::new("list", DataType::new_list(DataType::Int64, true), true), + Field::new("num", DataType::Int64, true), + ]) +} + +/// `st` is mostly `{list: [g, g + 1], num: g}` for group `g`, with a sprinkle +/// of null lists, empty lists, null nums and null structs so that keys of +/// different shapes meet in the same batches. `v` is unique. +fn nested_key_table() -> MemTable { + let schema = Arc::new(Schema::new(vec![ + Field::new_struct("st", nested_key_struct_fields(), true), + Field::new("v", DataType::Int64, false), + ])); + let batches = (0..NESTED_KEY_ROWS) + .step_by(NESTED_KEY_BATCH_ROWS) + .map(|start| { + let rows = start..(start + NESTED_KEY_BATCH_ROWS).min(NESTED_KEY_ROWS); + let mut list = ListBuilder::new(Int64Builder::new()); + let mut num = Vec::with_capacity(rows.len()); + let mut valid = Vec::with_capacity(rows.len()); + for row in rows.clone() { + let group = row as i64 % NESTED_KEY_GROUPS; + match row % 37 { + 0 => list.append_null(), + 1 => list.append(true), + _ => { + list.values().append_value(group); + list.values().append_value(group + 1); + list.append(true); + } + } + num.push((row % 41 != 0).then_some(group)); + valid.push(row % 43 != 0); + } + let st = StructArray::new( + nested_key_struct_fields(), + vec![Arc::new(list.finish()), Arc::new(Int64Array::from(num))], + Some(NullBuffer::from(valid)), + ); + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(st), + Arc::new(Int64Array::from_iter_values(rows.map(|row| row as i64))), + ], + ) + .unwrap() + }) + .collect(); + MemTable::try_new(schema, vec![batches]).unwrap() +} + +const NESTED_KEY_QUERY: &str = "select st, count(v), count(distinct v), sum(v), avg(v), min(v), max(v) \ + from t group by st"; + +/// Runs the query on the legacy stream, with or without a memory limit. +async fn run_nested_key_query(memory_limit: Option) -> String { + let mut runtime = + RuntimeEnvBuilder::new().with_disk_manager_builder(DiskManagerBuilder::default()); + if let Some(limit) = memory_limit { + runtime = runtime.with_memory_pool(Arc::new(FairSpillPool::new(limit))); + } + let config = SessionConfig::new() + .with_target_partitions(4) + // small batches: the merged spill stream arrives in many batches and + // groups span batch boundaries + .with_batch_size(64) + .set_bool("datafusion.execution.enable_migration_aggregate", false); + let ctx = SessionContext::new_with_config_rt(config, runtime.build_arc().unwrap()); + ctx.register_table("t", Arc::new(nested_key_table())) + .unwrap(); + let batches = ctx + .sql(NESTED_KEY_QUERY) + .await + .unwrap() + .collect() + .await + .unwrap(); + batches_to_sort_string(&batches) +} + +#[tokio::test] +async fn legacy_stream_nested_key_spill_keeps_groups_unique() { + let expected = run_nested_key_query(None).await; + let actual = run_nested_key_query(Some(NESTED_KEY_MEMORY_LIMIT)).await; + // A duplicated group shows up as extra rows with the counts split + assert_eq!(actual, expected); +} + #[tokio::test] async fn join_by_key_multiple_partitions() { let config = SessionConfig::new().with_target_partitions(2); diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index 3f6f3f8ce815b..25dc1086841b8 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -1372,15 +1372,16 @@ impl GroupedHashAggregateStream { // Recreate `group_values` for streaming merge so group ids are assigned // in first-seen order, as required by `GroupOrderingFull`. - // The pre-spill multi-column collector may use `vectorized_intern`, which - // can assign new group ids out of input order under hash collisions. + // The pre-spill collector may use `vectorized_intern`, which can assign + // new group ids out of input order under hash collisions. That is the + // multi-column collector, which also serves a single group column + // whose type has no specialized single-column collector (for example + // `Struct` or `Map`), so recreate unconditionally. let group_schema = self .spill_state .merging_group_by .group_schema(&self.spill_state.spill_schema)?; - if group_schema.fields().len() > 1 { - self.group_values = new_group_values(group_schema, &self.group_ordering)?; - } + self.group_values = new_group_values(group_schema, &self.group_ordering)?; // Use `OutOfMemoryMode::ReportError` from this point on // to ensure we don't spill the spilled data to disk again. From 1bce7180b97dd517f0303e9d57b9fb347136cf6e Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:30:30 +0300 Subject: [PATCH 02/21] test: add aggregate fuzz tests for all shapes --- .../tests/fuzz_cases/aggregate_chain_fuzz.rs | 1387 +++++++++++++++++ .../aggregation_fuzzer/AGGREGATE_CHAINS.md | 323 ++++ datafusion/core/tests/fuzz_cases/mod.rs | 1 + 3 files changed, 1711 insertions(+) create mode 100644 datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs create mode 100644 datafusion/core/tests/fuzz_cases/aggregation_fuzzer/AGGREGATE_CHAINS.md diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs new file mode 100644 index 0000000000000..3c37780793b63 --- /dev/null +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -0,0 +1,1387 @@ +// 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. + +//! Fuzz test that runs every valid `AggregateExec` chain over the same data +//! and asserts identical results. +//! +//! See `aggregation_fuzzer/AGGREGATE_CHAINS.md` for the chain catalogue. Each +//! chain there is a `Shape` below plus a source `Order`. Cases are generated +//! as `shapes × orders × migration flag × cardinality × memory`. + +use std::num::NonZeroUsize; +use std::sync::Arc; +use std::time::Duration; + +use arrow::array::{ + BooleanArray, Int64Array, Int64Builder, ListBuilder, RecordBatch, StringArray, + StringViewArray, StructArray, UInt32Array, +}; +use arrow::buffer::NullBuffer; +use arrow::compute::{SortColumn, lexsort_to_indices, take_record_batch}; +use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef, SortOptions}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +use datafusion::prelude::SessionConfig; +use datafusion_common::Result; +use datafusion_common::test_util::batches_to_sort_string; +use datafusion_common_runtime::JoinSet; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{FairSpillPool, TrackConsumersPool}; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_functions_aggregate::average::avg_udaf; +use datafusion_functions_aggregate::count::count_udaf; +use datafusion_functions_aggregate::min_max::{max_udaf, min_udaf}; +use datafusion_functions_aggregate::sum::sum_udaf; +use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr}; +use datafusion_physical_expr::expressions::{cast, col}; +use datafusion_physical_expr::{ + LexOrdering, Partitioning, PhysicalExpr, PhysicalSortExpr, +}; +use datafusion_physical_plan::aggregates::{ + AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy, +}; +use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion_physical_plan::repartition::RepartitionExec; +use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion_physical_plan::{ExecutionPlan, InputOrderMode, collect, displayable}; +use rand::rngs::StdRng; +use rand::seq::SliceRandom; +use rand::{Rng, SeedableRng}; + +use AggregateMode::*; +use Operator::*; + +// --------------------------------------------------------------------------- +// Case space +// --------------------------------------------------------------------------- + +const ROWS: usize = 64 * 1024; +const PARTITIONS: usize = 4; +const BATCH_SIZE: usize = 64; +/// The fair pool caps every spillable consumer at `pool / consumers`, and a +/// chain registers up to twenty consumers (aggregate streams plus one per +/// repartition channel). The cap has to clear a small table's legitimate +/// footprint, which at very low cardinality is dominated by the `count +/// distinct` sets and grows in steps of roughly 100 KB, while a final table at +/// very high cardinality must still exceed it. +const LIMITED_POOL_BYTES: usize = 4 * 1024 * 1024; + +/// How the source data is ordered relative to the group keys `(k1, k2)`. +#[derive(Clone, Copy, Debug, PartialEq)] +enum Order { + /// Not ordered. Aggregates see `InputOrderMode::Linear`. + Unordered, + /// Sorted by `k1` only. Aggregates see `InputOrderMode::PartiallySorted([0])`. + SortedByFirstKey, + /// Sorted by `k1, k2`. Aggregates see `InputOrderMode::Sorted`. + SortedByAllKeys, +} + +#[derive(Clone, Copy, Debug, PartialEq)] +enum Cardinality { + VeryHigh, + High, + Medium, + Low, + VeryLow, +} + +impl Cardinality { + const ALL: [Self; 5] = [ + Self::VeryHigh, + Self::High, + Self::Medium, + Self::Low, + Self::VeryLow, + ]; + + /// Number of distinct `(k1, k2)` groups. + fn groups(self) -> usize { + match self { + Self::VeryHigh => ROWS, + Self::High => ROWS / 2, + Self::Medium => ROWS / 32, + Self::Low => 16, + Self::VeryLow => 2, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq)] +enum Memory { + /// Unlimited pool. Nothing spills or emits early. + Unlimited, + /// Pool sized so final and single hash tables cannot fit. + Limited, +} + +/// One operator in a chain, listed bottom to top. +#[derive(Clone, Copy, Debug)] +enum Operator { + Aggregate(AggregateMode), + /// `AggregateExec` with `limit_options` set, which selects + /// `GroupedTopKAggregateStream` regardless of mode. + TopK(AggregateMode), + /// `RepartitionExec` hashed on the group keys. Destroys ordering. + HashRepartition, + /// `RepartitionExec` hashed on the group keys with `preserve_order`. + OrderPreservingHashRepartition, + /// `CoalescePartitionsExec`. Destroys ordering. + CoalescePartitions, + /// `SortPreservingMergeExec` on the current ordering. + SortPreservingMerge, +} + +/// The logical query a chain computes. +#[derive(Clone, Copy, Debug, PartialEq)] +enum Query { + /// `GROUP BY k1, k2` with count, count distinct, sum, avg, min, max. + /// Two primitive keys, handled by `GroupValuesColumn`. + Grouped, + /// The same aggregates without `GROUP BY`. + NoGrouping, + /// `GROUP BY k1` with `max(v)` only, the shape the TopK stream supports. + /// Chains using `Operator::TopK` set a limit larger than any possible + /// group count, so the result must still be the complete aggregate. + TopK, + /// `GROUP BY b` (Boolean), handled by `GroupValuesBoolean`. + BooleanKey, + /// `GROUP BY s` (Utf8), handled by `GroupValuesBytes`. + BytesKey, + /// `GROUP BY sv` (Utf8View), handled by `GroupValuesBytesView`. + BytesViewKey, + /// `GROUP BY p` (Int64 with as many distinct values as groups), handled + /// by `GroupValuesPrimitive`. + PrimitiveKey, + /// `GROUP BY b, s, sv, p`, handled by `GroupValuesColumn` with mixed + /// column types. + MixedKeys, + /// `GROUP BY st` (Struct of a List and an Int64), which no + /// specialized implementation supports, so it falls back to the row format + /// `GroupValuesRows`. + StructKey, +} + +impl Query { + /// Group key columns, in `GROUP BY` order. + fn keys(self) -> &'static [&'static str] { + match self { + Query::Grouped => &["k1", "k2"], + Query::NoGrouping => &[], + Query::TopK => &["k1"], + Query::BooleanKey => &["b"], + Query::BytesKey => &["s"], + Query::BytesViewKey => &["sv"], + Query::PrimitiveKey => &["p"], + Query::MixedKeys => &["b", "s", "sv", "p"], + Query::StructKey => &["st"], + } + } + + /// Whether the source can be sorted by the keys. Struct columns cannot be + /// sorted by the arrow sort kernels, so that query only runs unordered. + fn sortable(self) -> bool { + self != Query::StructKey + } +} + +/// Larger than any possible number of groups, so TopK keeps every group. +const TOP_K_LIMIT: usize = 2 * ROWS; + +/// A plan shape. +#[derive(Debug)] +struct Shape { + name: &'static str, + operators: &'static [Operator], + /// Source partition count. + source_partitions: usize, + query: Query, +} + +const fn shape( + name: &'static str, + operators: &'static [Operator], + source_partitions: usize, + query: Query, +) -> Shape { + Shape { + name, + operators, + source_partitions, + query, + } +} + +/// Every shape from AGGREGATE_CHAINS.md. The ordering variants there come from +/// crossing a shape with `Order`, so one entry here covers several rows. +const SHAPES: &[Shape] = &[ + shape("single", &[Aggregate(Single)], 1, Query::Grouped), + shape( + "single_partitioned", + &[HashRepartition, Aggregate(SinglePartitioned)], + PARTITIONS, + Query::Grouped, + ), + shape( + "single_partitioned_order_preserving", + &[OrderPreservingHashRepartition, Aggregate(SinglePartitioned)], + PARTITIONS, + Query::Grouped, + ), + shape( + "partial_repartition_final", + &[ + Aggregate(Partial), + HashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + Query::Grouped, + ), + shape( + "partial_coalesce_final", + &[Aggregate(Partial), CoalescePartitions, Aggregate(Final)], + PARTITIONS, + Query::Grouped, + ), + shape( + "partial_order_preserving_repartition_final", + &[ + Aggregate(Partial), + OrderPreservingHashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + Query::Grouped, + ), + shape( + "partial_sort_preserving_merge_final", + &[Aggregate(Partial), SortPreservingMerge, Aggregate(Final)], + PARTITIONS, + Query::Grouped, + ), + shape( + "partial_final_single_partition", + &[Aggregate(Partial), Aggregate(Final)], + 1, + Query::Grouped, + ), + shape( + "partial_repartition_reduce_repartition_final", + &[ + Aggregate(Partial), + HashRepartition, + Aggregate(PartialReduce), + HashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + Query::Grouped, + ), + shape( + "partial_repartition_reduce_coalesce_final", + &[ + Aggregate(Partial), + HashRepartition, + Aggregate(PartialReduce), + CoalescePartitions, + Aggregate(Final), + ], + PARTITIONS, + Query::Grouped, + ), + shape( + "partial_local_reduce_repartition_final", + &[ + Aggregate(Partial), + Aggregate(PartialReduce), + HashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + Query::Grouped, + ), + // ordered PartialReduce has no dedicated stream, lands on the fallback + shape( + "partial_reduce_final_order_preserving", + &[ + Aggregate(Partial), + OrderPreservingHashRepartition, + Aggregate(PartialReduce), + OrderPreservingHashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + Query::Grouped, + ), + shape( + "no_grouping_single", + &[Aggregate(Single)], + 1, + Query::NoGrouping, + ), + // TopK: same query without a limit is the reference for the TopK chains + shape("top_k_query_single", &[Aggregate(Single)], 1, Query::TopK), + shape("top_k_single", &[TopK(Single)], 1, Query::TopK), + // planner shape: the limit lands on the aggregate under the sort + shape( + "top_k_partial_repartition_final", + &[Aggregate(Partial), HashRepartition, TopK(FinalPartitioned)], + PARTITIONS, + Query::TopK, + ), + shape( + "top_k_partial_coalesce_final", + &[Aggregate(Partial), CoalescePartitions, TopK(Final)], + PARTITIONS, + Query::TopK, + ), + shape( + "top_k_both_stages", + &[TopK(Partial), HashRepartition, TopK(FinalPartitioned)], + PARTITIONS, + Query::TopK, + ), + // group key types: single stage and the default two-stage planner shape + shape( + "boolean_key_single", + &[Aggregate(Single)], + 1, + Query::BooleanKey, + ), + shape( + "boolean_key_partial_repartition_final", + &[ + Aggregate(Partial), + HashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + Query::BooleanKey, + ), + shape("bytes_key_single", &[Aggregate(Single)], 1, Query::BytesKey), + shape( + "bytes_key_partial_repartition_final", + &[ + Aggregate(Partial), + HashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + Query::BytesKey, + ), + shape( + "bytes_view_key_single", + &[Aggregate(Single)], + 1, + Query::BytesViewKey, + ), + shape( + "bytes_view_key_partial_repartition_final", + &[ + Aggregate(Partial), + HashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + Query::BytesViewKey, + ), + shape( + "primitive_key_single", + &[Aggregate(Single)], + 1, + Query::PrimitiveKey, + ), + shape( + "primitive_key_partial_repartition_final", + &[ + Aggregate(Partial), + HashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + Query::PrimitiveKey, + ), + shape( + "mixed_keys_single", + &[Aggregate(Single)], + 1, + Query::MixedKeys, + ), + shape( + "mixed_keys_partial_repartition_final", + &[ + Aggregate(Partial), + HashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + Query::MixedKeys, + ), + // ordered multi-column group values + shape( + "mixed_keys_partial_order_preserving_repartition_final", + &[ + Aggregate(Partial), + OrderPreservingHashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + Query::MixedKeys, + ), + shape( + "struct_key_single", + &[Aggregate(Single)], + 1, + Query::StructKey, + ), + shape( + "struct_key_partial_repartition_final", + &[ + Aggregate(Partial), + HashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + Query::StructKey, + ), + shape( + "no_grouping_partial_coalesce_final", + &[Aggregate(Partial), CoalescePartitions, Aggregate(Final)], + PARTITIONS, + Query::NoGrouping, + ), + shape( + "no_grouping_partial_reduce_final", + &[ + Aggregate(Partial), + CoalescePartitions, + Aggregate(PartialReduce), + CoalescePartitions, + Aggregate(Final), + ], + PARTITIONS, + Query::NoGrouping, + ), + shape( + "no_grouping_partial_final_single_partition", + &[Aggregate(Partial), Aggregate(Final)], + 1, + Query::NoGrouping, + ), +]; + +fn shape_by_name(name: &str) -> &'static Shape { + SHAPES.iter().find(|shape| shape.name == name).unwrap() +} + +impl Shape { + /// Source orders that make sense for this shape. Order-preserving shuffles + /// need an ordering to preserve; no-grouping chains ignore ordering. + fn orders(&self) -> Vec { + let needs_ordered_input = self.operators.iter().any(|operator| { + matches!( + operator, + OrderPreservingHashRepartition | SortPreservingMerge + ) + }); + let keys = self.query.keys(); + let mut orders = vec![]; + if !needs_ordered_input { + orders.push(Order::Unordered); + } + if self.query.sortable() && !keys.is_empty() { + // With a single key, sorting by the first key is already sorting + // by all keys. + if keys.len() > 1 { + orders.push(Order::SortedByFirstKey); + } + orders.push(Order::SortedByAllKeys); + } + orders + } +} + +#[derive(Clone, Debug)] +struct Case { + shape: &'static Shape, + order: Order, + migration_enabled: bool, + cardinality: Cardinality, + memory: Memory, + /// Whether the skip-partial probe may fire. Only varied for shapes with a + /// grouped `Partial` stage on Linear input, since nothing else runs the + /// probe. + skip_partial_enabled: bool, +} + +impl Shape { + /// Whether some `Partial` stage of this shape runs the skip-partial probe + /// for the given source order: grouped, not TopK, and Linear input. + fn has_skip_partial_candidate(&self, order: Order) -> bool { + if self.query == Query::NoGrouping { + return false; + } + let mut current = order; + for operator in self.operators { + match operator { + HashRepartition | CoalescePartitions => current = Order::Unordered, + Aggregate(Partial) if current == Order::Unordered => return true, + _ => {} + } + } + false + } +} + +fn all_cases() -> Vec { + // `AGGREGATE_CHAIN_SHAPES=a,b` restricts the run to shapes whose name + // contains one of the given substrings, to reproduce or bisect quickly. + let shape_filter: Vec = std::env::var("AGGREGATE_CHAIN_SHAPES") + .map(|value| value.split(',').map(str::to_string).collect()) + .unwrap_or_default(); + let mut cases = vec![]; + for shape in SHAPES.iter().filter(|shape| { + shape_filter.is_empty() + || shape_filter + .iter() + .any(|needle| shape.name.contains(needle)) + }) { + for order in shape.orders() { + let skip_partial_variants: &[bool] = + if shape.has_skip_partial_candidate(order) { + &[true, false] + } else { + &[true] + }; + for migration_enabled in [true, false] { + for cardinality in Cardinality::ALL { + for memory in [Memory::Unlimited, Memory::Limited] { + for &skip_partial_enabled in skip_partial_variants { + cases.push(Case { + shape, + order, + migration_enabled, + cardinality, + memory, + skip_partial_enabled, + }); + } + } + } + } + } + } + cases +} + +// --------------------------------------------------------------------------- +// Data generation +// --------------------------------------------------------------------------- + +/// `k1 Int64 nullable, k2 Int64 nullable, v Int64` +/// `k1, k2 Int64` (two-key query), `v Int64` (aggregated), and one column per +/// key type: `b Boolean`, `s Utf8`, `sv Utf8View`, `p Int64`, and +/// `st Struct, num: Int64>`. +/// Every key column is nullable. +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("k1", DataType::Int64, true), + Field::new("k2", DataType::Int64, true), + Field::new("v", DataType::Int64, false), + Field::new("b", DataType::Boolean, true), + Field::new("s", DataType::Utf8, true), + Field::new("sv", DataType::Utf8View, true), + Field::new("p", DataType::Int64, true), + Field::new_struct("st", struct_fields(), true), + ])) +} + +/// About 3% nulls. +fn not_null(rng: &mut StdRng) -> bool { + rng.random_range(0..100) >= 3 +} + +fn struct_fields() -> Fields { + Fields::from(vec![ + Field::new("list", DataType::new_list(DataType::Int64, true), true), + Field::new("num", DataType::Int64, true), + ]) +} + +/// The raw rows for one cardinality, deterministic per seed. The same multiset +/// is used for every `Order` and `Shape` so results are comparable. +/// +/// Requirements: +/// - exactly `ROWS` rows +/// - `cardinality.groups()` distinct `(k1, k2)` pairs, spread so that `k1` +/// alone has fewer distinct values than `(k1, k2)`. Otherwise +/// `SortedByFirstKey` degenerates into `SortedByAllKeys`. +/// - some nulls in `k1` and `k2` +fn generate_rows(cardinality: Cardinality, seed: u64) -> RecordBatch { + let mut rng = StdRng::seed_from_u64(seed); + let groups = cardinality.groups(); + // `k2` cycles through at most sqrt(groups) values, so `k1` alone has fewer + // distinct values than the `(k1, k2)` pair. + let k2_values = (groups as f64).sqrt().ceil().max(2.0) as i64; + + let mut k1 = Vec::with_capacity(ROWS); + let mut k2 = Vec::with_capacity(ROWS); + let mut v = Vec::with_capacity(ROWS); + let mut b = Vec::with_capacity(ROWS); + let mut s = Vec::with_capacity(ROWS); + let mut sv = Vec::with_capacity(ROWS); + let mut p = Vec::with_capacity(ROWS); + let mut st_list = ListBuilder::new(Int64Builder::new()); + let mut st_num = Vec::with_capacity(ROWS); + let mut st_valid = Vec::with_capacity(ROWS); + for row in 0..ROWS { + let group = (row % groups) as i64; + k1.push(not_null(&mut rng).then_some(group / k2_values)); + k2.push(not_null(&mut rng).then_some(group % k2_values)); + v.push(rng.random_range(-1_000i64..1_000)); + // every key-type column has `groups` distinct values (boolean: two) + b.push(not_null(&mut rng).then_some(group % 2 == 0)); + s.push(not_null(&mut rng).then(|| format!("s{group:06}"))); + sv.push(not_null(&mut rng).then(|| format!("sv{group:06}"))); + p.push(not_null(&mut rng).then_some(group)); + // struct { list: [group, group + 1], [] or null; num: group or null } + st_valid.push(not_null(&mut rng)); + match rng.random_range(0..100) { + 0..3 => st_list.append_null(), + 3..6 => st_list.append(true), + _ => { + st_list.values().append_value(group); + st_list.values().append_value(group + 1); + st_list.append(true); + } + } + st_num.push(not_null(&mut rng).then_some(group)); + } + let st = StructArray::try_new( + struct_fields(), + vec![ + Arc::new(st_list.finish()), + Arc::new(Int64Array::from(st_num)), + ], + Some(NullBuffer::from(st_valid)), + ) + .unwrap(); + + RecordBatch::try_new( + schema(), + vec![ + Arc::new(Int64Array::from(k1)), + Arc::new(Int64Array::from(k2)), + Arc::new(Int64Array::from(v)), + Arc::new(BooleanArray::from(b)), + Arc::new(StringArray::from(s)), + Arc::new(StringViewArray::from(sv)), + Arc::new(Int64Array::from(p)), + Arc::new(st), + ], + ) + .unwrap() +} + +/// Arrange `rows` for the given `order` and split into `partitions` partitions +/// of `BATCH_SIZE` batches. +/// +/// - `Unordered`: shuffle rows, round-robin into partitions +/// - `SortedByFirstKey`: sort by `k1` (nulls first), contiguous slice per partition +/// - `SortedByAllKeys`: sort by `k1, k2` (nulls first), contiguous slice per partition +/// +/// Every partition individually satisfies the ordering. +fn arrange( + rows: &RecordBatch, + query: Query, + order: Order, + partitions: usize, +) -> Vec> { + let schema = rows.schema(); + let per_partition: Vec = match source_ordering(&schema, query, order) { + None => { + let mut permutation: Vec = (0..rows.num_rows() as u32).collect(); + permutation.shuffle(&mut StdRng::seed_from_u64(0)); + let shuffled = + take_record_batch(rows, &UInt32Array::from(permutation)).unwrap(); + (0..partitions) + .map(|partition| { + let indices: UInt32Array = (partition as u32 + ..shuffled.num_rows() as u32) + .step_by(partitions) + .collect(); + take_record_batch(&shuffled, &indices).unwrap() + }) + .collect() + } + Some(ordering) => { + let sort_columns: Vec = ordering + .iter() + .map(|sort_expr| SortColumn { + values: sort_expr + .expr + .evaluate(rows) + .unwrap() + .into_array(rows.num_rows()) + .unwrap(), + options: Some(sort_expr.options), + }) + .collect(); + let indices = lexsort_to_indices(&sort_columns, None).unwrap(); + let sorted = take_record_batch(rows, &indices).unwrap(); + let per_partition = sorted.num_rows().div_ceil(partitions); + (0..partitions) + .map(|partition| { + let start = (partition * per_partition).min(sorted.num_rows()); + let length = per_partition.min(sorted.num_rows() - start); + sorted.slice(start, length) + }) + .collect() + } + }; + + // Copy every batch into its own buffers, as a real scan would produce. + // A slice shares the whole partition's buffers, and operators that + // account batches by `get_array_memory_size` (RepartitionExec, the merge) + // would charge every 64-row batch the size of the entire partition. + per_partition + .iter() + .map(|partition| { + (0..partition.num_rows()) + .step_by(BATCH_SIZE) + .map(|start| { + let slice = partition + .slice(start, BATCH_SIZE.min(partition.num_rows() - start)); + { + // `take` copies; `concat_batches` of one batch only slices + let indices = + UInt32Array::from_iter_values(0..slice.num_rows() as u32); + take_record_batch(&slice, &indices).unwrap() + } + }) + .collect() + }) + .collect() +} + +// --------------------------------------------------------------------------- +// Plan construction +// --------------------------------------------------------------------------- + +fn sort_expr(schema: &Schema, column: &str) -> PhysicalSortExpr { + PhysicalSortExpr::new( + col(column, schema).unwrap(), + SortOptions { + descending: false, + nulls_first: true, + }, + ) +} + +/// The ordering the source declares for `order`. +fn source_ordering(schema: &Schema, query: Query, order: Order) -> Option { + let keys = query.keys(); + let sort_columns: &[&str] = match order { + Order::Unordered => return None, + Order::SortedByFirstKey => &keys[..1], + Order::SortedByAllKeys => keys, + }; + LexOrdering::new(sort_columns.iter().map(|column| sort_expr(schema, column))) +} + +fn source( + partitions: &[Vec], + query: Query, + order: Order, +) -> Arc { + let schema = schema(); + let mut memory_source = + MemorySourceConfig::try_new(partitions, Arc::clone(&schema), None).unwrap(); + if let Some(ordering) = source_ordering(&schema, query, order) { + memory_source = memory_source + .try_with_sort_information(vec![ordering]) + .unwrap(); + } + DataSourceExec::from_data_source(memory_source) +} + +fn group_by(schema: &Schema, query: Query) -> PhysicalGroupBy { + PhysicalGroupBy::new_single( + query + .keys() + .iter() + .map(|key| (col(key, schema).unwrap(), key.to_string())) + .collect(), + ) +} + +/// Aggregates with non-trivial partial state so the Partial, PartialReduce and +/// Final stages are actually exercised. `avg` (two-field state) and +/// `count distinct` (set state) matter most. +fn aggregates(schema: &SchemaRef, query: Query) -> Vec> { + let value_column = || vec![col("v", schema).unwrap()]; + let build = |builder: AggregateExprBuilder, alias: &str| { + Arc::new( + builder + .schema(Arc::clone(schema)) + .alias(alias) + .build() + .unwrap(), + ) + }; + if query == Query::TopK { + // TopK supports exactly one min/max aggregate over a non-nullable input + return vec![build( + AggregateExprBuilder::new(max_udaf(), value_column()), + "max", + )]; + } + vec![ + build( + AggregateExprBuilder::new(count_udaf(), value_column()), + "count", + ), + build( + AggregateExprBuilder::new(count_udaf(), value_column()).distinct(), + "count_distinct", + ), + build(AggregateExprBuilder::new(sum_udaf(), value_column()), "sum"), + // avg has no Int64 groups accumulator; the values are small integers so + // the Float64 sum stays exact and the result is order-independent. + build( + AggregateExprBuilder::new( + avg_udaf(), + vec![cast(col("v", schema).unwrap(), schema, DataType::Float64).unwrap()], + ), + "avg", + ), + build(AggregateExprBuilder::new(min_udaf(), value_column()), "min"), + build(AggregateExprBuilder::new(max_udaf(), value_column()), "max"), + ] +} + +/// Folds `shape.operators` bottom-up into a plan. The group-by, aggregate +/// expressions and hash keys are rewritten after every aggregate stage so the +/// next stage consumes that stage's output. +fn build_plan(shape: &Shape, input: Arc) -> Arc { + let input_schema = schema(); + let mut plan = input; + let mut group_by = group_by(&input_schema, shape.query); + let mut aggregates = aggregates(&input_schema, shape.query); + let mut hash_keys: Vec> = group_by.input_exprs(); + + for operator in shape.operators { + plan = match operator { + Aggregate(mode) | TopK(mode) => { + let limit_options = matches!(operator, TopK(_)) + .then(|| LimitOptions::new_with_order(TOP_K_LIMIT, true)); + let aggregate = Arc::new( + AggregateExec::try_new( + *mode, + group_by.clone(), + aggregates.clone(), + vec![None; aggregates.len()], + plan, + Arc::clone(&input_schema), + ) + .unwrap() + .with_limit_options(limit_options), + ); + group_by = aggregate.group_expr().as_final(); + aggregates = aggregate.aggr_expr().to_vec(); + hash_keys = aggregate.output_group_expr(); + aggregate + } + HashRepartition => Arc::new( + RepartitionExec::try_new( + plan, + Partitioning::Hash(hash_keys.clone(), PARTITIONS), + ) + .unwrap(), + ), + OrderPreservingHashRepartition => Arc::new( + RepartitionExec::try_new( + plan, + Partitioning::Hash(hash_keys.clone(), PARTITIONS), + ) + .unwrap() + .with_preserve_order(), + ), + CoalescePartitions => Arc::new(CoalescePartitionsExec::new(plan)), + SortPreservingMerge => { + let ordering = plan.properties().output_ordering().cloned().unwrap(); + Arc::new(SortPreservingMergeExec::new(ordering, plan)) + } + }; + } + plan +} + +// --------------------------------------------------------------------------- +// Execution context +// --------------------------------------------------------------------------- + +fn task_context(case: &Case) -> Arc { + let config = SessionConfig::new() + .with_batch_size(BATCH_SIZE) + .with_target_partitions(PARTITIONS) + .set_bool( + "datafusion.execution.enable_migration_aggregate", + case.migration_enabled, + ) + // The default is 100k rows. Lower it so the skip-partial probe can + // fire on our per-partition row counts. A ratio threshold of 1.0 + // disables the probe entirely. + .set_usize( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + 1024, + ); + let mut config = config; + config + .options_mut() + .execution + .skip_partial_aggregation_probe_ratio_threshold = + if case.skip_partial_enabled { 0.8 } else { 1.0 }; + + let runtime = match case.memory { + Memory::Unlimited => RuntimeEnvBuilder::new(), + // Small enough that a very-high-cardinality final table spills, large + // enough that the legacy stream can still reserve its sort headroom + // and that RepartitionExec / SortPreservingMergeExec succeed. The + // fair pool keeps one stage from starving the others. + Memory::Limited => { + RuntimeEnvBuilder::new().with_memory_pool(Arc::new(TrackConsumersPool::new( + FairSpillPool::new(LIMITED_POOL_BYTES), + NonZeroUsize::new(5).unwrap(), + ))) + } + } + .build_arc() + .unwrap(); + + Arc::new( + TaskContext::default() + .with_session_config(config) + .with_runtime(runtime), + ) +} + +// --------------------------------------------------------------------------- +// Assertions +// --------------------------------------------------------------------------- + +/// All `AggregateExec` nodes in the plan, bottom-up. +fn aggregate_nodes(plan: &Arc) -> Vec> { + let mut nodes = vec![]; + let mut node = Arc::clone(plan); + loop { + if node.downcast_ref::().is_some() { + nodes.push(Arc::clone(&node)); + } + match node.children().first() { + Some(child) => node = Arc::clone(child), + None => break, + } + } + nodes.reverse(); + nodes +} + +fn as_aggregate(node: &Arc) -> &AggregateExec { + node.downcast_ref::().unwrap() +} + +/// Expected source order seen by each aggregate stage, bottom-up. Ordering is +/// lost at `HashRepartition` and `CoalescePartitions`, and kept by the +/// order-preserving shuffles and by aggregate stages themselves. +fn expected_orders(shape: &Shape, source_order: Order) -> Vec { + let mut current = source_order; + let mut expected = vec![]; + for operator in shape.operators { + match operator { + HashRepartition | CoalescePartitions => current = Order::Unordered, + Aggregate(_) | TopK(_) => expected.push(current), + OrderPreservingHashRepartition | SortPreservingMerge => {} + } + } + expected +} + +fn order_matches(query: Query, expected: Order, actual: &InputOrderMode) -> bool { + // With a single group key, sorting by the first key already covers every + // group key. + let single_key = query.keys().len() == 1; + match (expected, actual) { + (Order::Unordered, InputOrderMode::Linear) => true, + (Order::SortedByFirstKey, InputOrderMode::PartiallySorted(indices)) => { + !single_key && indices == &[0] + } + (Order::SortedByFirstKey, InputOrderMode::Sorted) => single_key, + (Order::SortedByAllKeys, InputOrderMode::Sorted) => true, + _ => false, + } +} + +/// Whether this stage's stream is allowed to spill. See the memory table in +/// AGGREGATE_CHAINS.md. +fn can_spill(case: &Case, aggregate: &AggregateExec) -> bool { + if aggregate.limit_options().is_some() { + // GroupedTopKAggregateStream keeps a bounded heap and never spills + return false; + } + let spilling_mode = match aggregate.mode() { + Final | FinalPartitioned | Single | SinglePartitioned => true, + // The dedicated PartialReduce stream emits early and only exists for + // Linear input; ordered input or migration off run the legacy stream, + // which spills. + PartialReduce => { + !case.migration_enabled + || *aggregate.input_order_mode() != InputOrderMode::Linear + } + Partial => false, + }; + let has_groups = !aggregate.group_expr().is_empty(); + spilling_mode && has_groups && *aggregate.input_order_mode() != InputOrderMode::Sorted +} + +/// Whether this stage runs the skip-partial probe. +fn runs_skip_partial_probe(aggregate: &AggregateExec) -> bool { + *aggregate.mode() == Partial + && aggregate.limit_options().is_none() + && !aggregate.group_expr().is_empty() + && *aggregate.input_order_mode() == InputOrderMode::Linear +} + +fn check_plan_shape(case: &Case, plan: &Arc) { + if case.shape.query == Query::NoGrouping { + return; + } + let nodes = aggregate_nodes(plan); + let expected = expected_orders(case.shape, case.order); + assert_eq!(nodes.len(), expected.len(), "{case:?}"); + for (node, expected_order) in nodes.iter().zip(expected) { + let aggregate = as_aggregate(node); + assert!( + order_matches( + case.shape.query, + expected_order, + aggregate.input_order_mode() + ), + "{case:?}: expected {expected_order:?} got {:?}\n{}", + aggregate.input_order_mode(), + displayable(plan.as_ref()).indent(true) + ); + } +} + +/// Returns a description of every stage that spilled, bottom-up, such as +/// `Final(Linear)`. +fn check_metrics(case: &Case, plan: &Arc) -> Vec { + let mut spilled = vec![]; + for node in aggregate_nodes(plan) { + let aggregate = as_aggregate(&node); + let mode = aggregate.mode(); + let metrics = node.metrics().unwrap(); + let spill_count = metrics.spill_count().unwrap_or(0); + if spill_count > 0 { + spilled.push(format!("{mode:?}({:?})", aggregate.input_order_mode())); + } + let skipped_rows = metrics + .sum_by_name("skipped_aggregation_rows") + .map(|metric| metric.as_usize()) + .unwrap_or(0); + + match case.memory { + Memory::Unlimited => { + assert_eq!(spill_count, 0, "{case:?}: unexpected spill in {mode:?}"); + } + Memory::Limited => { + // Whether a spilling-capable stage actually spills depends on + // the pool geometry, so only the run-wide coverage check in the + // driver requires it. Streams that cannot spill must not. + if !can_spill(case, aggregate) { + assert_eq!(spill_count, 0, "{case:?}: {mode:?} must never spill"); + } + } + } + + // Only the two-key query has as many groups as `cardinality` says; + // the TopK query groups by `k1` alone and stays far below the ratio. + if case.memory == Memory::Unlimited + && case.cardinality == Cardinality::VeryHigh + && case.shape.query == Query::Grouped + && case.skip_partial_enabled + && runs_skip_partial_probe(aggregate) + { + assert!( + skipped_rows > 0, + "{case:?}: skip-partial probe did not fire" + ); + } + if !case.skip_partial_enabled || !runs_skip_partial_probe(aggregate) { + assert_eq!(skipped_rows, 0, "{case:?}: skip-partial fired in {mode:?}"); + } + } + spilled +} + +// --------------------------------------------------------------------------- +// Driver +// --------------------------------------------------------------------------- + +/// Sorted output plus the stages that spilled, empty if none did. +struct Outcome { + output: String, + spilled: Vec, +} + +/// Runs one case, checks plan shape and metrics, and returns its outcome. +/// +/// Running out of memory is never accepted, not even under the limited pool: +/// every stream either spills, emits early, or is bounded, so an error there +/// is a bug in a stream's memory handling or in how stages share the pool. +async fn run_case(case: Case, rows: Arc) -> Outcome { + log::debug!("start {case:?}"); + let outcome = run_case_inner(&case, rows).await; + log::debug!("done {case:?}"); + outcome +} + +async fn run_case_inner(case: &Case, rows: Arc) -> Outcome { + let partitions = arrange( + &rows, + case.shape.query, + case.order, + case.shape.source_partitions, + ); + let plan = build_plan( + case.shape, + source(&partitions, case.shape.query, case.order), + ); + check_plan_shape(case, &plan); + + // A hang is a failure too: name the case instead of stalling the run. + let collected = tokio::time::timeout( + Duration::from_secs(CASE_TIMEOUT_SECS), + collect(Arc::clone(&plan), task_context(case)), + ) + .await + .unwrap_or_else(|_| { + panic!( + "{case:?} did not finish within {CASE_TIMEOUT_SECS}s\n{}", + displayable(plan.as_ref()).indent(true) + ) + }); + let batches = match collected { + Ok(batches) => batches, + Err(error) => panic!( + "{case:?} failed: {error}\n{}", + displayable(plan.as_ref()).indent(true) + ), + }; + let spilled = check_metrics(case, &plan); + Outcome { + output: batches_to_sort_string(&batches), + spilled, + } +} + +/// Reference result: the single-stage shape of `query` without a limit, one +/// partition, unordered input, unlimited memory and migration on. +async fn reference( + query: Query, + rows: Arc, + cardinality: Cardinality, +) -> String { + let shape = SHAPES + .iter() + .find(|shape| { + shape.query == query && matches!(shape.operators, [Aggregate(Single)]) + }) + .unwrap(); + let outcome = run_case( + Case { + shape, + order: Order::Unordered, + migration_enabled: true, + cardinality, + memory: Memory::Unlimited, + skip_partial_enabled: true, + }, + rows, + ) + .await; + outcome.output +} + +#[tokio::test(flavor = "multi_thread")] +async fn aggregate_chain_fuzz() { + let seed = rand::random::(); + println!("aggregate_chain_fuzz seed = {seed}"); + let mut total_spilled = 0; + let mut failures: Vec = vec![]; + + for cardinality in Cardinality::ALL { + let rows = Arc::new(generate_rows(cardinality, seed)); + let mut expected_by_query: Vec<(Query, String)> = Vec::new(); + for shape in SHAPES { + if expected_by_query + .iter() + .any(|(query, _)| *query == shape.query) + { + continue; + } + let expected = reference(shape.query, Arc::clone(&rows), cardinality).await; + expected_by_query.push((shape.query, expected)); + } + + let mut join_set = JoinSet::new(); + let (mut spilled, mut finished) = (vec![], vec![]); + for case in all_cases() + .into_iter() + .filter(|case| case.cardinality == cardinality) + { + let rows = Arc::clone(&rows); + let expected = expected_by_query + .iter() + .find(|(query, _)| *query == case.shape.query) + .map(|(_, expected)| expected.clone()) + .unwrap(); + // Every in-flight case holds several copies of the dataset, so + // bound the concurrency instead of spawning the whole matrix. + while join_set.len() >= MAX_CONCURRENT_CASES { + collect_finished( + &mut join_set, + &mut spilled, + &mut finished, + &mut failures, + ) + .await; + } + join_set.spawn(async move { + let outcome = run_case(case.clone(), rows).await; + assert_eq!(outcome.output, expected, "{case:?} (seed {seed})"); + (case, outcome.spilled) + }); + } + while !join_set.is_empty() { + collect_finished(&mut join_set, &mut spilled, &mut finished, &mut failures) + .await; + } + print_cases(cardinality, "spilled", &spilled); + print_cases(cardinality, "finished without spilling", &finished); + total_spilled += spilled.len(); + } + // A shape filter may select only shapes that cannot spill + if std::env::var("AGGREGATE_CHAIN_SHAPES").is_err() { + assert!(total_spilled > 0, "no case exercised the spill path"); + } + assert!( + failures.is_empty(), + "{} cases failed (seed {seed}):\n\n{}", + failures.len(), + failures.join("\n\n") + ); +} + +/// One line per case; `spilled_stages` names the aggregate operators that +/// spilled and flags when more than one did. +const MAX_CONCURRENT_CASES: usize = 16; +/// Generous: a case takes well under a second in debug builds. +const CASE_TIMEOUT_SECS: u64 = 60; + +/// Waits for one case and files it under spilled, finished or failed. A +/// failure does not stop the run, so one run reports every failing case. +async fn collect_finished( + join_set: &mut JoinSet<(Case, Vec)>, + spilled: &mut Vec<(Case, Vec)>, + finished: &mut Vec<(Case, Vec)>, + failures: &mut Vec, +) { + let Some(result) = join_set.join_next().await else { + return; + }; + match result { + Ok((case, stages)) if stages.is_empty() => finished.push((case, stages)), + Ok((case, stages)) => spilled.push((case, stages)), + Err(error) => failures.push(error.to_string()), + } +} + +fn print_cases(cardinality: Cardinality, outcome: &str, cases: &[(Case, Vec)]) { + let mut lines: Vec = cases + .iter() + .map(|(case, spilled_stages)| { + let spilled = match spilled_stages.len() { + 0 => String::new(), + 1 => format!(" spilled: {}", spilled_stages[0]), + _ => format!( + " spilled: {} (multiple stages)", + spilled_stages.join(" + ") + ), + }; + let skip_partial = if case.shape.has_skip_partial_candidate(case.order) { + format!(" skip_partial={:<5}", case.skip_partial_enabled) + } else { + " ".repeat(19) + }; + format!( + " {:<45} {:<17} migration={:<5} memory={:<9}{skip_partial}{spilled}", + case.shape.name, + format!("{:?}", case.order), + case.migration_enabled, + format!("{:?}", case.memory), + ) + }) + .collect(); + lines.sort(); + // Enable with `RUST_LOG=debug` + log::debug!("{cardinality:?}: {} cases {outcome}", lines.len()); + for line in lines { + log::debug!("{line}"); + } +} + +/// Reproduces one failing cell from its seed. +#[expect(dead_code)] +async fn run_single_case( + shape_name: &str, + order: Order, + migration_enabled: bool, + cardinality: Cardinality, + memory: Memory, + seed: u64, +) -> Result<()> { + let shape = shape_by_name(shape_name); + let rows = Arc::new(generate_rows(cardinality, seed)); + let expected = reference(shape.query, Arc::clone(&rows), cardinality).await; + let actual = run_case( + Case { + shape, + order, + migration_enabled, + cardinality, + memory, + skip_partial_enabled: true, + }, + rows, + ) + .await; + assert_eq!(actual.output, expected); + Ok(()) +} diff --git a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/AGGREGATE_CHAINS.md b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/AGGREGATE_CHAINS.md new file mode 100644 index 0000000000000..5f9c47d692862 --- /dev/null +++ b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/AGGREGATE_CHAINS.md @@ -0,0 +1,323 @@ +# Aggregate execution chains + +Every valid way to wire `AggregateExec` operators into a plan, the stream each +stage resolves to in `execute_typed` (`datafusion/physical-plan/src/aggregates/mod.rs`), +and what the input and config must look like to reach it. All chains compute +the same logical aggregation and must produce identical results once the +output is sorted by the group keys. + +TopK (`limit_options`) is covered with a limit larger than any possible group +count, so the TopK stream must return the complete aggregate. + +## How an aggregate stage picks its order mode + +Each `AggregateExec` computes its own `InputOrderMode` in `AggregateExec::try_new` +by comparing its group-by expressions against the ordering its input reports: + +- **Linear**: the input ordering covers none of the group keys. Plain hash + aggregation. +- **PartiallySorted**: the input ordering covers some of the group keys but not + all. The stream can flush groups whenever the sorted prefix changes. +- **Sorted**: the input ordering covers all group keys. The stream emits each + group as soon as the key changes and never holds more than one group. + +Because this is per operator, the same source ordering can produce different +modes at different stages: the shuffle between stages either keeps the ordering +(order-preserving `RepartitionExec`, `SortPreservingMergeExec`) or destroys it +(plain `RepartitionExec`, `CoalescePartitionsExec`). A final stage can never be +more ordered than the partial stage feeding it, because a Linear stage reports +no output ordering. + +Two kinds of rows flow through a chain: + +- **raw rows**: source data. Consumed by `Partial`, `Single`, `SinglePartitioned`. +- **partial state rows**: intermediate accumulator state produced by `Partial` + or `PartialReduce`. Consumed by `PartialReduce`, `Final`, `FinalPartitioned`. + +## Preconditions for the dedicated streams + +All rows below assume these hold. If any fails, every grouped stage falls back +to `GroupedHashAggregateStream` (see "Fallback chains"). + +- `datafusion.execution.enable_migration_aggregate = true` (default) +- exactly one grouping set (no `GROUPING SETS` / `CUBE` / `ROLLUP`) +- `limit_options = None` + +## Memory behavior per stream + +| Stream | Under memory pressure | Needs disk manager | +|---|---|---| +| `AggregateStream` (no grouping) | never allocates per group, nothing to do | no | +| `PartialHashAggregateStream` | emits current states downstream, resets table | no | +| `OrderedPartialAggregateStream` | emits current states downstream, resets table (on Sorted input too, since one group can outgrow the budget) | no | +| `PartialReduceHashAggregateStream` | emits current states downstream, resets table | no | +| `FinalHashAggregateStream` | sorts by full group key, spills, merges runs at end | yes | +| `SingleHashAggregateStream` | same as final | yes | +| `OrderedFinalAggregateStream`, Sorted input | streams groups out as key changes, cannot spill | no | +| `OrderedFinalAggregateStream`, PartiallySorted input | spills like final hash | yes | +| `OrderedSingleAggregateStream`, Sorted input | like ordered final with Sorted input | no | +| `OrderedSingleAggregateStream`, PartiallySorted input | spills like single hash | yes | +| `GroupedTopKAggregateStream` | bounded heap, never spills; with a limit above the group count it holds every group | no | +| `GroupedHashAggregateStream` (fallback) | `Partial` emits early: completed prefix groups on ordered input, or everything when there are none; `PartialReduce`, final and single modes spill on Linear or PartiallySorted input and report an error on Sorted input; spilling needs extra sort headroom, so it can fail under a tight pool where the dedicated streams succeed | yes except `Partial` | + +`PartialHashAggregateStream` and the fallback in `Partial` mode also have the +skip-partial probe (`skip_partial_aggregation_probe_*` options): when +cardinality is high relative to input rows they stop aggregating and convert raw +rows straight into state rows. This changes what the final stage receives, not +the result. The ordered partial stream has no such probe. + +--- + +## Chains: grouped aggregation + +Chains are written bottom to top with `→`. + +### Single stage + +| Case | Chain | Stream | Input requirements | +|---|---|---|---| +| `single, unordered` | `Single`, Linear | `SingleHashAggregateStream` | raw, 1 partition, no ordering on group keys | +| `single, sorted by all keys` | `Single`, Sorted | `OrderedSingleAggregateStream` | raw, 1 partition, ordered by all group keys | +| `single, sorted by first key` | `Single`, PartiallySorted | `OrderedSingleAggregateStream` with spill | raw, 1 partition, ordered by a subset of the group keys | +| `single_partitioned, unordered` | `SinglePartitioned`, Linear | `SingleHashAggregateStream` per partition | raw, several partitions hash-partitioned by group keys, no ordering | +| `single_partitioned_order_preserving, sorted by all keys` | `SinglePartitioned`, Sorted | `OrderedSingleAggregateStream` per partition | raw, several partitions hash-partitioned by group keys, each ordered by all group keys | +| `single_partitioned_order_preserving, sorted by first key` | `SinglePartitioned`, PartiallySorted | `OrderedSingleAggregateStream` with spill per partition | raw, several partitions hash-partitioned by group keys, each ordered by a subset | + +The `single_partitioned` cases are what the `CombinePartialFinalAggregate` optimizer rule produces +from adjacent `Partial` + `FinalPartitioned`. To hit them directly, build the +plan by hand with a hash `RepartitionExec` (order-preserving for the sorted variants) +directly under the aggregate. + +### Two stage, ordering lost or absent + +| Case | Chain | Streams | Input requirements | +|---|---|---|---| +| `partial_repartition_final, unordered` | `Partial` Linear → hash `RepartitionExec` → `FinalPartitioned` Linear | `PartialHashAggregateStream` → `FinalHashAggregateStream` | raw, several partitions, no ordering. **Default planner output.** | +| `partial_coalesce_final, unordered` | `Partial` Linear → `CoalescePartitionsExec` → `Final` Linear | same | raw, several partitions, no ordering. Planner output when `repartition_aggregations = false` | +| `partial_repartition_final, sorted by all keys` | `Partial` Sorted → hash `RepartitionExec` → `FinalPartitioned` Linear | `OrderedPartialAggregateStream` → `FinalHashAggregateStream` | raw ordered by all group keys, non-preserving shuffle | +| `partial_repartition_final, sorted by first key` | `Partial` PartiallySorted → hash `RepartitionExec` → `FinalPartitioned` Linear | same | raw ordered by a subset of group keys, non-preserving shuffle | +| `partial_coalesce_final, sorted by all keys` | `Partial` Sorted → `CoalescePartitionsExec` → `Final` Linear | same | raw ordered by all group keys | +| `partial_coalesce_final, sorted by first key` | `Partial` PartiallySorted → `CoalescePartitionsExec` → `Final` Linear | same | raw ordered by a subset of group keys | + +### Two stage, ordering preserved + +| Case | Chain | Streams | Input requirements | +|---|---|---|---| +| `partial_order_preserving_repartition_final, sorted by all keys` | `Partial` Sorted → order-preserving hash `RepartitionExec` → `FinalPartitioned` Sorted | `OrderedPartialAggregateStream` → `OrderedFinalAggregateStream`, no spill | raw, several partitions, each ordered by all group keys | +| `partial_order_preserving_repartition_final, sorted by first key` | `Partial` PartiallySorted → order-preserving hash `RepartitionExec` → `FinalPartitioned` PartiallySorted | `OrderedPartialAggregateStream` → `OrderedFinalAggregateStream` with spill | raw, several partitions, each ordered by a subset | +| `partial_sort_preserving_merge_final, sorted by all keys` | `Partial` Sorted → `SortPreservingMergeExec` → `Final` Sorted | `OrderedPartialAggregateStream` → `OrderedFinalAggregateStream`, no spill | raw, several partitions, each ordered by all group keys | +| `partial_sort_preserving_merge_final, sorted by first key` | `Partial` PartiallySorted → `SortPreservingMergeExec` → `Final` PartiallySorted | `OrderedPartialAggregateStream` → `OrderedFinalAggregateStream` with spill | raw, several partitions, each ordered by a subset | +| `partial_final_single_partition, sorted by all keys` | `Partial` Sorted → `Final` Sorted | `OrderedPartialAggregateStream` → `OrderedFinalAggregateStream` | raw, 1 partition ordered by all group keys, no shuffle | +| `partial_final_single_partition, sorted by first key` | `Partial` PartiallySorted → `Final` PartiallySorted | `OrderedPartialAggregateStream` → `OrderedFinalAggregateStream` with spill | raw, 1 partition ordered by a subset | + +The `partial_order_preserving_repartition_final` cases are the planner's ordered fast path when `prefer_existing_sort = true` +and the source reports the ordering. + +### Three stage (tree reduce) + +Not produced by the SQL planner. Build by hand. + +| Case | Chain | Streams | Input requirements | +|---|---|---|---| +| `partial_repartition_reduce_repartition_final` | `Partial` → hash `RepartitionExec` → `PartialReduce` → hash `RepartitionExec` → `FinalPartitioned`, all Linear | `PartialHashAggregateStream` → `PartialReduceHashAggregateStream` → `FinalHashAggregateStream` | raw, several partitions, no ordering | +| `partial_repartition_reduce_coalesce_final` | `Partial` → hash `RepartitionExec` → `PartialReduce` → `CoalescePartitionsExec` → `Final`, all Linear | same | raw, several partitions | +| `partial_local_reduce_repartition_final` | `Partial` → `PartialReduce` → hash `RepartitionExec` → `FinalPartitioned`, all Linear | same | raw, several partitions; the reduce stage merges within each partition only | +| `partial_reduce_final_order_preserving` | `Partial` ordered → order-preserving hash `RepartitionExec` → `PartialReduce` ordered → ... | `OrderedPartialAggregateStream` → **fallback `GroupedHashAggregateStream`** → ... | ordered input; `PartialReduce` has no ordered stream, so this is a fallback chain | + +`PartialReduce` is the only stage whose input and output are both partial +state. Its input partitioning does not matter for correctness; hash +partitioning just makes it reduce more. + +## Chains: TopK + +`GROUP BY k1` with a single `max(v)` and a limit larger than any possible group +count. The planner would only produce the shape with the limit on the final +stage; the others are built by hand. + +| Case | Chain | Streams | Input requirements | +|---|---|---|---| +| `top_k_query_single` | `Single`, no limit | `SingleHashAggregateStream` or `OrderedSingleAggregateStream` | reference for the TopK chains | +| `top_k_single` | `Single` with limit | `GroupedTopKAggregateStream` | raw, 1 partition, any ordering | +| `top_k_partial_repartition_final` | `Partial` → hash `RepartitionExec` → `FinalPartitioned` with limit | partial stream → `GroupedTopKAggregateStream` | raw, several partitions. **Planner shape** | +| `top_k_partial_coalesce_final` | `Partial` → `CoalescePartitionsExec` → `Final` with limit | partial stream → `GroupedTopKAggregateStream` | raw, several partitions | +| `top_k_both_stages` | `Partial` with limit → hash `RepartitionExec` → `FinalPartitioned` with limit | `GroupedTopKAggregateStream` → `GroupedTopKAggregateStream` | raw, several partitions | + +Sorting by `k1` alone already makes this query's input Sorted, so these chains +have no PartiallySorted variant. + +## Chains: group key types + +Every chain above groups by two `Int64` keys, which `GroupValuesColumn` +handles. These queries run the single-stage shape and the default two-stage +planner shape with the same six aggregates but different key columns, so each +`GroupValues` implementation is exercised across partial, final and single +stages, with and without spilling. + +| Query | Keys | Implementation | Orders | +|---|---|---|---| +| `boolean_key_*` | `b Boolean` | `GroupValuesBoolean` | unordered, sorted | +| `bytes_key_*` | `s Utf8` | `GroupValuesBytes` | unordered, sorted | +| `bytes_view_key_*` | `sv Utf8View` | `GroupValuesBytesView` | unordered, sorted | +| `primitive_key_*` | `p Int64` | `GroupValuesPrimitive` | unordered, sorted | +| `mixed_keys_*` | `b, s, sv, p` | `GroupValuesColumn` with mixed types; the order-preserving shape reaches its ordered variant | unordered, sorted by first key, sorted | +| `struct_key_*` | `st Struct, num: Int64>` | row fallback `GroupValuesRows` | unordered only, structs cannot be sorted by the sort kernels | + +Every key column is nullable and carries `cardinality` distinct values, except +the boolean one which has two. + +## Chains: no grouping + +Every stage is `AggregateStream` regardless of mode or ordering. Input ordering +is irrelevant. Nothing spills. + +| Case | Chain | Input requirements | +|---|---|---| +| `no_grouping_single` | `Single` | 1 partition | +| `no_grouping_partial_coalesce_final` | `Partial` → `CoalescePartitionsExec` → `Final` | several partitions. **Default planner output.** | +| `no_grouping_partial_reduce_final` | `Partial` → `CoalescePartitionsExec` → `PartialReduce` → `CoalescePartitionsExec` → `Final` | several partitions, by hand | +| `no_grouping_partial_final_single_partition` | `Partial` → `Final` | 1 partition | + +## Fallback chains + +Any chain above with one of these flipped runs the same shape on +`GroupedHashAggregateStream` for every grouped stage: + +| Case | Trigger | Notes | +|---|---|---| +| migration disabled | `enable_migration_aggregate = false` | Run every grouped chain twice, once per flag value. This is the main old-vs-new comparison. | +| grouping sets | grouping sets present | Order mode is forced to Linear; only the unordered chains are reachable | +| ordered partial reduce | `PartialReduce` with ordered input | `partial_reduce_final_order_preserving` above | + +--- + +## Scenario dimensions + +### Cardinality (distinct group keys divided by input rows) + +| Scenario | Ratio | What it exercises | +|---|---|---| +| very high | about 1.0, nearly all rows unique | skip-partial probe fires in the Linear partial stage; final and single stages spill hardest; Sorted streams emit one row per group | +| high | about 0.5 | skip-partial probe likely fires; large final tables | +| medium | about 0.05 | partial actually reduces; final table moderate | +| low | about 0.001 | few groups, partial reduces almost everything; PartiallySorted streams see many rows per prefix | +| very low | 1 to 10 groups | degenerate hash tables; `PartialReduce` merges to almost nothing | + +Skip-partial only fires in a Linear `Partial` stage, so the very high and high +rows are where the unordered partial chains differ most from the ordered chains. + +### Memory + +| Scenario | Config | Expected behavior | +|---|---|---| +| high | unlimited pool, or pool much larger than the data | no spill, no early emit anywhere | +| medium | `RuntimeEnvBuilder::with_memory_limit(...)` sized so the final or single table cannot fit, disk manager enabled | partial stages emit early; `FinalHashAggregateStream`, `SingleHashAggregateStream` and the PartiallySorted ordered streams spill; Sorted ordered streams are unaffected | + +A resources-exhausted error is never accepted, in either scenario. Every stream +is expected to spill, emit early, or stay bounded, so running out of memory +under the medium pool is treated as a bug in that stream's memory handling or in +how the stages share the pool. The run as a whole must also have spilled +somewhere. Known failures at the time of writing are listed under "Findings". + +For the medium scenario also run once with the disk manager disabled: spilling +streams must return a resources-exhausted error rather than wrong results, and +Sorted ordered streams must still succeed. + +### Skip partial aggregation + +| Scenario | Config | Applies to | +|---|---|---| +| enabled | `skip_partial_aggregation_probe_ratio_threshold = 0.8`, rows threshold lowered to 1024 | chains with a grouped `Partial` stage on Linear input, the only place the probe runs | +| disabled | ratio threshold `1.0`, which turns the probe off | same chains | + +Chains without such a stage run with the probe enabled only, since it can never +fire there. With it disabled the `skipped_aggregation_rows` metric must stay at +zero everywhere. + +### Full matrix + +Every chain in the tables above, plus the migration-disabled fallback over all +grouped ones, times 5 cardinalities, times 2 memory settings, times 2 skip +partial settings where the probe applies. + +Every cell must produce the same result as `single, unordered` with unlimited +memory for the same dataset and query, compared after sorting by group keys. + +## Data generation attributes needed + +To reach each row above the generator must be able to control: + +| Attribute | Values | Selects | +|---|---|---| +| group key columns | 1 or more | PartiallySorted is only reachable with at least 2 keys | +| sort of the source | none, all group keys, or a subset of group keys | Linear, Sorted, PartiallySorted | +| partition count | 1 or several | single versus partitioned chains | +| partitioning of the source | arbitrary or hash by group keys | the `single_partitioned` cases need hash partitioning, either from the source or from a `RepartitionExec` under the aggregate | +| cardinality | see scenario table | | +| batch size | small values, 1 to 64, matter | forces group boundaries to cross batches in the ordered streams | +| aggregate set | must include multi-field state such as `avg`, `count distinct`, `median`, `array_agg` | otherwise `Partial`, `Single` and `PartialReduce` are indistinguishable | +| nulls in group keys | yes or no | null groups are a separate code path in group values | + +## Per-chain checklist + +For each chain above, assert before comparing results: + +1. `explain` shows the expected `mode=` and `ordering_mode=` on every `AggregateExec`. +2. In the medium memory scenario, `spill_count > 0` on the stages listed as spilling and `== 0` on the rest. +3. In very high cardinality with a Linear `Partial` stage, the `skipped_aggregation_rows` metric is `> 0`. +4. Results equal the reference (`single, unordered`, unlimited memory) after sorting by group keys. + +## Findings from running the matrix + +Open: none at the time of writing. The occasional hangs seen earlier under the +limited pool stopped once the source batches were copied instead of sliced (see +the accounting item below); they were only ever observed while every batch was +charged the whole partition and `RepartitionExec` spilled constantly. The +per-case timeout stays in place to catch a recurrence. + +Fixed or documented: + +- `OrderedPartialAggregateStream` refused to emit early on Sorted input and + reported an error instead. It now emits its partial state like it does for + PartiallySorted input, and registers as able to handle memory pressure. +- `PartialReduceHashAggregateStream` registered its reservation as not + spillable although it emits early under pressure, so `FairSpillPool` neither + capped it nor counted it against the spillers, and several of them could hold + the whole pool while the final stage below starved. It now registers as + spillable. +- The legacy `GroupedHashAggregateStream` in `Partial` mode could only emit + completed prefix groups on ordered input and reported an error when there + were none (a boolean or low-cardinality first key spans whole partitions). It + now falls back to emitting everything and restarting its ordering. +- `GroupValuesColumn::emit(EmitTo::All)` replaced the column builders but left + the hash map populated, so the next `intern` dereferenced stale group indices + into empty builders (index out of bounds). The map is now cleared too. +- `RepartitionExec` and the order-preserving merge account a batch by + `get_array_memory_size`, which for a slice is the size of the whole backing + buffers. Memory-table sources hand out 64-row slices of a partition, so every + batch was charged the entire partition and the merge could never reserve one. + The fuzz test now copies each batch, as a real scan would. The accounting + itself is unchanged. + +- The legacy `GroupedHashAggregateStream` with a single nested group key + (`Struct`, `Map`) produced duplicate groups after spilling. When it switches + to merging the sorted spill files it relies on `GroupOrderingFull`, which + requires group ids in first-seen order, and it recreated the group values + collector to guarantee that only when there was more than one group column. + A single nested column has no specialized single-column collector and is + handled by `GroupValuesColumn` through a row-backed column, whose vectorized + interning assigns ids out of input order, so the ordering emitted a group + that was still in progress and the next batch reopened it as a new group. + Fixed by always recreating the collector for the merge phase. + +- `PrimitiveDistinctCountGroupsAccumulator` reported its capacity in `size()` + but only cleared its buffers when emitting everything, so a table looked as + large after `take_state_batch` as before. Every dedicated stream then failed + its post-emit or post-spill resize, and the legacy stream could not reserve + its sort headroom. Fixed by releasing the buffers on `EmitTo::All`. +- `PartialReduceHashAggregateStream` registers its reservation as not + spillable, so `FairSpillPool` neither caps it nor counts it against the + spillers, and several of them can hold the pool while the final stage below + starves. +- When a final or single stream is starved with nothing reserved, its + post-spill `try_resize` is really a grow from zero, and the error is reported + as "Decreasing allocation after spilling should succeed". diff --git a/datafusion/core/tests/fuzz_cases/mod.rs b/datafusion/core/tests/fuzz_cases/mod.rs index 3e425e48c7a0c..14dbacfa1a9a8 100644 --- a/datafusion/core/tests/fuzz_cases/mod.rs +++ b/datafusion/core/tests/fuzz_cases/mod.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +mod aggregate_chain_fuzz; #[expect(clippy::needless_pass_by_value)] mod aggregate_fuzz; mod distinct_count_string_fuzz; From 07dd46fc98b3e0a9f1b5bb060fb73a5816c17acf Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 7 Sep 2026 11:37:29 +0300 Subject: [PATCH 03/21] mark some cases as allowed to OOM and set memory limit to not be UNBOUNDED --- .../tests/fuzz_cases/aggregate_chain_fuzz.rs | 277 +++++++++++++----- 1 file changed, 211 insertions(+), 66 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs index 3c37780793b63..b639d9a8c3b3a 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -36,8 +36,8 @@ use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef, SortOptions}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::prelude::SessionConfig; -use datafusion_common::Result; use datafusion_common::test_util::batches_to_sort_string; +use datafusion_common::{DataFusionError, Result}; use datafusion_common_runtime::JoinSet; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{FairSpillPool, TrackConsumersPool}; @@ -202,6 +202,19 @@ impl Query { /// Larger than any possible number of groups, so TopK keeps every group. const TOP_K_LIMIT: usize = 2 * ROWS; +/// Everything that varies for a case apart from the shape itself. Passed as a +/// struct so a new dimension does not change every shape predicate. +#[derive(Clone, Copy, Debug, PartialEq)] +struct CaseParams { + order: Order, + migration_enabled: bool, + cardinality: Cardinality, + memory: Memory, + /// Whether the skip-partial probe may fire. Only varied for shapes with a + /// grouped `Partial` stage on Linear input, since nothing else runs it. + skip_partial_enabled: bool, +} + /// A plan shape. #[derive(Debug)] struct Shape { @@ -210,6 +223,32 @@ struct Shape { /// Source partition count. source_partitions: usize, query: Query, + /// Whether a resources-exhausted error is an accepted outcome for this + /// chain instead of a result, given the rest of the case parameters. + /// Every shape declares its own, so nothing is accepted by derivation. + accepts_out_of_memory: fn(&CaseParams) -> bool, +} + +/// A chain never runs out of memory: every stage spills, emits early, or is +/// bounded. +fn never_out_of_memory(_params: &CaseParams) -> bool { + false +} + +/// The chain starts with a grouped `Partial` aggregate reading the source, so +/// with `SortedByFirstKey` input that stage sees `InputOrderMode::PartiallySorted`. +/// +/// The legacy `GroupedHashAggregateStream` answers memory pressure there by +/// emitting the groups whose sort prefix is complete, and when one prefix value +/// spans the whole partition there is nothing it may emit, so it reports the +/// error instead of degrading. The dedicated `OrderedPartialAggregateStream` +/// emits everything. A `Partial` stage only reaches the legacy stream while +/// `enable_migration_aggregate` is off, an option that goes away once the +/// migration finishes, so this is accepted rather than fixed. +fn legacy_partial_on_partially_sorted_input(params: &CaseParams) -> bool { + !params.migration_enabled + && params.memory == Memory::Limited + && params.order == Order::SortedByFirstKey } const fn shape( @@ -217,30 +256,40 @@ const fn shape( operators: &'static [Operator], source_partitions: usize, query: Query, + accepts_out_of_memory: fn(&CaseParams) -> bool, ) -> Shape { Shape { name, operators, source_partitions, query, + accepts_out_of_memory, } } /// Every shape from AGGREGATE_CHAINS.md. The ordering variants there come from /// crossing a shape with `Order`, so one entry here covers several rows. const SHAPES: &[Shape] = &[ - shape("single", &[Aggregate(Single)], 1, Query::Grouped), + shape( + "single", + &[Aggregate(Single)], + 1, + Query::Grouped, + never_out_of_memory, + ), shape( "single_partitioned", &[HashRepartition, Aggregate(SinglePartitioned)], PARTITIONS, Query::Grouped, + never_out_of_memory, ), shape( "single_partitioned_order_preserving", &[OrderPreservingHashRepartition, Aggregate(SinglePartitioned)], PARTITIONS, Query::Grouped, + never_out_of_memory, ), shape( "partial_repartition_final", @@ -251,12 +300,14 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::Grouped, + legacy_partial_on_partially_sorted_input, ), shape( "partial_coalesce_final", &[Aggregate(Partial), CoalescePartitions, Aggregate(Final)], PARTITIONS, Query::Grouped, + legacy_partial_on_partially_sorted_input, ), shape( "partial_order_preserving_repartition_final", @@ -267,18 +318,21 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::Grouped, + legacy_partial_on_partially_sorted_input, ), shape( "partial_sort_preserving_merge_final", &[Aggregate(Partial), SortPreservingMerge, Aggregate(Final)], PARTITIONS, Query::Grouped, + legacy_partial_on_partially_sorted_input, ), shape( "partial_final_single_partition", &[Aggregate(Partial), Aggregate(Final)], 1, Query::Grouped, + legacy_partial_on_partially_sorted_input, ), shape( "partial_repartition_reduce_repartition_final", @@ -291,6 +345,7 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::Grouped, + legacy_partial_on_partially_sorted_input, ), shape( "partial_repartition_reduce_coalesce_final", @@ -303,6 +358,7 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::Grouped, + legacy_partial_on_partially_sorted_input, ), shape( "partial_local_reduce_repartition_final", @@ -314,6 +370,7 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::Grouped, + legacy_partial_on_partially_sorted_input, ), // ordered PartialReduce has no dedicated stream, lands on the fallback shape( @@ -327,34 +384,51 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::Grouped, + legacy_partial_on_partially_sorted_input, ), shape( "no_grouping_single", &[Aggregate(Single)], 1, Query::NoGrouping, + never_out_of_memory, ), // TopK: same query without a limit is the reference for the TopK chains - shape("top_k_query_single", &[Aggregate(Single)], 1, Query::TopK), - shape("top_k_single", &[TopK(Single)], 1, Query::TopK), + shape( + "top_k_query_single", + &[Aggregate(Single)], + 1, + Query::TopK, + never_out_of_memory, + ), + shape( + "top_k_single", + &[TopK(Single)], + 1, + Query::TopK, + never_out_of_memory, + ), // planner shape: the limit lands on the aggregate under the sort shape( "top_k_partial_repartition_final", &[Aggregate(Partial), HashRepartition, TopK(FinalPartitioned)], PARTITIONS, Query::TopK, + never_out_of_memory, ), shape( "top_k_partial_coalesce_final", &[Aggregate(Partial), CoalescePartitions, TopK(Final)], PARTITIONS, Query::TopK, + never_out_of_memory, ), shape( "top_k_both_stages", &[TopK(Partial), HashRepartition, TopK(FinalPartitioned)], PARTITIONS, Query::TopK, + never_out_of_memory, ), // group key types: single stage and the default two-stage planner shape shape( @@ -362,6 +436,7 @@ const SHAPES: &[Shape] = &[ &[Aggregate(Single)], 1, Query::BooleanKey, + never_out_of_memory, ), shape( "boolean_key_partial_repartition_final", @@ -372,8 +447,15 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::BooleanKey, + never_out_of_memory, + ), + shape( + "bytes_key_single", + &[Aggregate(Single)], + 1, + Query::BytesKey, + never_out_of_memory, ), - shape("bytes_key_single", &[Aggregate(Single)], 1, Query::BytesKey), shape( "bytes_key_partial_repartition_final", &[ @@ -383,12 +465,14 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::BytesKey, + never_out_of_memory, ), shape( "bytes_view_key_single", &[Aggregate(Single)], 1, Query::BytesViewKey, + never_out_of_memory, ), shape( "bytes_view_key_partial_repartition_final", @@ -399,12 +483,14 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::BytesViewKey, + never_out_of_memory, ), shape( "primitive_key_single", &[Aggregate(Single)], 1, Query::PrimitiveKey, + never_out_of_memory, ), shape( "primitive_key_partial_repartition_final", @@ -415,12 +501,14 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::PrimitiveKey, + never_out_of_memory, ), shape( "mixed_keys_single", &[Aggregate(Single)], 1, Query::MixedKeys, + never_out_of_memory, ), shape( "mixed_keys_partial_repartition_final", @@ -431,6 +519,7 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::MixedKeys, + legacy_partial_on_partially_sorted_input, ), // ordered multi-column group values shape( @@ -442,12 +531,14 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::MixedKeys, + legacy_partial_on_partially_sorted_input, ), shape( "struct_key_single", &[Aggregate(Single)], 1, Query::StructKey, + never_out_of_memory, ), shape( "struct_key_partial_repartition_final", @@ -458,12 +549,14 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::StructKey, + never_out_of_memory, ), shape( "no_grouping_partial_coalesce_final", &[Aggregate(Partial), CoalescePartitions, Aggregate(Final)], PARTITIONS, Query::NoGrouping, + never_out_of_memory, ), shape( "no_grouping_partial_reduce_final", @@ -476,12 +569,14 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::NoGrouping, + never_out_of_memory, ), shape( "no_grouping_partial_final_single_partition", &[Aggregate(Partial), Aggregate(Final)], 1, Query::NoGrouping, + never_out_of_memory, ), ]; @@ -519,14 +614,15 @@ impl Shape { #[derive(Clone, Debug)] struct Case { shape: &'static Shape, - order: Order, - migration_enabled: bool, - cardinality: Cardinality, - memory: Memory, - /// Whether the skip-partial probe may fire. Only varied for shapes with a - /// grouped `Partial` stage on Linear input, since nothing else runs the - /// probe. - skip_partial_enabled: bool, + params: CaseParams, +} + +impl Case { + /// Whether a resources-exhausted error is accepted for this case instead + /// of a result, as declared by its shape. + fn accepts_out_of_memory(&self) -> bool { + (self.shape.accepts_out_of_memory)(&self.params) + } } impl Shape { @@ -574,11 +670,13 @@ fn all_cases() -> Vec { for &skip_partial_enabled in skip_partial_variants { cases.push(Case { shape, - order, - migration_enabled, - cardinality, - memory, - skip_partial_enabled, + params: CaseParams { + order, + migration_enabled, + cardinality, + memory, + skip_partial_enabled, + }, }); } } @@ -942,7 +1040,7 @@ fn task_context(case: &Case) -> Arc { .with_target_partitions(PARTITIONS) .set_bool( "datafusion.execution.enable_migration_aggregate", - case.migration_enabled, + case.params.migration_enabled, ) // The default is 100k rows. Lower it so the skip-partial probe can // fire on our per-partition row counts. A ratio threshold of 1.0 @@ -956,10 +1054,15 @@ fn task_context(case: &Case) -> Arc { .options_mut() .execution .skip_partial_aggregation_probe_ratio_threshold = - if case.skip_partial_enabled { 0.8 } else { 1.0 }; + if case.params.skip_partial_enabled { + 0.8 + } else { + 1.0 + }; - let runtime = match case.memory { - Memory::Unlimited => RuntimeEnvBuilder::new(), + let runtime = match case.params.memory { + // Not using UnboundedMemoryPool, so users would still think that we have a valid pool, but just with enough memory + Memory::Unlimited => RuntimeEnvBuilder::new().with_memory_limit(usize::MAX, 1.0), // Small enough that a very-high-cardinality final table spills, large // enough that the legacy stream can still reserve its sort headroom // and that RepartitionExec / SortPreservingMergeExec succeed. The @@ -1050,7 +1153,7 @@ fn can_spill(case: &Case, aggregate: &AggregateExec) -> bool { // Linear input; ordered input or migration off run the legacy stream, // which spills. PartialReduce => { - !case.migration_enabled + !case.params.migration_enabled || *aggregate.input_order_mode() != InputOrderMode::Linear } Partial => false, @@ -1072,7 +1175,7 @@ fn check_plan_shape(case: &Case, plan: &Arc) { return; } let nodes = aggregate_nodes(plan); - let expected = expected_orders(case.shape, case.order); + let expected = expected_orders(case.shape, case.params.order); assert_eq!(nodes.len(), expected.len(), "{case:?}"); for (node, expected_order) in nodes.iter().zip(expected) { let aggregate = as_aggregate(node); @@ -1106,7 +1209,7 @@ fn check_metrics(case: &Case, plan: &Arc) -> Vec { .map(|metric| metric.as_usize()) .unwrap_or(0); - match case.memory { + match case.params.memory { Memory::Unlimited => { assert_eq!(spill_count, 0, "{case:?}: unexpected spill in {mode:?}"); } @@ -1122,10 +1225,10 @@ fn check_metrics(case: &Case, plan: &Arc) -> Vec { // Only the two-key query has as many groups as `cardinality` says; // the TopK query groups by `k1` alone and stays far below the ratio. - if case.memory == Memory::Unlimited - && case.cardinality == Cardinality::VeryHigh + if case.params.memory == Memory::Unlimited + && case.params.cardinality == Cardinality::VeryHigh && case.shape.query == Query::Grouped - && case.skip_partial_enabled + && case.params.skip_partial_enabled && runs_skip_partial_probe(aggregate) { assert!( @@ -1133,7 +1236,7 @@ fn check_metrics(case: &Case, plan: &Arc) -> Vec { "{case:?}: skip-partial probe did not fire" ); } - if !case.skip_partial_enabled || !runs_skip_partial_probe(aggregate) { + if !case.params.skip_partial_enabled || !runs_skip_partial_probe(aggregate) { assert_eq!(skipped_rows, 0, "{case:?}: skip-partial fired in {mode:?}"); } } @@ -1144,17 +1247,23 @@ fn check_metrics(case: &Case, plan: &Arc) -> Vec { // Driver // --------------------------------------------------------------------------- -/// Sorted output plus the stages that spilled, empty if none did. -struct Outcome { - output: String, - spilled: Vec, +enum Outcome { + /// Sorted output plus the stages that spilled, empty if none did. + Finished { + output: String, + spilled: Vec, + }, + /// Ran out of memory, and the case was built to accept that. See + /// [`Shape::accepts_out_of_memory`]. + OutOfMemory, } /// Runs one case, checks plan shape and metrics, and returns its outcome. /// -/// Running out of memory is never accepted, not even under the limited pool: -/// every stream either spills, emits early, or is bounded, so an error there -/// is a bug in a stream's memory handling or in how stages share the pool. +/// Running out of memory fails the run unless the case was built with +/// `accepts_out_of_memory`: every other stream either spills, emits early, or +/// is bounded, so an error there is a bug in a stream's memory handling or in +/// how stages share the pool. async fn run_case(case: Case, rows: Arc) -> Outcome { log::debug!("start {case:?}"); let outcome = run_case_inner(&case, rows).await; @@ -1166,12 +1275,12 @@ async fn run_case_inner(case: &Case, rows: Arc) -> Outcome { let partitions = arrange( &rows, case.shape.query, - case.order, + case.params.order, case.shape.source_partitions, ); let plan = build_plan( case.shape, - source(&partitions, case.shape.query, case.order), + source(&partitions, case.shape.query, case.params.order), ); check_plan_shape(case, &plan); @@ -1189,13 +1298,19 @@ async fn run_case_inner(case: &Case, rows: Arc) -> Outcome { }); let batches = match collected { Ok(batches) => batches, + Err(error) + if matches!(error.find_root(), DataFusionError::ResourcesExhausted(_)) + && case.accepts_out_of_memory() => + { + return Outcome::OutOfMemory; + } Err(error) => panic!( "{case:?} failed: {error}\n{}", displayable(plan.as_ref()).indent(true) ), }; let spilled = check_metrics(case, &plan); - Outcome { + Outcome::Finished { output: batches_to_sort_string(&batches), spilled, } @@ -1217,16 +1332,23 @@ async fn reference( let outcome = run_case( Case { shape, - order: Order::Unordered, - migration_enabled: true, - cardinality, - memory: Memory::Unlimited, - skip_partial_enabled: true, + params: CaseParams { + order: Order::Unordered, + migration_enabled: true, + cardinality, + memory: Memory::Unlimited, + skip_partial_enabled: true, + }, }, rows, ) .await; - outcome.output + match outcome { + Outcome::Finished { output, .. } => output, + Outcome::OutOfMemory => { + unreachable!("the reference runs with unlimited memory and migration on") + } + } } #[tokio::test(flavor = "multi_thread")] @@ -1251,10 +1373,10 @@ async fn aggregate_chain_fuzz() { } let mut join_set = JoinSet::new(); - let (mut spilled, mut finished) = (vec![], vec![]); + let (mut spilled, mut finished, mut out_of_memory) = (vec![], vec![], vec![]); for case in all_cases() .into_iter() - .filter(|case| case.cardinality == cardinality) + .filter(|case| case.params.cardinality == cardinality) { let rows = Arc::clone(&rows); let expected = expected_by_query @@ -1269,22 +1391,38 @@ async fn aggregate_chain_fuzz() { &mut join_set, &mut spilled, &mut finished, + &mut out_of_memory, &mut failures, ) .await; } join_set.spawn(async move { - let outcome = run_case(case.clone(), rows).await; - assert_eq!(outcome.output, expected, "{case:?} (seed {seed})"); - (case, outcome.spilled) + match run_case(case.clone(), rows).await { + Outcome::Finished { output, spilled } => { + assert_eq!(output, expected, "{case:?} (seed {seed})"); + (case, Some(spilled)) + } + Outcome::OutOfMemory => (case, None), + } }); } while !join_set.is_empty() { - collect_finished(&mut join_set, &mut spilled, &mut finished, &mut failures) - .await; + collect_finished( + &mut join_set, + &mut spilled, + &mut finished, + &mut out_of_memory, + &mut failures, + ) + .await; } print_cases(cardinality, "spilled", &spilled); print_cases(cardinality, "finished without spilling", &finished); + print_cases( + cardinality, + "ran out of memory (accepted, legacy partial stream)", + &out_of_memory, + ); total_spilled += spilled.len(); } // A shape filter may select only shapes that cannot spill @@ -1308,17 +1446,19 @@ const CASE_TIMEOUT_SECS: u64 = 60; /// Waits for one case and files it under spilled, finished or failed. A /// failure does not stop the run, so one run reports every failing case. async fn collect_finished( - join_set: &mut JoinSet<(Case, Vec)>, + join_set: &mut JoinSet<(Case, Option>)>, spilled: &mut Vec<(Case, Vec)>, finished: &mut Vec<(Case, Vec)>, + out_of_memory: &mut Vec<(Case, Vec)>, failures: &mut Vec, ) { let Some(result) = join_set.join_next().await else { return; }; match result { - Ok((case, stages)) if stages.is_empty() => finished.push((case, stages)), - Ok((case, stages)) => spilled.push((case, stages)), + Ok((case, None)) => out_of_memory.push((case, vec![])), + Ok((case, Some(stages))) if stages.is_empty() => finished.push((case, stages)), + Ok((case, Some(stages))) => spilled.push((case, stages)), Err(error) => failures.push(error.to_string()), } } @@ -1335,17 +1475,18 @@ fn print_cases(cardinality: Cardinality, outcome: &str, cases: &[(Case, Vec Date: Mon, 7 Sep 2026 13:05:59 +0300 Subject: [PATCH 04/21] update comment and bad oom function --- .../tests/fuzz_cases/aggregate_chain_fuzz.rs | 64 ++++++++++++------- 1 file changed, 40 insertions(+), 24 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs index b639d9a8c3b3a..e529c11fa2116 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -236,19 +236,35 @@ fn never_out_of_memory(_params: &CaseParams) -> bool { } /// The chain starts with a grouped `Partial` aggregate reading the source, so -/// with `SortedByFirstKey` input that stage sees `InputOrderMode::PartiallySorted`. -/// -/// The legacy `GroupedHashAggregateStream` answers memory pressure there by -/// emitting the groups whose sort prefix is complete, and when one prefix value -/// spans the whole partition there is nothing it may emit, so it reports the -/// error instead of degrading. The dedicated `OrderedPartialAggregateStream` -/// emits everything. A `Partial` stage only reaches the legacy stream while -/// `enable_migration_aggregate` is off, an option that goes away once the -/// migration finishes, so this is accepted rather than fixed. -fn legacy_partial_on_partially_sorted_input(params: &CaseParams) -> bool { - !params.migration_enabled - && params.memory == Memory::Limited - && params.order == Order::SortedByFirstKey +/// with ordered input that stage cannot always survive memory pressure. +fn partial_stage_starved_on_ordered_input(params: &CaseParams) -> bool { + params.memory == Memory::Limited + && match params.order { + // The stage sees `InputOrderMode::PartiallySorted`. The legacy + // `GroupedHashAggregateStream` emits only the groups whose sort + // prefix is complete, and when one prefix value spans the whole + // partition there is nothing it may emit, so it reports the error + // instead of degrading. The dedicated + // `OrderedPartialAggregateStream` emits everything there, so this + // only happens while `enable_migration_aggregate` is off, an option + // that goes away once the migration finishes. + Order::SortedByFirstKey => !params.migration_enabled, + // The stage sees `InputOrderMode::Sorted`, where + // `OrderedPartialAggregateStream` holds one group at a time and + // registers its reservation as unable to handle memory pressure. + // `FairSpillPool` caps the spillable consumers against each other + // but reserves nothing for the others, so whichever consumer took + // the pool first, a final hash stage or a `PartialReduce` stage, + // leaves nothing and this stage is refused the few kilobytes it + // needs. It is the victim rather than the cause. The legacy stream + // is spillable and can emit for a full ordering, so this one only + // happens with the migration enabled. + Order::SortedByAllKeys => params.migration_enabled, + // Linear input: the dedicated and the legacy stream both emit + // their state early and both register as spillable, so the stage + // always survives the pressure. + Order::Unordered => false, + } } const fn shape( @@ -300,14 +316,14 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::Grouped, - legacy_partial_on_partially_sorted_input, + partial_stage_starved_on_ordered_input, ), shape( "partial_coalesce_final", &[Aggregate(Partial), CoalescePartitions, Aggregate(Final)], PARTITIONS, Query::Grouped, - legacy_partial_on_partially_sorted_input, + partial_stage_starved_on_ordered_input, ), shape( "partial_order_preserving_repartition_final", @@ -318,21 +334,21 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::Grouped, - legacy_partial_on_partially_sorted_input, + partial_stage_starved_on_ordered_input, ), shape( "partial_sort_preserving_merge_final", &[Aggregate(Partial), SortPreservingMerge, Aggregate(Final)], PARTITIONS, Query::Grouped, - legacy_partial_on_partially_sorted_input, + partial_stage_starved_on_ordered_input, ), shape( "partial_final_single_partition", &[Aggregate(Partial), Aggregate(Final)], 1, Query::Grouped, - legacy_partial_on_partially_sorted_input, + partial_stage_starved_on_ordered_input, ), shape( "partial_repartition_reduce_repartition_final", @@ -345,7 +361,7 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::Grouped, - legacy_partial_on_partially_sorted_input, + partial_stage_starved_on_ordered_input, ), shape( "partial_repartition_reduce_coalesce_final", @@ -358,7 +374,7 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::Grouped, - legacy_partial_on_partially_sorted_input, + partial_stage_starved_on_ordered_input, ), shape( "partial_local_reduce_repartition_final", @@ -370,7 +386,7 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::Grouped, - legacy_partial_on_partially_sorted_input, + partial_stage_starved_on_ordered_input, ), // ordered PartialReduce has no dedicated stream, lands on the fallback shape( @@ -384,7 +400,7 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::Grouped, - legacy_partial_on_partially_sorted_input, + partial_stage_starved_on_ordered_input, ), shape( "no_grouping_single", @@ -519,7 +535,7 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::MixedKeys, - legacy_partial_on_partially_sorted_input, + partial_stage_starved_on_ordered_input, ), // ordered multi-column group values shape( @@ -531,7 +547,7 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::MixedKeys, - legacy_partial_on_partially_sorted_input, + partial_stage_starved_on_ordered_input, ), shape( "struct_key_single", From 6a5db1040fd81d13126618a9587bb8d052969694 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 7 Sep 2026 13:15:56 +0300 Subject: [PATCH 05/21] remove migration enabled flag --- .../tests/fuzz_cases/aggregate_chain_fuzz.rs | 71 ++-- .../aggregation_fuzzer/AGGREGATE_CHAINS.md | 323 ------------------ 2 files changed, 21 insertions(+), 373 deletions(-) delete mode 100644 datafusion/core/tests/fuzz_cases/aggregation_fuzzer/AGGREGATE_CHAINS.md diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs index e529c11fa2116..f375f8dc4a9b7 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -17,10 +17,6 @@ //! Fuzz test that runs every valid `AggregateExec` chain over the same data //! and asserts identical results. -//! -//! See `aggregation_fuzzer/AGGREGATE_CHAINS.md` for the chain catalogue. Each -//! chain there is a `Shape` below plus a source `Order`. Cases are generated -//! as `shapes × orders × migration flag × cardinality × memory`. use std::num::NonZeroUsize; use std::sync::Arc; @@ -207,7 +203,6 @@ const TOP_K_LIMIT: usize = 2 * ROWS; #[derive(Clone, Copy, Debug, PartialEq)] struct CaseParams { order: Order, - migration_enabled: bool, cardinality: Cardinality, memory: Memory, /// Whether the skip-partial probe may fire. Only varied for shapes with a @@ -240,15 +235,8 @@ fn never_out_of_memory(_params: &CaseParams) -> bool { fn partial_stage_starved_on_ordered_input(params: &CaseParams) -> bool { params.memory == Memory::Limited && match params.order { - // The stage sees `InputOrderMode::PartiallySorted`. The legacy - // `GroupedHashAggregateStream` emits only the groups whose sort - // prefix is complete, and when one prefix value spans the whole - // partition there is nothing it may emit, so it reports the error - // instead of degrading. The dedicated - // `OrderedPartialAggregateStream` emits everything there, so this - // only happens while `enable_migration_aggregate` is off, an option - // that goes away once the migration finishes. - Order::SortedByFirstKey => !params.migration_enabled, + // the impl emit its state early + Order::SortedByFirstKey => false, // The stage sees `InputOrderMode::Sorted`, where // `OrderedPartialAggregateStream` holds one group at a time and // registers its reservation as unable to handle memory pressure. @@ -256,12 +244,9 @@ fn partial_stage_starved_on_ordered_input(params: &CaseParams) -> bool { // but reserves nothing for the others, so whichever consumer took // the pool first, a final hash stage or a `PartialReduce` stage, // leaves nothing and this stage is refused the few kilobytes it - // needs. It is the victim rather than the cause. The legacy stream - // is spillable and can emit for a full ordering, so this one only - // happens with the migration enabled. - Order::SortedByAllKeys => params.migration_enabled, - // Linear input: the dedicated and the legacy stream both emit - // their state early and both register as spillable, so the stage + // needs. It is the victim rather than the cause. + Order::SortedByAllKeys => true, + // Linear input: the impl emit its state early and register as spillable, so the stage // always survives the pressure. Order::Unordered => false, } @@ -680,21 +665,18 @@ fn all_cases() -> Vec { } else { &[true] }; - for migration_enabled in [true, false] { - for cardinality in Cardinality::ALL { - for memory in [Memory::Unlimited, Memory::Limited] { - for &skip_partial_enabled in skip_partial_variants { - cases.push(Case { - shape, - params: CaseParams { - order, - migration_enabled, - cardinality, - memory, - skip_partial_enabled, - }, - }); - } + for cardinality in Cardinality::ALL { + for memory in [Memory::Unlimited, Memory::Limited] { + for &skip_partial_enabled in skip_partial_variants { + cases.push(Case { + shape, + params: CaseParams { + order, + cardinality, + memory, + skip_partial_enabled, + }, + }); } } } @@ -1054,10 +1036,6 @@ fn task_context(case: &Case) -> Arc { let config = SessionConfig::new() .with_batch_size(BATCH_SIZE) .with_target_partitions(PARTITIONS) - .set_bool( - "datafusion.execution.enable_migration_aggregate", - case.params.migration_enabled, - ) // The default is 100k rows. Lower it so the skip-partial probe can // fire on our per-partition row counts. A ratio threshold of 1.0 // disables the probe entirely. @@ -1168,10 +1146,7 @@ fn can_spill(case: &Case, aggregate: &AggregateExec) -> bool { // The dedicated PartialReduce stream emits early and only exists for // Linear input; ordered input or migration off run the legacy stream, // which spills. - PartialReduce => { - !case.params.migration_enabled - || *aggregate.input_order_mode() != InputOrderMode::Linear - } + PartialReduce => *aggregate.input_order_mode() != InputOrderMode::Linear, Partial => false, }; let has_groups = !aggregate.group_expr().is_empty(); @@ -1333,7 +1308,7 @@ async fn run_case_inner(case: &Case, rows: Arc) -> Outcome { } /// Reference result: the single-stage shape of `query` without a limit, one -/// partition, unordered input, unlimited memory and migration on. +/// partition, unordered input, unlimited memory. async fn reference( query: Query, rows: Arc, @@ -1350,7 +1325,6 @@ async fn reference( shape, params: CaseParams { order: Order::Unordered, - migration_enabled: true, cardinality, memory: Memory::Unlimited, skip_partial_enabled: true, @@ -1362,7 +1336,7 @@ async fn reference( match outcome { Outcome::Finished { output, .. } => output, Outcome::OutOfMemory => { - unreachable!("the reference runs with unlimited memory and migration on") + unreachable!("the reference runs with unlimited memory") } } } @@ -1498,10 +1472,9 @@ fn print_cases(cardinality: Cardinality, outcome: &str, cases: &[(Case, Vec, num: Int64>` | row fallback `GroupValuesRows` | unordered only, structs cannot be sorted by the sort kernels | - -Every key column is nullable and carries `cardinality` distinct values, except -the boolean one which has two. - -## Chains: no grouping - -Every stage is `AggregateStream` regardless of mode or ordering. Input ordering -is irrelevant. Nothing spills. - -| Case | Chain | Input requirements | -|---|---|---| -| `no_grouping_single` | `Single` | 1 partition | -| `no_grouping_partial_coalesce_final` | `Partial` → `CoalescePartitionsExec` → `Final` | several partitions. **Default planner output.** | -| `no_grouping_partial_reduce_final` | `Partial` → `CoalescePartitionsExec` → `PartialReduce` → `CoalescePartitionsExec` → `Final` | several partitions, by hand | -| `no_grouping_partial_final_single_partition` | `Partial` → `Final` | 1 partition | - -## Fallback chains - -Any chain above with one of these flipped runs the same shape on -`GroupedHashAggregateStream` for every grouped stage: - -| Case | Trigger | Notes | -|---|---|---| -| migration disabled | `enable_migration_aggregate = false` | Run every grouped chain twice, once per flag value. This is the main old-vs-new comparison. | -| grouping sets | grouping sets present | Order mode is forced to Linear; only the unordered chains are reachable | -| ordered partial reduce | `PartialReduce` with ordered input | `partial_reduce_final_order_preserving` above | - ---- - -## Scenario dimensions - -### Cardinality (distinct group keys divided by input rows) - -| Scenario | Ratio | What it exercises | -|---|---|---| -| very high | about 1.0, nearly all rows unique | skip-partial probe fires in the Linear partial stage; final and single stages spill hardest; Sorted streams emit one row per group | -| high | about 0.5 | skip-partial probe likely fires; large final tables | -| medium | about 0.05 | partial actually reduces; final table moderate | -| low | about 0.001 | few groups, partial reduces almost everything; PartiallySorted streams see many rows per prefix | -| very low | 1 to 10 groups | degenerate hash tables; `PartialReduce` merges to almost nothing | - -Skip-partial only fires in a Linear `Partial` stage, so the very high and high -rows are where the unordered partial chains differ most from the ordered chains. - -### Memory - -| Scenario | Config | Expected behavior | -|---|---|---| -| high | unlimited pool, or pool much larger than the data | no spill, no early emit anywhere | -| medium | `RuntimeEnvBuilder::with_memory_limit(...)` sized so the final or single table cannot fit, disk manager enabled | partial stages emit early; `FinalHashAggregateStream`, `SingleHashAggregateStream` and the PartiallySorted ordered streams spill; Sorted ordered streams are unaffected | - -A resources-exhausted error is never accepted, in either scenario. Every stream -is expected to spill, emit early, or stay bounded, so running out of memory -under the medium pool is treated as a bug in that stream's memory handling or in -how the stages share the pool. The run as a whole must also have spilled -somewhere. Known failures at the time of writing are listed under "Findings". - -For the medium scenario also run once with the disk manager disabled: spilling -streams must return a resources-exhausted error rather than wrong results, and -Sorted ordered streams must still succeed. - -### Skip partial aggregation - -| Scenario | Config | Applies to | -|---|---|---| -| enabled | `skip_partial_aggregation_probe_ratio_threshold = 0.8`, rows threshold lowered to 1024 | chains with a grouped `Partial` stage on Linear input, the only place the probe runs | -| disabled | ratio threshold `1.0`, which turns the probe off | same chains | - -Chains without such a stage run with the probe enabled only, since it can never -fire there. With it disabled the `skipped_aggregation_rows` metric must stay at -zero everywhere. - -### Full matrix - -Every chain in the tables above, plus the migration-disabled fallback over all -grouped ones, times 5 cardinalities, times 2 memory settings, times 2 skip -partial settings where the probe applies. - -Every cell must produce the same result as `single, unordered` with unlimited -memory for the same dataset and query, compared after sorting by group keys. - -## Data generation attributes needed - -To reach each row above the generator must be able to control: - -| Attribute | Values | Selects | -|---|---|---| -| group key columns | 1 or more | PartiallySorted is only reachable with at least 2 keys | -| sort of the source | none, all group keys, or a subset of group keys | Linear, Sorted, PartiallySorted | -| partition count | 1 or several | single versus partitioned chains | -| partitioning of the source | arbitrary or hash by group keys | the `single_partitioned` cases need hash partitioning, either from the source or from a `RepartitionExec` under the aggregate | -| cardinality | see scenario table | | -| batch size | small values, 1 to 64, matter | forces group boundaries to cross batches in the ordered streams | -| aggregate set | must include multi-field state such as `avg`, `count distinct`, `median`, `array_agg` | otherwise `Partial`, `Single` and `PartialReduce` are indistinguishable | -| nulls in group keys | yes or no | null groups are a separate code path in group values | - -## Per-chain checklist - -For each chain above, assert before comparing results: - -1. `explain` shows the expected `mode=` and `ordering_mode=` on every `AggregateExec`. -2. In the medium memory scenario, `spill_count > 0` on the stages listed as spilling and `== 0` on the rest. -3. In very high cardinality with a Linear `Partial` stage, the `skipped_aggregation_rows` metric is `> 0`. -4. Results equal the reference (`single, unordered`, unlimited memory) after sorting by group keys. - -## Findings from running the matrix - -Open: none at the time of writing. The occasional hangs seen earlier under the -limited pool stopped once the source batches were copied instead of sliced (see -the accounting item below); they were only ever observed while every batch was -charged the whole partition and `RepartitionExec` spilled constantly. The -per-case timeout stays in place to catch a recurrence. - -Fixed or documented: - -- `OrderedPartialAggregateStream` refused to emit early on Sorted input and - reported an error instead. It now emits its partial state like it does for - PartiallySorted input, and registers as able to handle memory pressure. -- `PartialReduceHashAggregateStream` registered its reservation as not - spillable although it emits early under pressure, so `FairSpillPool` neither - capped it nor counted it against the spillers, and several of them could hold - the whole pool while the final stage below starved. It now registers as - spillable. -- The legacy `GroupedHashAggregateStream` in `Partial` mode could only emit - completed prefix groups on ordered input and reported an error when there - were none (a boolean or low-cardinality first key spans whole partitions). It - now falls back to emitting everything and restarting its ordering. -- `GroupValuesColumn::emit(EmitTo::All)` replaced the column builders but left - the hash map populated, so the next `intern` dereferenced stale group indices - into empty builders (index out of bounds). The map is now cleared too. -- `RepartitionExec` and the order-preserving merge account a batch by - `get_array_memory_size`, which for a slice is the size of the whole backing - buffers. Memory-table sources hand out 64-row slices of a partition, so every - batch was charged the entire partition and the merge could never reserve one. - The fuzz test now copies each batch, as a real scan would. The accounting - itself is unchanged. - -- The legacy `GroupedHashAggregateStream` with a single nested group key - (`Struct`, `Map`) produced duplicate groups after spilling. When it switches - to merging the sorted spill files it relies on `GroupOrderingFull`, which - requires group ids in first-seen order, and it recreated the group values - collector to guarantee that only when there was more than one group column. - A single nested column has no specialized single-column collector and is - handled by `GroupValuesColumn` through a row-backed column, whose vectorized - interning assigns ids out of input order, so the ordering emitted a group - that was still in progress and the next batch reopened it as a new group. - Fixed by always recreating the collector for the merge phase. - -- `PrimitiveDistinctCountGroupsAccumulator` reported its capacity in `size()` - but only cleared its buffers when emitting everything, so a table looked as - large after `take_state_batch` as before. Every dedicated stream then failed - its post-emit or post-spill resize, and the legacy stream could not reserve - its sort headroom. Fixed by releasing the buffers on `EmitTo::All`. -- `PartialReduceHashAggregateStream` registers its reservation as not - spillable, so `FairSpillPool` neither caps it nor counts it against the - spillers, and several of them can hold the pool while the final stage below - starves. -- When a final or single stream is starved with nothing reserved, its - post-spill `try_resize` is really a grow from zero, and the error is reported - as "Decreasing allocation after spilling should succeed". From 316a15cc9154a195632515426c2d4545146c7d79 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:32:33 +0300 Subject: [PATCH 06/21] mark partial reduce as unordered --- .../tests/fuzz_cases/aggregate_chain_fuzz.rs | 19 ++- datafusion/core/tests/memory_limit/mod.rs | 145 +++++++++--------- 2 files changed, 84 insertions(+), 80 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs index f375f8dc4a9b7..9a1d9ca12681a 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -1112,6 +1112,14 @@ fn expected_orders(shape: &Shape, source_order: Order) -> Vec { for operator in shape.operators { match operator { HashRepartition | CoalescePartitions => current = Order::Unordered, + // `AggregateExec::try_new` forces `InputOrderMode::Linear` for + // partial reduce, since it emits its groups in hash table order, + // and it advertises no output ordering either. Everything above it + // is unordered until something sorts again. + Aggregate(PartialReduce) => { + expected.push(Order::Unordered); + current = Order::Unordered; + } Aggregate(_) | TopK(_) => expected.push(current), OrderPreservingHashRepartition | SortPreservingMerge => {} } @@ -1136,18 +1144,15 @@ fn order_matches(query: Query, expected: Order, actual: &InputOrderMode) -> bool /// Whether this stage's stream is allowed to spill. See the memory table in /// AGGREGATE_CHAINS.md. -fn can_spill(case: &Case, aggregate: &AggregateExec) -> bool { +fn can_spill(aggregate: &AggregateExec) -> bool { if aggregate.limit_options().is_some() { // GroupedTopKAggregateStream keeps a bounded heap and never spills return false; } let spilling_mode = match aggregate.mode() { Final | FinalPartitioned | Single | SinglePartitioned => true, - // The dedicated PartialReduce stream emits early and only exists for - // Linear input; ordered input or migration off run the legacy stream, - // which spills. - PartialReduce => *aggregate.input_order_mode() != InputOrderMode::Linear, - Partial => false, + // Both partial streams emit their state early instead of spilling. + PartialReduce | Partial => false, }; let has_groups = !aggregate.group_expr().is_empty(); spilling_mode && has_groups && *aggregate.input_order_mode() != InputOrderMode::Sorted @@ -1208,7 +1213,7 @@ fn check_metrics(case: &Case, plan: &Arc) -> Vec { // Whether a spilling-capable stage actually spills depends on // the pool geometry, so only the run-wide coverage check in the // driver requires it. Streams that cannot spill must not. - if !can_spill(case, aggregate) { + if !can_spill(aggregate) { assert_eq!(spill_count, 0, "{case:?}: {mode:?} must never spill"); } } diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index bdd5211c4f7a4..6d8303b87a92e 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -28,7 +28,7 @@ mod union_nullable_spill; mod view_spill_compaction; use arrow::array::{ ArrayRef, DictionaryArray, Int32Array, Int64Array, Int64Builder, ListBuilder, - RecordBatch, StringViewArray, StructArray, StringArray + RecordBatch, StringArray, StringViewArray, StructArray, }; use arrow::buffer::NullBuffer; use arrow::compute::SortOptions; @@ -286,7 +286,6 @@ async fn group_by_count_distinct_utf8_view() { .await } - /// `GROUP BY` on a single nested key in the legacy `GroupedHashAggregateStream` /// under a memory limit. /// @@ -308,57 +307,57 @@ const NESTED_KEY_BATCH_ROWS: usize = 8_192; const NESTED_KEY_MEMORY_LIMIT: usize = 4 * 1024 * 1024; fn nested_key_struct_fields() -> Fields { - Fields::from(vec![ - Field::new("list", DataType::new_list(DataType::Int64, true), true), - Field::new("num", DataType::Int64, true), - ]) + Fields::from(vec![ + Field::new("list", DataType::new_list(DataType::Int64, true), true), + Field::new("num", DataType::Int64, true), + ]) } /// `st` is mostly `{list: [g, g + 1], num: g}` for group `g`, with a sprinkle /// of null lists, empty lists, null nums and null structs so that keys of /// different shapes meet in the same batches. `v` is unique. fn nested_key_table() -> MemTable { - let schema = Arc::new(Schema::new(vec![ - Field::new_struct("st", nested_key_struct_fields(), true), - Field::new("v", DataType::Int64, false), - ])); - let batches = (0..NESTED_KEY_ROWS) - .step_by(NESTED_KEY_BATCH_ROWS) - .map(|start| { - let rows = start..(start + NESTED_KEY_BATCH_ROWS).min(NESTED_KEY_ROWS); - let mut list = ListBuilder::new(Int64Builder::new()); - let mut num = Vec::with_capacity(rows.len()); - let mut valid = Vec::with_capacity(rows.len()); - for row in rows.clone() { - let group = row as i64 % NESTED_KEY_GROUPS; - match row % 37 { - 0 => list.append_null(), - 1 => list.append(true), - _ => { - list.values().append_value(group); - list.values().append_value(group + 1); - list.append(true); - } - } - num.push((row % 41 != 0).then_some(group)); - valid.push(row % 43 != 0); - } - let st = StructArray::new( - nested_key_struct_fields(), - vec![Arc::new(list.finish()), Arc::new(Int64Array::from(num))], - Some(NullBuffer::from(valid)), - ); - RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(st), - Arc::new(Int64Array::from_iter_values(rows.map(|row| row as i64))), - ], - ) - .unwrap() - }) - .collect(); - MemTable::try_new(schema, vec![batches]).unwrap() + let schema = Arc::new(Schema::new(vec![ + Field::new_struct("st", nested_key_struct_fields(), true), + Field::new("v", DataType::Int64, false), + ])); + let batches = (0..NESTED_KEY_ROWS) + .step_by(NESTED_KEY_BATCH_ROWS) + .map(|start| { + let rows = start..(start + NESTED_KEY_BATCH_ROWS).min(NESTED_KEY_ROWS); + let mut list = ListBuilder::new(Int64Builder::new()); + let mut num = Vec::with_capacity(rows.len()); + let mut valid = Vec::with_capacity(rows.len()); + for row in rows.clone() { + let group = row as i64 % NESTED_KEY_GROUPS; + match row % 37 { + 0 => list.append_null(), + 1 => list.append(true), + _ => { + list.values().append_value(group); + list.values().append_value(group + 1); + list.append(true); + } + } + num.push((row % 41 != 0).then_some(group)); + valid.push(row % 43 != 0); + } + let st = StructArray::new( + nested_key_struct_fields(), + vec![Arc::new(list.finish()), Arc::new(Int64Array::from(num))], + Some(NullBuffer::from(valid)), + ); + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(st), + Arc::new(Int64Array::from_iter_values(rows.map(|row| row as i64))), + ], + ) + .unwrap() + }) + .collect(); + MemTable::try_new(schema, vec![batches]).unwrap() } const NESTED_KEY_QUERY: &str = "select st, count(v), count(distinct v), sum(v), avg(v), min(v), max(v) \ @@ -366,36 +365,36 @@ const NESTED_KEY_QUERY: &str = "select st, count(v), count(distinct v), sum(v), /// Runs the query on the legacy stream, with or without a memory limit. async fn run_nested_key_query(memory_limit: Option) -> String { - let mut runtime = - RuntimeEnvBuilder::new().with_disk_manager_builder(DiskManagerBuilder::default()); - if let Some(limit) = memory_limit { - runtime = runtime.with_memory_pool(Arc::new(FairSpillPool::new(limit))); - } - let config = SessionConfig::new() - .with_target_partitions(4) - // small batches: the merged spill stream arrives in many batches and - // groups span batch boundaries - .with_batch_size(64) - .set_bool("datafusion.execution.enable_migration_aggregate", false); - let ctx = SessionContext::new_with_config_rt(config, runtime.build_arc().unwrap()); - ctx.register_table("t", Arc::new(nested_key_table())) - .unwrap(); - let batches = ctx - .sql(NESTED_KEY_QUERY) - .await - .unwrap() - .collect() - .await - .unwrap(); - batches_to_sort_string(&batches) + let mut runtime = + RuntimeEnvBuilder::new().with_disk_manager_builder(DiskManagerBuilder::default()); + if let Some(limit) = memory_limit { + runtime = runtime.with_memory_pool(Arc::new(FairSpillPool::new(limit))); + } + let config = SessionConfig::new() + .with_target_partitions(4) + // small batches: the merged spill stream arrives in many batches and + // groups span batch boundaries + .with_batch_size(64) + .set_bool("datafusion.execution.enable_migration_aggregate", false); + let ctx = SessionContext::new_with_config_rt(config, runtime.build_arc().unwrap()); + ctx.register_table("t", Arc::new(nested_key_table())) + .unwrap(); + let batches = ctx + .sql(NESTED_KEY_QUERY) + .await + .unwrap() + .collect() + .await + .unwrap(); + batches_to_sort_string(&batches) } #[tokio::test] async fn legacy_stream_nested_key_spill_keeps_groups_unique() { - let expected = run_nested_key_query(None).await; - let actual = run_nested_key_query(Some(NESTED_KEY_MEMORY_LIMIT)).await; - // A duplicated group shows up as extra rows with the counts split - assert_eq!(actual, expected); + let expected = run_nested_key_query(None).await; + let actual = run_nested_key_query(Some(NESTED_KEY_MEMORY_LIMIT)).await; + // A duplicated group shows up as extra rows with the counts split + assert_eq!(actual, expected); } #[tokio::test] From ae44d67cb3fd2639ba5ec1eb0b7c0cae9492b781 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:46:50 +0300 Subject: [PATCH 07/21] remove allowing oom --- .../tests/fuzz_cases/aggregate_chain_fuzz.rs | 181 +++--------------- .../physical-plan/src/aggregates/mod.rs | 61 ++++++ .../src/aggregates/partial_reduce_stream.rs | 4 + 3 files changed, 89 insertions(+), 157 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs index 9a1d9ca12681a..d9896b2b321a5 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -32,8 +32,8 @@ use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef, SortOptions}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::prelude::SessionConfig; +use datafusion_common::Result; use datafusion_common::test_util::batches_to_sort_string; -use datafusion_common::{DataFusionError, Result}; use datafusion_common_runtime::JoinSet; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{FairSpillPool, TrackConsumersPool}; @@ -218,38 +218,6 @@ struct Shape { /// Source partition count. source_partitions: usize, query: Query, - /// Whether a resources-exhausted error is an accepted outcome for this - /// chain instead of a result, given the rest of the case parameters. - /// Every shape declares its own, so nothing is accepted by derivation. - accepts_out_of_memory: fn(&CaseParams) -> bool, -} - -/// A chain never runs out of memory: every stage spills, emits early, or is -/// bounded. -fn never_out_of_memory(_params: &CaseParams) -> bool { - false -} - -/// The chain starts with a grouped `Partial` aggregate reading the source, so -/// with ordered input that stage cannot always survive memory pressure. -fn partial_stage_starved_on_ordered_input(params: &CaseParams) -> bool { - params.memory == Memory::Limited - && match params.order { - // the impl emit its state early - Order::SortedByFirstKey => false, - // The stage sees `InputOrderMode::Sorted`, where - // `OrderedPartialAggregateStream` holds one group at a time and - // registers its reservation as unable to handle memory pressure. - // `FairSpillPool` caps the spillable consumers against each other - // but reserves nothing for the others, so whichever consumer took - // the pool first, a final hash stage or a `PartialReduce` stage, - // leaves nothing and this stage is refused the few kilobytes it - // needs. It is the victim rather than the cause. - Order::SortedByAllKeys => true, - // Linear input: the impl emit its state early and register as spillable, so the stage - // always survives the pressure. - Order::Unordered => false, - } } const fn shape( @@ -257,40 +225,30 @@ const fn shape( operators: &'static [Operator], source_partitions: usize, query: Query, - accepts_out_of_memory: fn(&CaseParams) -> bool, ) -> Shape { Shape { name, operators, source_partitions, query, - accepts_out_of_memory, } } /// Every shape from AGGREGATE_CHAINS.md. The ordering variants there come from /// crossing a shape with `Order`, so one entry here covers several rows. const SHAPES: &[Shape] = &[ - shape( - "single", - &[Aggregate(Single)], - 1, - Query::Grouped, - never_out_of_memory, - ), + shape("single", &[Aggregate(Single)], 1, Query::Grouped), shape( "single_partitioned", &[HashRepartition, Aggregate(SinglePartitioned)], PARTITIONS, Query::Grouped, - never_out_of_memory, ), shape( "single_partitioned_order_preserving", &[OrderPreservingHashRepartition, Aggregate(SinglePartitioned)], PARTITIONS, Query::Grouped, - never_out_of_memory, ), shape( "partial_repartition_final", @@ -301,14 +259,12 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::Grouped, - partial_stage_starved_on_ordered_input, ), shape( "partial_coalesce_final", &[Aggregate(Partial), CoalescePartitions, Aggregate(Final)], PARTITIONS, Query::Grouped, - partial_stage_starved_on_ordered_input, ), shape( "partial_order_preserving_repartition_final", @@ -319,21 +275,18 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::Grouped, - partial_stage_starved_on_ordered_input, ), shape( "partial_sort_preserving_merge_final", &[Aggregate(Partial), SortPreservingMerge, Aggregate(Final)], PARTITIONS, Query::Grouped, - partial_stage_starved_on_ordered_input, ), shape( "partial_final_single_partition", &[Aggregate(Partial), Aggregate(Final)], 1, Query::Grouped, - partial_stage_starved_on_ordered_input, ), shape( "partial_repartition_reduce_repartition_final", @@ -346,7 +299,6 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::Grouped, - partial_stage_starved_on_ordered_input, ), shape( "partial_repartition_reduce_coalesce_final", @@ -359,7 +311,6 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::Grouped, - partial_stage_starved_on_ordered_input, ), shape( "partial_local_reduce_repartition_final", @@ -371,7 +322,6 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::Grouped, - partial_stage_starved_on_ordered_input, ), // ordered PartialReduce has no dedicated stream, lands on the fallback shape( @@ -385,51 +335,34 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::Grouped, - partial_stage_starved_on_ordered_input, ), shape( "no_grouping_single", &[Aggregate(Single)], 1, Query::NoGrouping, - never_out_of_memory, ), // TopK: same query without a limit is the reference for the TopK chains - shape( - "top_k_query_single", - &[Aggregate(Single)], - 1, - Query::TopK, - never_out_of_memory, - ), - shape( - "top_k_single", - &[TopK(Single)], - 1, - Query::TopK, - never_out_of_memory, - ), + shape("top_k_query_single", &[Aggregate(Single)], 1, Query::TopK), + shape("top_k_single", &[TopK(Single)], 1, Query::TopK), // planner shape: the limit lands on the aggregate under the sort shape( "top_k_partial_repartition_final", &[Aggregate(Partial), HashRepartition, TopK(FinalPartitioned)], PARTITIONS, Query::TopK, - never_out_of_memory, ), shape( "top_k_partial_coalesce_final", &[Aggregate(Partial), CoalescePartitions, TopK(Final)], PARTITIONS, Query::TopK, - never_out_of_memory, ), shape( "top_k_both_stages", &[TopK(Partial), HashRepartition, TopK(FinalPartitioned)], PARTITIONS, Query::TopK, - never_out_of_memory, ), // group key types: single stage and the default two-stage planner shape shape( @@ -437,7 +370,6 @@ const SHAPES: &[Shape] = &[ &[Aggregate(Single)], 1, Query::BooleanKey, - never_out_of_memory, ), shape( "boolean_key_partial_repartition_final", @@ -448,15 +380,8 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::BooleanKey, - never_out_of_memory, - ), - shape( - "bytes_key_single", - &[Aggregate(Single)], - 1, - Query::BytesKey, - never_out_of_memory, ), + shape("bytes_key_single", &[Aggregate(Single)], 1, Query::BytesKey), shape( "bytes_key_partial_repartition_final", &[ @@ -466,14 +391,12 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::BytesKey, - never_out_of_memory, ), shape( "bytes_view_key_single", &[Aggregate(Single)], 1, Query::BytesViewKey, - never_out_of_memory, ), shape( "bytes_view_key_partial_repartition_final", @@ -484,14 +407,12 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::BytesViewKey, - never_out_of_memory, ), shape( "primitive_key_single", &[Aggregate(Single)], 1, Query::PrimitiveKey, - never_out_of_memory, ), shape( "primitive_key_partial_repartition_final", @@ -502,14 +423,12 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::PrimitiveKey, - never_out_of_memory, ), shape( "mixed_keys_single", &[Aggregate(Single)], 1, Query::MixedKeys, - never_out_of_memory, ), shape( "mixed_keys_partial_repartition_final", @@ -520,7 +439,6 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::MixedKeys, - partial_stage_starved_on_ordered_input, ), // ordered multi-column group values shape( @@ -532,14 +450,12 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::MixedKeys, - partial_stage_starved_on_ordered_input, ), shape( "struct_key_single", &[Aggregate(Single)], 1, Query::StructKey, - never_out_of_memory, ), shape( "struct_key_partial_repartition_final", @@ -550,14 +466,12 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::StructKey, - never_out_of_memory, ), shape( "no_grouping_partial_coalesce_final", &[Aggregate(Partial), CoalescePartitions, Aggregate(Final)], PARTITIONS, Query::NoGrouping, - never_out_of_memory, ), shape( "no_grouping_partial_reduce_final", @@ -570,14 +484,12 @@ const SHAPES: &[Shape] = &[ ], PARTITIONS, Query::NoGrouping, - never_out_of_memory, ), shape( "no_grouping_partial_final_single_partition", &[Aggregate(Partial), Aggregate(Final)], 1, Query::NoGrouping, - never_out_of_memory, ), ]; @@ -618,14 +530,6 @@ struct Case { params: CaseParams, } -impl Case { - /// Whether a resources-exhausted error is accepted for this case instead - /// of a result, as declared by its shape. - fn accepts_out_of_memory(&self) -> bool { - (self.shape.accepts_out_of_memory)(&self.params) - } -} - impl Shape { /// Whether some `Partial` stage of this shape runs the skip-partial probe /// for the given source order: grouped, not TopK, and Linear input. @@ -1243,23 +1147,17 @@ fn check_metrics(case: &Case, plan: &Arc) -> Vec { // Driver // --------------------------------------------------------------------------- -enum Outcome { - /// Sorted output plus the stages that spilled, empty if none did. - Finished { - output: String, - spilled: Vec, - }, - /// Ran out of memory, and the case was built to accept that. See - /// [`Shape::accepts_out_of_memory`]. - OutOfMemory, +/// Sorted output plus the stages that spilled, empty if none did. +struct Outcome { + output: String, + spilled: Vec, } /// Runs one case, checks plan shape and metrics, and returns its outcome. /// -/// Running out of memory fails the run unless the case was built with -/// `accepts_out_of_memory`: every other stream either spills, emits early, or -/// is bounded, so an error there is a bug in a stream's memory handling or in -/// how stages share the pool. +/// Running out of memory is never accepted: every stream either spills, emits +/// early, or is bounded, so an error there is a bug in a stream's memory +/// handling or in how the stages share the pool. async fn run_case(case: Case, rows: Arc) -> Outcome { log::debug!("start {case:?}"); let outcome = run_case_inner(&case, rows).await; @@ -1294,19 +1192,13 @@ async fn run_case_inner(case: &Case, rows: Arc) -> Outcome { }); let batches = match collected { Ok(batches) => batches, - Err(error) - if matches!(error.find_root(), DataFusionError::ResourcesExhausted(_)) - && case.accepts_out_of_memory() => - { - return Outcome::OutOfMemory; - } Err(error) => panic!( "{case:?} failed: {error}\n{}", displayable(plan.as_ref()).indent(true) ), }; let spilled = check_metrics(case, &plan); - Outcome::Finished { + Outcome { output: batches_to_sort_string(&batches), spilled, } @@ -1338,12 +1230,7 @@ async fn reference( rows, ) .await; - match outcome { - Outcome::Finished { output, .. } => output, - Outcome::OutOfMemory => { - unreachable!("the reference runs with unlimited memory") - } - } + outcome.output } #[tokio::test(flavor = "multi_thread")] @@ -1368,7 +1255,7 @@ async fn aggregate_chain_fuzz() { } let mut join_set = JoinSet::new(); - let (mut spilled, mut finished, mut out_of_memory) = (vec![], vec![], vec![]); + let (mut spilled, mut finished) = (vec![], vec![]); for case in all_cases() .into_iter() .filter(|case| case.params.cardinality == cardinality) @@ -1386,38 +1273,22 @@ async fn aggregate_chain_fuzz() { &mut join_set, &mut spilled, &mut finished, - &mut out_of_memory, &mut failures, ) .await; } join_set.spawn(async move { - match run_case(case.clone(), rows).await { - Outcome::Finished { output, spilled } => { - assert_eq!(output, expected, "{case:?} (seed {seed})"); - (case, Some(spilled)) - } - Outcome::OutOfMemory => (case, None), - } + let outcome = run_case(case.clone(), rows).await; + assert_eq!(outcome.output, expected, "{case:?} (seed {seed})"); + (case, outcome.spilled) }); } while !join_set.is_empty() { - collect_finished( - &mut join_set, - &mut spilled, - &mut finished, - &mut out_of_memory, - &mut failures, - ) - .await; + collect_finished(&mut join_set, &mut spilled, &mut finished, &mut failures) + .await; } print_cases(cardinality, "spilled", &spilled); print_cases(cardinality, "finished without spilling", &finished); - print_cases( - cardinality, - "ran out of memory (accepted, legacy partial stream)", - &out_of_memory, - ); total_spilled += spilled.len(); } // A shape filter may select only shapes that cannot spill @@ -1441,19 +1312,17 @@ const CASE_TIMEOUT_SECS: u64 = 60; /// Waits for one case and files it under spilled, finished or failed. A /// failure does not stop the run, so one run reports every failing case. async fn collect_finished( - join_set: &mut JoinSet<(Case, Option>)>, + join_set: &mut JoinSet<(Case, Vec)>, spilled: &mut Vec<(Case, Vec)>, finished: &mut Vec<(Case, Vec)>, - out_of_memory: &mut Vec<(Case, Vec)>, failures: &mut Vec, ) { let Some(result) = join_set.join_next().await else { return; }; match result { - Ok((case, None)) => out_of_memory.push((case, vec![])), - Ok((case, Some(stages))) if stages.is_empty() => finished.push((case, stages)), - Ok((case, Some(stages))) => spilled.push((case, stages)), + Ok((case, stages)) if stages.is_empty() => finished.push((case, stages)), + Ok((case, stages)) => spilled.push((case, stages)), Err(error) => failures.push(error.to_string()), } } @@ -1517,8 +1386,6 @@ async fn run_single_case( rows, ) .await; - if let Outcome::Finished { output, .. } = actual { - assert_eq!(output, expected); - } + assert_eq!(actual.output, expected); Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index ba08c4b003195..1857ceb4d930b 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1425,6 +1425,19 @@ impl AggregateExec { .equivalence_properties() .project(group_expr_mapping, schema); + // An aggregation that does not maintain its input order must not + // advertise the input's ordering either. `maintains_input_order` + // reports `false` for `InputOrderMode::Linear`, and a hash aggregation + // emits its groups in hash table order. `PartialReduce` is forced to + // `Linear` in `try_new` for exactly this reason, but the projection + // above still carries the input orderings over, so drop them here. + // Otherwise a consumer may take the early emit path on an ordering the + // aggregation does not keep, and flush a group before all of its rows + // have arrived. + if *input_order_mode == InputOrderMode::Linear { + eq_properties.clear_orderings(); + } + // True no-group aggregates produce only one row in each output // partition, so aggregate outputs are constants within the partition. // Grouping sets with empty grouping expressions are not covered here: @@ -4595,6 +4608,54 @@ mod tests { Ok(()) } + /// A `PartialReduce` aggregation emits its groups in hash table order, and + /// `try_new` forces its [`InputOrderMode`] to `Linear` to say so. It must + /// not advertise its input's ordering as its own output ordering either: + /// a consumer that believed it could take the early emit path and flush a + /// group before all of its rows had arrived. + #[test] + fn partial_reduce_does_not_advertise_input_ordering() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Float64, false), + ])); + let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new( + Column::new("a", 0), + ))]) + .unwrap(); + let input = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let input = Arc::new(TestMemoryExec::update_cache(&Arc::new(input))); + assert!( + input.properties().output_ordering().is_some(), + "test setup: the input is ordered by the group key" + ); + + let partial_reduce = AggregateExec::try_new( + AggregateMode::PartialReduce, + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]), + vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(b)") + .build()?, + )], + vec![None], + input, + Arc::clone(&schema), + )?; + + assert_eq!(partial_reduce.input_order_mode(), &InputOrderMode::Linear); + assert_eq!(partial_reduce.maintains_input_order(), vec![false]); + assert!( + partial_reduce.properties().output_ordering().is_none(), + "partial reduce advertised an ordering it does not maintain: {:?}", + partial_reduce.properties().output_ordering() + ); + + Ok(()) + } + fn partial_reduce_test_aggregate() -> Result { partial_reduce_test_aggregate_with_batches(1) } diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index d8f1447bc9521..6c3cd622f6bfd 100644 --- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs +++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs @@ -203,6 +203,10 @@ impl PartialReduceHashAggregateStream { let reservation = MemoryConsumer::new(format!("PartialReduceHashAggregateStream[{partition}]")) + // We interpret 'can spill' as 'can handle memory back pressure'. + // This value needs to be set to true for the default memory pool implementations + // to ensure fair application of back pressure amongst the memory consumers. + .with_can_spill(true) .register(context.memory_pool()); Ok(Self { From be1bd6c0bafbb70cec6be0a17d57aab63b9c9af6 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:12:36 +0300 Subject: [PATCH 08/21] update comment --- datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs index d9896b2b321a5..0758f53c2bd09 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -234,8 +234,7 @@ const fn shape( } } -/// Every shape from AGGREGATE_CHAINS.md. The ordering variants there come from -/// crossing a shape with `Order`, so one entry here covers several rows. +// Test cases const SHAPES: &[Shape] = &[ shape("single", &[Aggregate(Single)], 1, Query::Grouped), shape( @@ -1046,8 +1045,7 @@ fn order_matches(query: Query, expected: Order, actual: &InputOrderMode) -> bool } } -/// Whether this stage's stream is allowed to spill. See the memory table in -/// AGGREGATE_CHAINS.md. +/// Whether this stage's stream is allowed to spill. fn can_spill(aggregate: &AggregateExec) -> bool { if aggregate.limit_options().is_some() { // GroupedTopKAggregateStream keeps a bounded heap and never spills From c1faba2514b575a6033c3b942918e7b7883d7401 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 10 Sep 2026 11:02:15 +0300 Subject: [PATCH 09/21] Fix formatting and raise fuzz case timeout for slow CI runners The per-case timeout of 60s fired on the four-core CI runner for the order-preserving repartition shape at very high cardinality. The cases take about two seconds alone locally and pass under the full concurrent matrix, so this is CI oversubscription rather than a hang. --- datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs | 7 +++++-- .../physical-plan/src/aggregates/partial_reduce_stream.rs | 8 ++++---- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs index 0758f53c2bd09..b28d88445fe6b 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -1304,8 +1304,11 @@ async fn aggregate_chain_fuzz() { /// One line per case; `spilled_stages` names the aggregate operators that /// spilled and flags when more than one did. const MAX_CONCURRENT_CASES: usize = 16; -/// Generous: a case takes well under a second in debug builds. -const CASE_TIMEOUT_SECS: u64 = 60; +/// A case takes about two seconds alone in a debug build, but CI runs the +/// whole fuzz binary on a four-core runner with `MAX_CONCURRENT_CASES` of +/// them in flight, and has taken over a minute per case there. Generous, so +/// only a real hang fires it. +const CASE_TIMEOUT_SECS: u64 = 600; /// Waits for one case and files it under spilled, finished or failed. A /// failure does not stop the run, so one run reports every failing case. diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index 6c3cd622f6bfd..24db34c03cbd7 100644 --- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs +++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs @@ -203,10 +203,10 @@ impl PartialReduceHashAggregateStream { let reservation = MemoryConsumer::new(format!("PartialReduceHashAggregateStream[{partition}]")) - // We interpret 'can spill' as 'can handle memory back pressure'. - // This value needs to be set to true for the default memory pool implementations - // to ensure fair application of back pressure amongst the memory consumers. - .with_can_spill(true) + // We interpret 'can spill' as 'can handle memory back pressure'. + // This value needs to be set to true for the default memory pool implementations + // to ensure fair application of back pressure amongst the memory consumers. + .with_can_spill(true) .register(context.memory_pool()); Ok(Self { From dc6977818129ceda3d64d1707f71921bc101d071 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 10 Sep 2026 12:17:27 +0300 Subject: [PATCH 10/21] Fix formatting and keep fuzz case timeout off slow CI runners The per-case timeout of 60s fired on the four-core CI runner for the order-preserving repartition shape at very high cardinality. The cases take about two seconds alone locally and pass under the full concurrent matrix, so this is CI oversubscription rather than a hang. Bound the in-flight cases by the available parallelism instead of a fixed 16, and raise the timeout so only a real hang fires it. --- .../tests/fuzz_cases/aggregate_chain_fuzz.rs | 19 +-- datafusion/core/tests/memory_limit/mod.rs | 111 ------------------ 2 files changed, 10 insertions(+), 120 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs index b28d88445fe6b..cdc9b98d95775 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -34,6 +34,7 @@ use datafusion::datasource::source::DataSourceExec; use datafusion::prelude::SessionConfig; use datafusion_common::Result; use datafusion_common::test_util::batches_to_sort_string; +use datafusion_common::utils::get_available_parallelism; use datafusion_common_runtime::JoinSet; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{FairSpillPool, TrackConsumersPool}; @@ -1237,6 +1238,10 @@ async fn aggregate_chain_fuzz() { println!("aggregate_chain_fuzz seed = {seed}"); let mut total_spilled = 0; let mut failures: Vec = vec![]; + // Every in-flight case holds several copies of the dataset and its own + // partitioned streams, so bound the concurrency by the cores at hand + // instead of spawning the whole matrix. + let max_concurrent_cases = get_available_parallelism(); for cardinality in Cardinality::ALL { let rows = Arc::new(generate_rows(cardinality, seed)); @@ -1264,9 +1269,7 @@ async fn aggregate_chain_fuzz() { .find(|(query, _)| *query == case.shape.query) .map(|(_, expected)| expected.clone()) .unwrap(); - // Every in-flight case holds several copies of the dataset, so - // bound the concurrency instead of spawning the whole matrix. - while join_set.len() >= MAX_CONCURRENT_CASES { + while join_set.len() >= max_concurrent_cases { collect_finished( &mut join_set, &mut spilled, @@ -1301,13 +1304,9 @@ async fn aggregate_chain_fuzz() { ); } -/// One line per case; `spilled_stages` names the aggregate operators that -/// spilled and flags when more than one did. -const MAX_CONCURRENT_CASES: usize = 16; /// A case takes about two seconds alone in a debug build, but CI runs the -/// whole fuzz binary on a four-core runner with `MAX_CONCURRENT_CASES` of -/// them in flight, and has taken over a minute per case there. Generous, so -/// only a real hang fires it. +/// whole fuzz binary on a four-core runner, and has taken over a minute per +/// case there. Generous, so only a real hang fires it. const CASE_TIMEOUT_SECS: u64 = 600; /// Waits for one case and files it under spilled, finished or failed. A @@ -1328,6 +1327,8 @@ async fn collect_finished( } } +/// One line per case; `spilled_stages` names the aggregate operators that +/// spilled and flags when more than one did. fn print_cases(cardinality: Cardinality, outcome: &str, cases: &[(Case, Vec)]) { let mut lines: Vec = cases .iter() diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index fdc3d545b7966..1636b676ba023 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -409,117 +409,6 @@ async fn group_by_count_distinct_utf8_view() { .await } -/// `GROUP BY` on a single nested key in the legacy `GroupedHashAggregateStream` -/// under a memory limit. -/// -/// After spilling, the legacy stream re-aggregates the merged spill files with -/// `GroupOrderingFull`, which requires group ids in first-seen order. A single -/// nested key has no specialized single-column group values implementation and -/// is handled by the multi-column one, whose vectorized interning does not -/// guarantee that order. If the stream keeps that implementation for the merge -/// phase, groups are emitted while still in progress and come out duplicated. -const NESTED_KEY_ROWS: usize = 200_000; -const NESTED_KEY_GROUPS: i64 = 16; -const NESTED_KEY_BATCH_ROWS: usize = 8_192; - -/// Small enough that the final stages must spill their `count(distinct)` -/// state, large enough for the merge of the spilled runs. A `FairSpillPool` -/// caps every stage at its share, so the partial stages emit early and the -/// final stages spill repeatedly, which produces the many small merged batches -/// this bug needs. -const NESTED_KEY_MEMORY_LIMIT: usize = 4 * 1024 * 1024; - -fn nested_key_struct_fields() -> Fields { - Fields::from(vec![ - Field::new("list", DataType::new_list(DataType::Int64, true), true), - Field::new("num", DataType::Int64, true), - ]) -} - -/// `st` is mostly `{list: [g, g + 1], num: g}` for group `g`, with a sprinkle -/// of null lists, empty lists, null nums and null structs so that keys of -/// different shapes meet in the same batches. `v` is unique. -fn nested_key_table() -> MemTable { - let schema = Arc::new(Schema::new(vec![ - Field::new_struct("st", nested_key_struct_fields(), true), - Field::new("v", DataType::Int64, false), - ])); - let batches = (0..NESTED_KEY_ROWS) - .step_by(NESTED_KEY_BATCH_ROWS) - .map(|start| { - let rows = start..(start + NESTED_KEY_BATCH_ROWS).min(NESTED_KEY_ROWS); - let mut list = ListBuilder::new(Int64Builder::new()); - let mut num = Vec::with_capacity(rows.len()); - let mut valid = Vec::with_capacity(rows.len()); - for row in rows.clone() { - let group = row as i64 % NESTED_KEY_GROUPS; - match row % 37 { - 0 => list.append_null(), - 1 => list.append(true), - _ => { - list.values().append_value(group); - list.values().append_value(group + 1); - list.append(true); - } - } - num.push((row % 41 != 0).then_some(group)); - valid.push(row % 43 != 0); - } - let st = StructArray::new( - nested_key_struct_fields(), - vec![Arc::new(list.finish()), Arc::new(Int64Array::from(num))], - Some(NullBuffer::from(valid)), - ); - RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(st), - Arc::new(Int64Array::from_iter_values(rows.map(|row| row as i64))), - ], - ) - .unwrap() - }) - .collect(); - MemTable::try_new(schema, vec![batches]).unwrap() -} - -const NESTED_KEY_QUERY: &str = "select st, count(v), count(distinct v), sum(v), avg(v), min(v), max(v) \ - from t group by st"; - -/// Runs the query on the legacy stream, with or without a memory limit. -async fn run_nested_key_query(memory_limit: Option) -> String { - let mut runtime = - RuntimeEnvBuilder::new().with_disk_manager_builder(DiskManagerBuilder::default()); - if let Some(limit) = memory_limit { - runtime = runtime.with_memory_pool(Arc::new(FairSpillPool::new(limit))); - } - let config = SessionConfig::new() - .with_target_partitions(4) - // small batches: the merged spill stream arrives in many batches and - // groups span batch boundaries - .with_batch_size(64) - .set_bool("datafusion.execution.enable_migration_aggregate", false); - let ctx = SessionContext::new_with_config_rt(config, runtime.build_arc().unwrap()); - ctx.register_table("t", Arc::new(nested_key_table())) - .unwrap(); - let batches = ctx - .sql(NESTED_KEY_QUERY) - .await - .unwrap() - .collect() - .await - .unwrap(); - batches_to_sort_string(&batches) -} - -#[tokio::test] -async fn legacy_stream_nested_key_spill_keeps_groups_unique() { - let expected = run_nested_key_query(None).await; - let actual = run_nested_key_query(Some(NESTED_KEY_MEMORY_LIMIT)).await; - // A duplicated group shows up as extra rows with the counts split - assert_eq!(actual, expected); -} - #[tokio::test] async fn join_by_key_multiple_partitions() { let config = SessionConfig::new().with_target_partitions(2); From c8006167b1cc6ab7e22a386e323bcc7277da592a Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 10 Sep 2026 12:56:38 +0300 Subject: [PATCH 11/21] Cover group by without aggregate expressions in aggregate chain fuzz Adds a Distinct query (GROUP BY k1, k2 with no aggregates) with the single, partial/final and partial/reduce/final chains. --- .../tests/fuzz_cases/aggregate_chain_fuzz.rs | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs index cdc9b98d95775..5f888668e7ef3 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -151,6 +151,9 @@ enum Query { Grouped, /// The same aggregates without `GROUP BY`. NoGrouping, + /// `GROUP BY k1, k2` with no aggregate expressions, as `SELECT DISTINCT` + /// plans: the accumulator-free path of every stream. + Distinct, /// `GROUP BY k1` with `max(v)` only, the shape the TopK stream supports. /// Chains using `Operator::TopK` set a limit larger than any possible /// group count, so the result must still be the complete aggregate. @@ -179,6 +182,7 @@ impl Query { match self { Query::Grouped => &["k1", "k2"], Query::NoGrouping => &[], + Query::Distinct => &["k1", "k2"], Query::TopK => &["k1"], Query::BooleanKey => &["b"], Query::BytesKey => &["s"], @@ -342,6 +346,29 @@ const SHAPES: &[Shape] = &[ 1, Query::NoGrouping, ), + shape("distinct_single", &[Aggregate(Single)], 1, Query::Distinct), + shape( + "distinct_partial_repartition_final", + &[ + Aggregate(Partial), + HashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + Query::Distinct, + ), + shape( + "distinct_partial_repartition_reduce_repartition_final", + &[ + Aggregate(Partial), + HashRepartition, + Aggregate(PartialReduce), + HashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + Query::Distinct, + ), // TopK: same query without a limit is the reference for the TopK chains shape("top_k_query_single", &[Aggregate(Single)], 1, Query::TopK), shape("top_k_single", &[TopK(Single)], 1, Query::TopK), @@ -844,6 +871,9 @@ fn aggregates(schema: &SchemaRef, query: Query) -> Vec Date: Thu, 10 Sep 2026 13:09:47 +0300 Subject: [PATCH 12/21] Run every aggregate chain over every group key type and aggregate set The shapes were a hand-picked list of (chain, query) pairs. Split them into independent axes: an operator chain, the group keys (one per GroupValues implementation) and the aggregate expressions, and generate every combination that can be planned. No-key queries skip chains that hash or sort, TopK chains take one primitive or string key with a single max or no aggregates, and max alone is dropped elsewhere as a subset of the full list. --- .../tests/fuzz_cases/aggregate_chain_fuzz.rs | 551 ++++++++---------- 1 file changed, 250 insertions(+), 301 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs index 5f888668e7ef3..2eb22bcdc2825 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -17,6 +17,12 @@ //! Fuzz test that runs every valid `AggregateExec` chain over the same data //! and asserts identical results. +//! +//! The case space is the cross product of independent axes: the operator +//! [`Chain`], the group [`Keys`] (one per `GroupValues` implementation), the +//! [`Aggregates`], the source [`Order`], the group [`Cardinality`], the +//! [`Memory`] budget and whether the skip-partial probe may fire. Only the +//! combinations that cannot be planned are left out, see [`all_shapes`]. use std::num::NonZeroUsize; use std::sync::Arc; @@ -77,14 +83,15 @@ const BATCH_SIZE: usize = 64; /// very high cardinality must still exceed it. const LIMITED_POOL_BYTES: usize = 4 * 1024 * 1024; -/// How the source data is ordered relative to the group keys `(k1, k2)`. +/// How the source data is ordered relative to the group keys. #[derive(Clone, Copy, Debug, PartialEq)] enum Order { /// Not ordered. Aggregates see `InputOrderMode::Linear`. Unordered, - /// Sorted by `k1` only. Aggregates see `InputOrderMode::PartiallySorted([0])`. + /// Sorted by the first key only. Aggregates see + /// `InputOrderMode::PartiallySorted([0])`. SortedByFirstKey, - /// Sorted by `k1, k2`. Aggregates see `InputOrderMode::Sorted`. + /// Sorted by all keys. Aggregates see `InputOrderMode::Sorted`. SortedByAllKeys, } @@ -143,63 +150,103 @@ enum Operator { SortPreservingMerge, } -/// The logical query a chain computes. -#[derive(Clone, Copy, Debug, PartialEq)] -enum Query { - /// `GROUP BY k1, k2` with count, count distinct, sum, avg, min, max. - /// Two primitive keys, handled by `GroupValuesColumn`. - Grouped, - /// The same aggregates without `GROUP BY`. - NoGrouping, - /// `GROUP BY k1, k2` with no aggregate expressions, as `SELECT DISTINCT` - /// plans: the accumulator-free path of every stream. - Distinct, - /// `GROUP BY k1` with `max(v)` only, the shape the TopK stream supports. - /// Chains using `Operator::TopK` set a limit larger than any possible - /// group count, so the result must still be the complete aggregate. - TopK, - /// `GROUP BY b` (Boolean), handled by `GroupValuesBoolean`. - BooleanKey, - /// `GROUP BY s` (Utf8), handled by `GroupValuesBytes`. - BytesKey, - /// `GROUP BY sv` (Utf8View), handled by `GroupValuesBytesView`. - BytesViewKey, - /// `GROUP BY p` (Int64 with as many distinct values as groups), handled - /// by `GroupValuesPrimitive`. - PrimitiveKey, - /// `GROUP BY b, s, sv, p`, handled by `GroupValuesColumn` with mixed - /// column types. - MixedKeys, - /// `GROUP BY st` (Struct of a List and an Int64), which no - /// specialized implementation supports, so it falls back to the row format +/// The `GROUP BY` keys. Every key type has its own `GroupValues` +/// implementation, so each is a value of this axis. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Keys { + /// No `GROUP BY`. + None, + /// `k1, k2` (two Int64), handled by `GroupValuesColumn`. + TwoInts, + /// `b` (Boolean), handled by `GroupValuesBoolean`. + Boolean, + /// `s` (Utf8), handled by `GroupValuesBytes`. + Bytes, + /// `sv` (Utf8View), handled by `GroupValuesBytesView`. + BytesView, + /// `p` (Int64 with as many distinct values as groups), handled by + /// `GroupValuesPrimitive`. + Primitive, + /// `b, s, sv, p`, handled by `GroupValuesColumn` with mixed column types. + Mixed, + /// `st` (Struct of a List and an Int64), which no specialized + /// implementation supports, so it falls back to the row format /// `GroupValuesRows`. - StructKey, + Struct, } -impl Query { - /// Group key columns, in `GROUP BY` order. - fn keys(self) -> &'static [&'static str] { +impl Keys { + const ALL: [Self; 8] = [ + Self::None, + Self::TwoInts, + Self::Boolean, + Self::Bytes, + Self::BytesView, + Self::Primitive, + Self::Mixed, + Self::Struct, + ]; + + /// Key columns, in `GROUP BY` order. + fn columns(self) -> &'static [&'static str] { match self { - Query::Grouped => &["k1", "k2"], - Query::NoGrouping => &[], - Query::Distinct => &["k1", "k2"], - Query::TopK => &["k1"], - Query::BooleanKey => &["b"], - Query::BytesKey => &["s"], - Query::BytesViewKey => &["sv"], - Query::PrimitiveKey => &["p"], - Query::MixedKeys => &["b", "s", "sv", "p"], - Query::StructKey => &["st"], + Keys::None => &[], + Keys::TwoInts => &["k1", "k2"], + Keys::Boolean => &["b"], + Keys::Bytes => &["s"], + Keys::BytesView => &["sv"], + Keys::Primitive => &["p"], + Keys::Mixed => &["b", "s", "sv", "p"], + Keys::Struct => &["st"], } } /// Whether the source can be sorted by the keys. Struct columns cannot be - /// sorted by the arrow sort kernels, so that query only runs unordered. + /// sorted by the arrow sort kernels, so those keys only run unordered. fn sortable(self) -> bool { - self != Query::StructKey + self != Keys::Struct + } + + /// Whether the number of groups is `Cardinality::groups()`. Every key + /// column has one distinct value per group except the Boolean one. + fn tracks_cardinality(self) -> bool { + !matches!(self, Keys::None | Keys::Boolean) + } + + /// Whether `GroupedTopKAggregateStream` supports these keys: exactly one + /// primitive or string column. + fn top_k_supported(self) -> bool { + matches!(self, Keys::Bytes | Keys::BytesView | Keys::Primitive) } } +/// The aggregate expressions. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Aggregates { + /// count, count distinct, sum, avg, min, max: non-trivial partial state so + /// the Partial, PartialReduce and Final stages are actually exercised. + /// `avg` (two-field state) and `count distinct` (set state) matter most. + All, + /// No aggregate expressions, as `SELECT DISTINCT` plans: the + /// accumulator-free path of every stream. + None, + /// `max(v)` only, the one aggregate the TopK stream supports. Chains using + /// `Operator::TopK` set a limit larger than any possible group count, so + /// the result must still be the complete aggregate. + Max, +} + +impl Aggregates { + const ALL: [Self; 3] = [Self::All, Self::None, Self::Max]; +} + +/// The logical query a chain computes. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct Query { + keys: Keys, + aggregates: Aggregates, +} + /// Larger than any possible number of groups, so TopK keeps every group. const TOP_K_LIMIT: usize = 2 * ROWS; @@ -215,46 +262,42 @@ struct CaseParams { skip_partial_enabled: bool, } -/// A plan shape. +/// An operator chain, independent of the query it computes. #[derive(Debug)] -struct Shape { +struct Chain { name: &'static str, operators: &'static [Operator], /// Source partition count. source_partitions: usize, - query: Query, } -const fn shape( +const fn chain( name: &'static str, operators: &'static [Operator], source_partitions: usize, - query: Query, -) -> Shape { - Shape { +) -> Chain { + Chain { name, operators, source_partitions, - query, } } -// Test cases -const SHAPES: &[Shape] = &[ - shape("single", &[Aggregate(Single)], 1, Query::Grouped), - shape( +/// Every chain the planner can produce. Each runs over every query it can be +/// planned for, see [`all_shapes`]. +const CHAINS: &[Chain] = &[ + chain("single", &[Aggregate(Single)], 1), + chain( "single_partitioned", &[HashRepartition, Aggregate(SinglePartitioned)], PARTITIONS, - Query::Grouped, ), - shape( + chain( "single_partitioned_order_preserving", &[OrderPreservingHashRepartition, Aggregate(SinglePartitioned)], PARTITIONS, - Query::Grouped, ), - shape( + chain( "partial_repartition_final", &[ Aggregate(Partial), @@ -262,15 +305,13 @@ const SHAPES: &[Shape] = &[ Aggregate(FinalPartitioned), ], PARTITIONS, - Query::Grouped, ), - shape( + chain( "partial_coalesce_final", &[Aggregate(Partial), CoalescePartitions, Aggregate(Final)], PARTITIONS, - Query::Grouped, ), - shape( + chain( "partial_order_preserving_repartition_final", &[ Aggregate(Partial), @@ -278,21 +319,18 @@ const SHAPES: &[Shape] = &[ Aggregate(FinalPartitioned), ], PARTITIONS, - Query::Grouped, ), - shape( + chain( "partial_sort_preserving_merge_final", &[Aggregate(Partial), SortPreservingMerge, Aggregate(Final)], PARTITIONS, - Query::Grouped, ), - shape( + chain( "partial_final_single_partition", &[Aggregate(Partial), Aggregate(Final)], 1, - Query::Grouped, ), - shape( + chain( "partial_repartition_reduce_repartition_final", &[ Aggregate(Partial), @@ -302,9 +340,8 @@ const SHAPES: &[Shape] = &[ Aggregate(FinalPartitioned), ], PARTITIONS, - Query::Grouped, ), - shape( + chain( "partial_repartition_reduce_coalesce_final", &[ Aggregate(Partial), @@ -314,232 +351,115 @@ const SHAPES: &[Shape] = &[ Aggregate(Final), ], PARTITIONS, - Query::Grouped, ), - shape( - "partial_local_reduce_repartition_final", + chain( + "partial_coalesce_reduce_coalesce_final", &[ Aggregate(Partial), + CoalescePartitions, Aggregate(PartialReduce), - HashRepartition, - Aggregate(FinalPartitioned), + CoalescePartitions, + Aggregate(Final), ], PARTITIONS, - Query::Grouped, ), - // ordered PartialReduce has no dedicated stream, lands on the fallback - shape( - "partial_reduce_final_order_preserving", + chain( + "partial_local_reduce_repartition_final", &[ Aggregate(Partial), - OrderPreservingHashRepartition, Aggregate(PartialReduce), - OrderPreservingHashRepartition, - Aggregate(FinalPartitioned), - ], - PARTITIONS, - Query::Grouped, - ), - shape( - "no_grouping_single", - &[Aggregate(Single)], - 1, - Query::NoGrouping, - ), - shape("distinct_single", &[Aggregate(Single)], 1, Query::Distinct), - shape( - "distinct_partial_repartition_final", - &[ - Aggregate(Partial), HashRepartition, Aggregate(FinalPartitioned), ], PARTITIONS, - Query::Distinct, ), - shape( - "distinct_partial_repartition_reduce_repartition_final", + // ordered PartialReduce has no dedicated stream, lands on the fallback + chain( + "partial_reduce_final_order_preserving", &[ Aggregate(Partial), - HashRepartition, + OrderPreservingHashRepartition, Aggregate(PartialReduce), - HashRepartition, + OrderPreservingHashRepartition, Aggregate(FinalPartitioned), ], PARTITIONS, - Query::Distinct, ), - // TopK: same query without a limit is the reference for the TopK chains - shape("top_k_query_single", &[Aggregate(Single)], 1, Query::TopK), - shape("top_k_single", &[TopK(Single)], 1, Query::TopK), + chain("top_k_single", &[TopK(Single)], 1), // planner shape: the limit lands on the aggregate under the sort - shape( + chain( "top_k_partial_repartition_final", &[Aggregate(Partial), HashRepartition, TopK(FinalPartitioned)], PARTITIONS, - Query::TopK, ), - shape( + chain( "top_k_partial_coalesce_final", &[Aggregate(Partial), CoalescePartitions, TopK(Final)], PARTITIONS, - Query::TopK, ), - shape( + chain( "top_k_both_stages", &[TopK(Partial), HashRepartition, TopK(FinalPartitioned)], PARTITIONS, - Query::TopK, - ), - // group key types: single stage and the default two-stage planner shape - shape( - "boolean_key_single", - &[Aggregate(Single)], - 1, - Query::BooleanKey, - ), - shape( - "boolean_key_partial_repartition_final", - &[ - Aggregate(Partial), - HashRepartition, - Aggregate(FinalPartitioned), - ], - PARTITIONS, - Query::BooleanKey, - ), - shape("bytes_key_single", &[Aggregate(Single)], 1, Query::BytesKey), - shape( - "bytes_key_partial_repartition_final", - &[ - Aggregate(Partial), - HashRepartition, - Aggregate(FinalPartitioned), - ], - PARTITIONS, - Query::BytesKey, - ), - shape( - "bytes_view_key_single", - &[Aggregate(Single)], - 1, - Query::BytesViewKey, - ), - shape( - "bytes_view_key_partial_repartition_final", - &[ - Aggregate(Partial), - HashRepartition, - Aggregate(FinalPartitioned), - ], - PARTITIONS, - Query::BytesViewKey, - ), - shape( - "primitive_key_single", - &[Aggregate(Single)], - 1, - Query::PrimitiveKey, - ), - shape( - "primitive_key_partial_repartition_final", - &[ - Aggregate(Partial), - HashRepartition, - Aggregate(FinalPartitioned), - ], - PARTITIONS, - Query::PrimitiveKey, - ), - shape( - "mixed_keys_single", - &[Aggregate(Single)], - 1, - Query::MixedKeys, - ), - shape( - "mixed_keys_partial_repartition_final", - &[ - Aggregate(Partial), - HashRepartition, - Aggregate(FinalPartitioned), - ], - PARTITIONS, - Query::MixedKeys, - ), - // ordered multi-column group values - shape( - "mixed_keys_partial_order_preserving_repartition_final", - &[ - Aggregate(Partial), - OrderPreservingHashRepartition, - Aggregate(FinalPartitioned), - ], - PARTITIONS, - Query::MixedKeys, - ), - shape( - "struct_key_single", - &[Aggregate(Single)], - 1, - Query::StructKey, - ), - shape( - "struct_key_partial_repartition_final", - &[ - Aggregate(Partial), - HashRepartition, - Aggregate(FinalPartitioned), - ], - PARTITIONS, - Query::StructKey, - ), - shape( - "no_grouping_partial_coalesce_final", - &[Aggregate(Partial), CoalescePartitions, Aggregate(Final)], - PARTITIONS, - Query::NoGrouping, - ), - shape( - "no_grouping_partial_reduce_final", - &[ - Aggregate(Partial), - CoalescePartitions, - Aggregate(PartialReduce), - CoalescePartitions, - Aggregate(Final), - ], - PARTITIONS, - Query::NoGrouping, - ), - shape( - "no_grouping_partial_final_single_partition", - &[Aggregate(Partial), Aggregate(Final)], - 1, - Query::NoGrouping, ), ]; -fn shape_by_name(name: &str) -> &'static Shape { - SHAPES.iter().find(|shape| shape.name == name).unwrap() +fn chain_by_name(name: &str) -> &'static Chain { + CHAINS.iter().find(|chain| chain.name == name).unwrap() +} + +impl Chain { + /// Whether the chain hashes or sorts on the group keys, so it cannot run + /// without any. + fn needs_keys(&self) -> bool { + self.operators.iter().any(|operator| { + matches!( + operator, + HashRepartition + | OrderPreservingHashRepartition + | SortPreservingMerge + | TopK(_) + ) + }) + } + + fn is_top_k(&self) -> bool { + self.operators + .iter() + .any(|operator| matches!(operator, TopK(_))) + } +} + +/// A plan shape: a chain computing a query. +#[derive(Clone, Copy, Debug)] +struct Shape { + chain: &'static Chain, + query: Query, } impl Shape { + fn name(&self) -> String { + format!( + "{} {:?} {:?}", + self.chain.name, self.query.keys, self.query.aggregates + ) + } + /// Source orders that make sense for this shape. Order-preserving shuffles /// need an ordering to preserve; no-grouping chains ignore ordering. fn orders(&self) -> Vec { - let needs_ordered_input = self.operators.iter().any(|operator| { + let needs_ordered_input = self.chain.operators.iter().any(|operator| { matches!( operator, OrderPreservingHashRepartition | SortPreservingMerge ) }); - let keys = self.query.keys(); + let keys = self.query.keys.columns(); let mut orders = vec![]; if !needs_ordered_input { orders.push(Order::Unordered); } - if self.query.sortable() && !keys.is_empty() { + if self.query.keys.sortable() && !keys.is_empty() { // With a single key, sorting by the first key is already sorting // by all keys. if keys.len() > 1 { @@ -549,23 +469,15 @@ impl Shape { } orders } -} - -#[derive(Clone, Debug)] -struct Case { - shape: &'static Shape, - params: CaseParams, -} -impl Shape { /// Whether some `Partial` stage of this shape runs the skip-partial probe /// for the given source order: grouped, not TopK, and Linear input. fn has_skip_partial_candidate(&self, order: Order) -> bool { - if self.query == Query::NoGrouping { + if self.query.keys == Keys::None { return false; } let mut current = order; - for operator in self.operators { + for operator in self.chain.operators { match operator { HashRepartition | CoalescePartitions => current = Order::Unordered, Aggregate(Partial) if current == Order::Unordered => return true, @@ -576,6 +488,44 @@ impl Shape { } } +/// Every chain over every query it can be planned for. +/// +/// - Without keys a chain can neither hash nor sort, and `max` alone is a +/// subset of the full aggregate list, so only that list runs. +/// - The TopK stream needs one primitive or string key with either a single +/// `max` or no aggregates at all (the `DISTINCT ... LIMIT` form). +/// - Every other chain runs the full aggregate list and the accumulator-free +/// form; `max` alone adds nothing there. +fn all_shapes() -> Vec { + let mut shapes = vec![]; + for chain in CHAINS { + for keys in Keys::ALL { + for aggregates in Aggregates::ALL { + let valid = if keys == Keys::None { + !chain.needs_keys() && aggregates == Aggregates::All + } else if chain.is_top_k() { + keys.top_k_supported() && aggregates != Aggregates::All + } else { + aggregates != Aggregates::Max + }; + if valid { + shapes.push(Shape { + chain, + query: Query { keys, aggregates }, + }); + } + } + } + } + shapes +} + +#[derive(Clone, Debug)] +struct Case { + shape: Shape, + params: CaseParams, +} + fn all_cases() -> Vec { // `AGGREGATE_CHAIN_SHAPES=a,b` restricts the run to shapes whose name // contains one of the given substrings, to reproduce or bisect quickly. @@ -583,11 +533,11 @@ fn all_cases() -> Vec { .map(|value| value.split(',').map(str::to_string).collect()) .unwrap_or_default(); let mut cases = vec![]; - for shape in SHAPES.iter().filter(|shape| { + for shape in all_shapes().into_iter().filter(|shape| { shape_filter.is_empty() || shape_filter .iter() - .any(|needle| shape.name.contains(needle)) + .any(|needle| shape.name().contains(needle)) }) { for order in shape.orders() { let skip_partial_variants: &[bool] = @@ -735,12 +685,12 @@ fn generate_rows(cardinality: Cardinality, seed: u64) -> RecordBatch { /// Every partition individually satisfies the ordering. fn arrange( rows: &RecordBatch, - query: Query, + keys: Keys, order: Order, partitions: usize, ) -> Vec> { let schema = rows.schema(); - let per_partition: Vec = match source_ordering(&schema, query, order) { + let per_partition: Vec = match source_ordering(&schema, keys, order) { None => { let mut permutation: Vec = (0..rows.num_rows() as u32).collect(); permutation.shuffle(&mut StdRng::seed_from_u64(0)); @@ -821,8 +771,8 @@ fn sort_expr(schema: &Schema, column: &str) -> PhysicalSortExpr { } /// The ordering the source declares for `order`. -fn source_ordering(schema: &Schema, query: Query, order: Order) -> Option { - let keys = query.keys(); +fn source_ordering(schema: &Schema, keys: Keys, order: Order) -> Option { + let keys = keys.columns(); let sort_columns: &[&str] = match order { Order::Unordered => return None, Order::SortedByFirstKey => &keys[..1], @@ -833,13 +783,13 @@ fn source_ordering(schema: &Schema, query: Query, order: Order) -> Option], - query: Query, + keys: Keys, order: Order, ) -> Arc { let schema = schema(); let mut memory_source = MemorySourceConfig::try_new(partitions, Arc::clone(&schema), None).unwrap(); - if let Some(ordering) = source_ordering(&schema, query, order) { + if let Some(ordering) = source_ordering(&schema, keys, order) { memory_source = memory_source .try_with_sort_information(vec![ordering]) .unwrap(); @@ -847,19 +797,15 @@ fn source( DataSourceExec::from_data_source(memory_source) } -fn group_by(schema: &Schema, query: Query) -> PhysicalGroupBy { +fn group_by(schema: &Schema, keys: Keys) -> PhysicalGroupBy { PhysicalGroupBy::new_single( - query - .keys() + keys.columns() .iter() .map(|key| (col(key, schema).unwrap(), key.to_string())) .collect(), ) } -/// Aggregates with non-trivial partial state so the Partial, PartialReduce and -/// Final stages are actually exercised. `avg` (two-field state) and -/// `count distinct` (set state) matter most. fn aggregates(schema: &SchemaRef, query: Query) -> Vec> { let value_column = || vec![col("v", schema).unwrap()]; let build = |builder: AggregateExprBuilder, alias: &str| { @@ -871,10 +817,10 @@ fn aggregates(schema: &SchemaRef, query: Query) -> Vec Vec) -> Arc { let input_schema = schema(); let mut plan = input; - let mut group_by = group_by(&input_schema, shape.query); + let mut group_by = group_by(&input_schema, shape.query.keys); let mut aggregates = aggregates(&input_schema, shape.query); let mut hash_keys: Vec> = group_by.input_exprs(); - for operator in shape.operators { + for operator in shape.chain.operators { plan = match operator { Aggregate(mode) | TopK(mode) => { let limit_options = matches!(operator, TopK(_)) @@ -1043,7 +989,7 @@ fn as_aggregate(node: &Arc) -> &AggregateExec { fn expected_orders(shape: &Shape, source_order: Order) -> Vec { let mut current = source_order; let mut expected = vec![]; - for operator in shape.operators { + for operator in shape.chain.operators { match operator { HashRepartition | CoalescePartitions => current = Order::Unordered, // `AggregateExec::try_new` forces `InputOrderMode::Linear` for @@ -1064,7 +1010,7 @@ fn expected_orders(shape: &Shape, source_order: Order) -> Vec { fn order_matches(query: Query, expected: Order, actual: &InputOrderMode) -> bool { // With a single group key, sorting by the first key already covers every // group key. - let single_key = query.keys().len() == 1; + let single_key = query.keys.columns().len() == 1; match (expected, actual) { (Order::Unordered, InputOrderMode::Linear) => true, (Order::SortedByFirstKey, InputOrderMode::PartiallySorted(indices)) => { @@ -1100,11 +1046,11 @@ fn runs_skip_partial_probe(aggregate: &AggregateExec) -> bool { } fn check_plan_shape(case: &Case, plan: &Arc) { - if case.shape.query == Query::NoGrouping { + if case.shape.query.keys == Keys::None { return; } let nodes = aggregate_nodes(plan); - let expected = expected_orders(case.shape, case.params.order); + let expected = expected_orders(&case.shape, case.params.order); assert_eq!(nodes.len(), expected.len(), "{case:?}"); for (node, expected_order) in nodes.iter().zip(expected) { let aggregate = as_aggregate(node); @@ -1152,11 +1098,11 @@ fn check_metrics(case: &Case, plan: &Arc) -> Vec { } } - // Only the two-key query has as many groups as `cardinality` says; - // the TopK query groups by `k1` alone and stays far below the ratio. + // Boolean keys have two groups whatever `cardinality` says, far + // below the ratio. if case.params.memory == Memory::Unlimited && case.params.cardinality == Cardinality::VeryHigh - && case.shape.query == Query::Grouped + && case.shape.query.keys.tracks_cardinality() && case.params.skip_partial_enabled && runs_skip_partial_probe(aggregate) { @@ -1197,13 +1143,13 @@ async fn run_case(case: Case, rows: Arc) -> Outcome { async fn run_case_inner(case: &Case, rows: Arc) -> Outcome { let partitions = arrange( &rows, - case.shape.query, + case.shape.query.keys, case.params.order, - case.shape.source_partitions, + case.shape.chain.source_partitions, ); let plan = build_plan( - case.shape, - source(&partitions, case.shape.query, case.params.order), + &case.shape, + source(&partitions, case.shape.query.keys, case.params.order), ); check_plan_shape(case, &plan); @@ -1233,19 +1179,17 @@ async fn run_case_inner(case: &Case, rows: Arc) -> Outcome { } } -/// Reference result: the single-stage shape of `query` without a limit, one -/// partition, unordered input, unlimited memory. +/// Reference result: `query` computed by the `single` chain, one partition, +/// unordered input, unlimited memory. async fn reference( query: Query, rows: Arc, cardinality: Cardinality, ) -> String { - let shape = SHAPES - .iter() - .find(|shape| { - shape.query == query && matches!(shape.operators, [Aggregate(Single)]) - }) - .unwrap(); + let shape = Shape { + chain: chain_by_name("single"), + query, + }; let outcome = run_case( Case { shape, @@ -1276,7 +1220,7 @@ async fn aggregate_chain_fuzz() { for cardinality in Cardinality::ALL { let rows = Arc::new(generate_rows(cardinality, seed)); let mut expected_by_query: Vec<(Query, String)> = Vec::new(); - for shape in SHAPES { + for shape in all_shapes() { if expected_by_query .iter() .any(|(query, _)| *query == shape.query) @@ -1379,7 +1323,7 @@ fn print_cases(cardinality: Cardinality, outcome: &str, cases: &[(Case, Vec Result<()> { - let shape = shape_by_name(shape_name); + let shape = Shape { + chain: chain_by_name(chain_name), + query: Query { keys, aggregates }, + }; let rows = Arc::new(generate_rows(cardinality, seed)); let expected = reference(shape.query, Arc::clone(&rows), cardinality).await; let actual = run_case( From 06bc82f807e61b4254003c442f7083a480d4fe21 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:41:25 +0300 Subject: [PATCH 13/21] Share arranged inputs across cases and halve the fuzz row count Arranging the source took about a third of every case and depends only on the keys, order and partition count, so arrange once per cardinality and share. 32k rows keep the spill coverage (about a quarter of the cases at very high cardinality still spill) at less than half the time. --- .../tests/fuzz_cases/aggregate_chain_fuzz.rs | 157 ++++++++++-------- 1 file changed, 89 insertions(+), 68 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs index 2eb22bcdc2825..d78653ad4848e 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -24,6 +24,7 @@ //! [`Memory`] budget and whether the skip-partial probe may fire. Only the //! combinations that cannot be planned are left out, see [`all_shapes`]. +use std::collections::HashMap; use std::num::NonZeroUsize; use std::sync::Arc; use std::time::Duration; @@ -72,7 +73,10 @@ use Operator::*; // Case space // --------------------------------------------------------------------------- -const ROWS: usize = 64 * 1024; +/// About 2 MB per cardinality, half of `LIMITED_POOL_BYTES`. The +/// order-preserving repartition holds a sorted low-cardinality input almost +/// whole in its merges, which cannot spill, so the input has to fit. +const ROWS: usize = 32 * 1024; const PARTITIONS: usize = 4; const BATCH_SIZE: usize = 64; /// The fair pool caps every spillable consumer at `pool / consumers`, and a @@ -84,7 +88,7 @@ const BATCH_SIZE: usize = 64; const LIMITED_POOL_BYTES: usize = 4 * 1024 * 1024; /// How the source data is ordered relative to the group keys. -#[derive(Clone, Copy, Debug, PartialEq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] enum Order { /// Not ordered. Aggregates see `InputOrderMode::Linear`. Unordered, @@ -152,7 +156,7 @@ enum Operator { /// The `GROUP BY` keys. Every key type has its own `GroupValues` /// implementation, so each is a value of this axis. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] enum Keys { /// No `GROUP BY`. None, @@ -1128,28 +1132,51 @@ struct Outcome { spilled: Vec, } +/// Arranged source partitions by `(keys, order, partition count)`, the only +/// case dimensions the arrangement depends on. Arranging costs about a third +/// of a case, so it is shared across chains, memory budgets and skip-partial +/// settings. +type Inputs = HashMap<(Keys, Order, usize), Arc>>>; + +fn input_key(case: &Case) -> (Keys, Order, usize) { + ( + case.shape.query.keys, + case.params.order, + case.shape.chain.source_partitions, + ) +} + +fn arrange_all<'a>(rows: &RecordBatch, cases: impl Iterator) -> Inputs { + let mut inputs = Inputs::new(); + for case in cases { + inputs.entry(input_key(case)).or_insert_with(|| { + Arc::new(arrange( + rows, + case.shape.query.keys, + case.params.order, + case.shape.chain.source_partitions, + )) + }); + } + inputs +} + /// Runs one case, checks plan shape and metrics, and returns its outcome. /// /// Running out of memory is never accepted: every stream either spills, emits /// early, or is bounded, so an error there is a bug in a stream's memory /// handling or in how the stages share the pool. -async fn run_case(case: Case, rows: Arc) -> Outcome { +async fn run_case(case: Case, inputs: Arc) -> Outcome { log::debug!("start {case:?}"); - let outcome = run_case_inner(&case, rows).await; + let outcome = run_case_inner(&case, &inputs[&input_key(&case)]).await; log::debug!("done {case:?}"); outcome } -async fn run_case_inner(case: &Case, rows: Arc) -> Outcome { - let partitions = arrange( - &rows, - case.shape.query.keys, - case.params.order, - case.shape.chain.source_partitions, - ); +async fn run_case_inner(case: &Case, partitions: &[Vec]) -> Outcome { let plan = build_plan( &case.shape, - source(&partitions, case.shape.query.keys, case.params.order), + source(partitions, case.shape.query.keys, case.params.order), ); check_plan_shape(case, &plan); @@ -1179,31 +1206,21 @@ async fn run_case_inner(case: &Case, rows: Arc) -> Outcome { } } -/// Reference result: `query` computed by the `single` chain, one partition, -/// unordered input, unlimited memory. -async fn reference( - query: Query, - rows: Arc, - cardinality: Cardinality, -) -> String { - let shape = Shape { - chain: chain_by_name("single"), - query, - }; - let outcome = run_case( - Case { - shape, - params: CaseParams { - order: Order::Unordered, - cardinality, - memory: Memory::Unlimited, - skip_partial_enabled: true, - }, +/// The case whose output is the reference for `query`: the `single` chain, +/// one partition, unordered input, unlimited memory. +fn reference_case(query: Query, cardinality: Cardinality) -> Case { + Case { + shape: Shape { + chain: chain_by_name("single"), + query, }, - rows, - ) - .await; - outcome.output + params: CaseParams { + order: Order::Unordered, + cardinality, + memory: Memory::Unlimited, + skip_partial_enabled: true, + }, + } } #[tokio::test(flavor = "multi_thread")] @@ -1218,26 +1235,31 @@ async fn aggregate_chain_fuzz() { let max_concurrent_cases = get_available_parallelism(); for cardinality in Cardinality::ALL { - let rows = Arc::new(generate_rows(cardinality, seed)); - let mut expected_by_query: Vec<(Query, String)> = Vec::new(); - for shape in all_shapes() { - if expected_by_query - .iter() - .any(|(query, _)| *query == shape.query) - { - continue; + let rows = generate_rows(cardinality, seed); + let cases: Vec = all_cases() + .into_iter() + .filter(|case| case.params.cardinality == cardinality) + .collect(); + let mut reference_cases: Vec = vec![]; + for case in &cases { + let query = case.shape.query; + if !reference_cases.iter().any(|case| case.shape.query == query) { + reference_cases.push(reference_case(query, cardinality)); } - let expected = reference(shape.query, Arc::clone(&rows), cardinality).await; - expected_by_query.push((shape.query, expected)); + } + let inputs = Arc::new(arrange_all(&rows, cases.iter().chain(&reference_cases))); + + let mut expected_by_query: Vec<(Query, String)> = Vec::new(); + for case in reference_cases { + let query = case.shape.query; + let outcome = run_case(case, Arc::clone(&inputs)).await; + expected_by_query.push((query, outcome.output)); } let mut join_set = JoinSet::new(); let (mut spilled, mut finished) = (vec![], vec![]); - for case in all_cases() - .into_iter() - .filter(|case| case.params.cardinality == cardinality) - { - let rows = Arc::clone(&rows); + for case in cases { + let inputs = Arc::clone(&inputs); let expected = expected_by_query .iter() .find(|(query, _)| *query == case.shape.query) @@ -1253,7 +1275,7 @@ async fn aggregate_chain_fuzz() { .await; } join_set.spawn(async move { - let outcome = run_case(case.clone(), rows).await; + let outcome = run_case(case.clone(), inputs).await; assert_eq!(outcome.output, expected, "{case:?} (seed {seed})"); (case, outcome.spilled) }); @@ -1352,21 +1374,20 @@ async fn run_single_case( chain: chain_by_name(chain_name), query: Query { keys, aggregates }, }; - let rows = Arc::new(generate_rows(cardinality, seed)); - let expected = reference(shape.query, Arc::clone(&rows), cardinality).await; - let actual = run_case( - Case { - shape, - params: CaseParams { - order, - cardinality, - memory, - skip_partial_enabled: true, - }, + let rows = generate_rows(cardinality, seed); + let reference = reference_case(shape.query, cardinality); + let case = Case { + shape, + params: CaseParams { + order, + cardinality, + memory, + skip_partial_enabled: true, }, - rows, - ) - .await; - assert_eq!(actual.output, expected); + }; + let inputs = Arc::new(arrange_all(&rows, [&case, &reference].into_iter())); + let expected = run_case(reference, Arc::clone(&inputs)).await; + let actual = run_case(case, inputs).await; + assert_eq!(actual.output, expected.output); Ok(()) } From fe6b6538a33a81e12bc2e9945211c1e8ab5aa46a Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:55:24 +0300 Subject: [PATCH 14/21] Drop the High cardinality level from the aggregate chain fuzz It spills at the same rate as VeryHigh and exercises nothing VeryHigh does not, at a third of the run time. --- .../core/tests/fuzz_cases/aggregate_chain_fuzz.rs | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs index d78653ad4848e..86aa3bcc360d1 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -102,26 +102,18 @@ enum Order { #[derive(Clone, Copy, Debug, PartialEq)] enum Cardinality { VeryHigh, - High, Medium, Low, VeryLow, } impl Cardinality { - const ALL: [Self; 5] = [ - Self::VeryHigh, - Self::High, - Self::Medium, - Self::Low, - Self::VeryLow, - ]; + const ALL: [Self; 4] = [Self::VeryHigh, Self::Medium, Self::Low, Self::VeryLow]; /// Number of distinct `(k1, k2)` groups. fn groups(self) -> usize { match self { Self::VeryHigh => ROWS, - Self::High => ROWS / 2, Self::Medium => ROWS / 32, Self::Low => 16, Self::VeryLow => 2, From 4e5766d4c05572a89fcbf71e382db8e0bd967f79 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 10 Sep 2026 14:12:23 +0300 Subject: [PATCH 15/21] Copy fuzz batches from the unsliced batch so memory accounting is real take on a sliced list sizes the new values buffer from the child length divided by the slice length, so a 64-row slice of a 32k-row list got a values buffer for the whole child and get_array_memory_size, which reports capacity, charged every batch about 500 KB instead of 5 KB. Eight batches in flight filled the 4 MB pool, which is what the order-preserving repartition merge was running out of. --- .../tests/fuzz_cases/aggregate_chain_fuzz.rs | 31 ++++++++++++------- 1 file changed, 19 insertions(+), 12 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs index 86aa3bcc360d1..709971e7cd5ce 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -73,9 +73,6 @@ use Operator::*; // Case space // --------------------------------------------------------------------------- -/// About 2 MB per cardinality, half of `LIMITED_POOL_BYTES`. The -/// order-preserving repartition holds a sorted low-cardinality input almost -/// whole in its merges, which cannot spill, so the input has to fit. const ROWS: usize = 32 * 1024; const PARTITIONS: usize = 4; const BATCH_SIZE: usize = 64; @@ -722,7 +719,7 @@ fn arrange( .map(|partition| { let start = (partition * per_partition).min(sorted.num_rows()); let length = per_partition.min(sorted.num_rows() - start); - sorted.slice(start, length) + copy_rows(&sorted, start, length) }) .collect() } @@ -738,20 +735,30 @@ fn arrange( (0..partition.num_rows()) .step_by(BATCH_SIZE) .map(|start| { - let slice = partition - .slice(start, BATCH_SIZE.min(partition.num_rows() - start)); - { - // `take` copies; `concat_batches` of one batch only slices - let indices = - UInt32Array::from_iter_values(0..slice.num_rows() as u32); - take_record_batch(&slice, &indices).unwrap() - } + copy_rows( + partition, + start, + BATCH_SIZE.min(partition.num_rows() - start), + ) }) .collect() }) .collect() } +/// `batch[start..start + length]` in its own buffers. `take` copies where +/// `slice` shares and `concat_batches` of one batch only slices. +/// +/// Take from the unsliced batch: `take` on a list sizes the new values buffer +/// as child length / list length * taken rows, so taking 64 rows out of a +/// 64-row slice of a 32k-row list allocates a values buffer for the whole +/// child, and `get_array_memory_size` reports capacity. That charged every +/// batch about 500 KB instead of 5 KB. +fn copy_rows(batch: &RecordBatch, start: usize, length: usize) -> RecordBatch { + let indices = UInt32Array::from_iter_values(start as u32..(start + length) as u32); + take_record_batch(batch, &indices).unwrap() +} + // --------------------------------------------------------------------------- // Plan construction // --------------------------------------------------------------------------- From 109dfe187c33eaf2da11f6948462ef4cd830ba2b Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 17 Sep 2026 13:13:28 +0300 Subject: [PATCH 16/21] change to fixed seed --- datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs index 709971e7cd5ce..47fa4f6b114b3 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -1224,8 +1224,7 @@ fn reference_case(query: Query, cardinality: Cardinality) -> Case { #[tokio::test(flavor = "multi_thread")] async fn aggregate_chain_fuzz() { - let seed = rand::random::(); - println!("aggregate_chain_fuzz seed = {seed}"); + const SEED: u64 = 42; let mut total_spilled = 0; let mut failures: Vec = vec![]; // Every in-flight case holds several copies of the dataset and its own @@ -1234,7 +1233,7 @@ async fn aggregate_chain_fuzz() { let max_concurrent_cases = get_available_parallelism(); for cardinality in Cardinality::ALL { - let rows = generate_rows(cardinality, seed); + let rows = generate_rows(cardinality, SEED); let cases: Vec = all_cases() .into_iter() .filter(|case| case.params.cardinality == cardinality) @@ -1275,7 +1274,7 @@ async fn aggregate_chain_fuzz() { } join_set.spawn(async move { let outcome = run_case(case.clone(), inputs).await; - assert_eq!(outcome.output, expected, "{case:?} (seed {seed})"); + assert_eq!(outcome.output, expected, "{case:?}"); (case, outcome.spilled) }); } @@ -1293,7 +1292,7 @@ async fn aggregate_chain_fuzz() { } assert!( failures.is_empty(), - "{} cases failed (seed {seed}):\n\n{}", + "{} cases failed:\n\n{}", failures.len(), failures.join("\n\n") ); From fb024b535d4f4f16ebcca67598b5a754ec197a75 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 17 Sep 2026 13:23:27 +0300 Subject: [PATCH 17/21] changed ordered partial stream to mark as can spill for early emit --- .../src/aggregates/ordered_partial_stream.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index 8c2315588ad7a..aed1005ce3c25 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -139,10 +139,10 @@ impl OrderedPartialAggregateStream { )?; let reservation = MemoryConsumer::new(format!("OrderedPartialAggregateStream[{partition}]")) - .with_can_spill(matches!( - table.group_ordering(), - GroupOrdering::Partial(_) - )) + // We interpret 'can spill' as 'can handle memory back pressure'. + // This value needs to be set to true and for every ordering except full, which fail on OOM we early emit. + // to ensure fair application of back pressure amongst the memory consumers. + .with_can_spill(!matches!(table.group_ordering(), GroupOrdering::Full(_))) .register(context.memory_pool()); Ok(Self { From 939b2512dbae2359140c29d6fbeeb25be553e1f1 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 17 Sep 2026 16:38:13 +0300 Subject: [PATCH 18/21] extract --- .../tests/fuzz_cases/aggregate_chain_fuzz.rs | 1361 +++++------------ .../aggregate_chain_fuzz/assertions.rs | 181 +++ .../aggregate_chain_fuzz/case_space.rs | 400 +++++ .../aggregate_chain_fuzz/context.rs | 65 + .../fuzz_cases/aggregate_chain_fuzz/data.rs | 215 +++ .../fuzz_cases/aggregate_chain_fuzz/plan.rs | 177 +++ datafusion/core/tests/fuzz_cases/mod.rs | 2 +- 7 files changed, 1389 insertions(+), 1012 deletions(-) create mode 100644 datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/assertions.rs create mode 100644 datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/case_space.rs create mode 100644 datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/context.rs create mode 100644 datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/data.rs create mode 100644 datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/plan.rs diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs index 47fa4f6b114b3..1d733d9f21d14 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -15,14 +15,15 @@ // specific language governing permissions and limitations // under the License. -//! Fuzz test that runs every valid `AggregateExec` chain over the same data -//! and asserts identical results. +//! Fuzz tests that run every valid `AggregateExec` chain over the same data +//! and assert identical results. //! -//! The case space is the cross product of independent axes: the operator -//! [`Chain`], the group [`Keys`] (one per `GroupValues` implementation), the -//! [`Aggregates`], the source [`Order`], the group [`Cardinality`], the -//! [`Memory`] budget and whether the skip-partial probe may fire. Only the -//! combinations that cannot be planned are left out, see [`all_shapes`]. +//! One test per operator [`Chain`], each running the chain over the cross +//! product of the other axes: the group [`Keys`] (one per `GroupValues` +//! implementation), the [`Aggregates`], the source [`Order`], the group +//! [`Cardinality`], the [`Memory`] budget and whether the skip-partial probe +//! may fire. Only the combinations that cannot be planned are left out, see +//! [`shapes`]. use std::collections::HashMap; use std::num::NonZeroUsize; @@ -39,7 +40,6 @@ use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef, SortOptions}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::prelude::SessionConfig; -use datafusion_common::Result; use datafusion_common::test_util::batches_to_sort_string; use datafusion_common::utils::get_available_parallelism; use datafusion_common_runtime::JoinSet; @@ -69,228 +69,137 @@ use rand::{Rng, SeedableRng}; use AggregateMode::*; use Operator::*; -// --------------------------------------------------------------------------- -// Case space -// --------------------------------------------------------------------------- - -const ROWS: usize = 32 * 1024; -const PARTITIONS: usize = 4; -const BATCH_SIZE: usize = 64; -/// The fair pool caps every spillable consumer at `pool / consumers`, and a -/// chain registers up to twenty consumers (aggregate streams plus one per -/// repartition channel). The cap has to clear a small table's legitimate -/// footprint, which at very low cardinality is dominated by the `count -/// distinct` sets and grows in steps of roughly 100 KB, while a final table at -/// very high cardinality must still exceed it. -const LIMITED_POOL_BYTES: usize = 4 * 1024 * 1024; - -/// How the source data is ordered relative to the group keys. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] -enum Order { - /// Not ordered. Aggregates see `InputOrderMode::Linear`. - Unordered, - /// Sorted by the first key only. Aggregates see - /// `InputOrderMode::PartiallySorted([0])`. - SortedByFirstKey, - /// Sorted by all keys. Aggregates see `InputOrderMode::Sorted`. - SortedByAllKeys, -} - -#[derive(Clone, Copy, Debug, PartialEq)] -enum Cardinality { - VeryHigh, - Medium, - Low, - VeryLow, -} - -impl Cardinality { - const ALL: [Self; 4] = [Self::VeryHigh, Self::Medium, Self::Low, Self::VeryLow]; - - /// Number of distinct `(k1, k2)` groups. - fn groups(self) -> usize { - match self { - Self::VeryHigh => ROWS, - Self::Medium => ROWS / 32, - Self::Low => 16, - Self::VeryLow => 2, - } - } -} - -#[derive(Clone, Copy, Debug, PartialEq)] -enum Memory { - /// Unlimited pool. Nothing spills or emits early. - Unlimited, - /// Pool sized so final and single hash tables cannot fit. - Limited, -} - -/// One operator in a chain, listed bottom to top. -#[derive(Clone, Copy, Debug)] -enum Operator { - Aggregate(AggregateMode), - /// `AggregateExec` with `limit_options` set, which selects - /// `GroupedTopKAggregateStream` regardless of mode. - TopK(AggregateMode), - /// `RepartitionExec` hashed on the group keys. Destroys ordering. - HashRepartition, - /// `RepartitionExec` hashed on the group keys with `preserve_order`. - OrderPreservingHashRepartition, - /// `CoalescePartitionsExec`. Destroys ordering. - CoalescePartitions, - /// `SortPreservingMergeExec` on the current ordering. - SortPreservingMerge, -} - -/// The `GROUP BY` keys. Every key type has its own `GroupValues` -/// implementation, so each is a value of this axis. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] -enum Keys { - /// No `GROUP BY`. - None, - /// `k1, k2` (two Int64), handled by `GroupValuesColumn`. - TwoInts, - /// `b` (Boolean), handled by `GroupValuesBoolean`. - Boolean, - /// `s` (Utf8), handled by `GroupValuesBytes`. - Bytes, - /// `sv` (Utf8View), handled by `GroupValuesBytesView`. - BytesView, - /// `p` (Int64 with as many distinct values as groups), handled by - /// `GroupValuesPrimitive`. - Primitive, - /// `b, s, sv, p`, handled by `GroupValuesColumn` with mixed column types. - Mixed, - /// `st` (Struct of a List and an Int64), which no specialized - /// implementation supports, so it falls back to the row format - /// `GroupValuesRows`. - Struct, -} +mod assertions; +mod case_space; +mod context; +mod data; +mod plan; -impl Keys { - const ALL: [Self; 8] = [ - Self::None, - Self::TwoInts, - Self::Boolean, - Self::Bytes, - Self::BytesView, - Self::Primitive, - Self::Mixed, - Self::Struct, - ]; +use assertions::*; +use case_space::*; +use context::*; +use data::*; +use plan::*; - /// Key columns, in `GROUP BY` order. - fn columns(self) -> &'static [&'static str] { - match self { - Keys::None => &[], - Keys::TwoInts => &["k1", "k2"], - Keys::Boolean => &["b"], - Keys::Bytes => &["s"], - Keys::BytesView => &["sv"], - Keys::Primitive => &["p"], - Keys::Mixed => &["b", "s", "sv", "p"], - Keys::Struct => &["st"], - } - } - - /// Whether the source can be sorted by the keys. Struct columns cannot be - /// sorted by the arrow sort kernels, so those keys only run unordered. - fn sortable(self) -> bool { - self != Keys::Struct - } - - /// Whether the number of groups is `Cardinality::groups()`. Every key - /// column has one distinct value per group except the Boolean one. - fn tracks_cardinality(self) -> bool { - !matches!(self, Keys::None | Keys::Boolean) - } - - /// Whether `GroupedTopKAggregateStream` supports these keys: exactly one - /// primitive or string column. - fn top_k_supported(self) -> bool { - matches!(self, Keys::Bytes | Keys::BytesView | Keys::Primitive) - } -} - -/// The aggregate expressions. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum Aggregates { - /// count, count distinct, sum, avg, min, max: non-trivial partial state so - /// the Partial, PartialReduce and Final stages are actually exercised. - /// `avg` (two-field state) and `count distinct` (set state) matter most. - All, - /// No aggregate expressions, as `SELECT DISTINCT` plans: the - /// accumulator-free path of every stream. - None, - /// `max(v)` only, the one aggregate the TopK stream supports. Chains using - /// `Operator::TopK` set a limit larger than any possible group count, so - /// the result must still be the complete aggregate. - Max, -} - -impl Aggregates { - const ALL: [Self; 3] = [Self::All, Self::None, Self::Max]; -} - -/// The logical query a chain computes. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -struct Query { - keys: Keys, - aggregates: Aggregates, -} - -/// Larger than any possible number of groups, so TopK keeps every group. -const TOP_K_LIMIT: usize = 2 * ROWS; - -/// Everything that varies for a case apart from the shape itself. Passed as a -/// struct so a new dimension does not change every shape predicate. -#[derive(Clone, Copy, Debug, PartialEq)] -struct CaseParams { - order: Order, - cardinality: Cardinality, - memory: Memory, - /// Whether the skip-partial probe may fire. Only varied for shapes with a - /// grouped `Partial` stage on Linear input, since nothing else runs it. - skip_partial_enabled: bool, -} - -/// An operator chain, independent of the query it computes. -#[derive(Debug)] -struct Chain { - name: &'static str, - operators: &'static [Operator], - /// Source partition count. - source_partitions: usize, -} - -const fn chain( - name: &'static str, - operators: &'static [Operator], - source_partitions: usize, -) -> Chain { - Chain { - name, - operators, - source_partitions, - } +// What is tested +// ============== +// +// Every test below is one physical plan shape, an `AggregateExec` chain. +// `chain` lists its operators bottom-up, source first; the doc comment shows +// the same plan as DataFusion prints it. The test runs that plan for every +// query and input below and asserts it returns the same rows as the plain +// single-stage aggregate `SINGLE`. +// +// The table (`data.rs`), 32K rows, generated from a fixed seed: +// +// k1 Int64, k2 Int64 two-column key, k1 alone has fewer distinct +// values than (k1, k2) +// v Int64 the aggregated value, -1000..1000 +// b Boolean one column per `GroupValues` implementation, +// s Utf8 each with one distinct value per (k1, k2) +// sv Utf8View group (Boolean: two) +// p Int64 +// st Struct, num: Int64> +// +// Every key column has about 3% nulls. The number of groups is the +// cardinality axis: 32K (one row per group), 1K, 16 or 2. +// +// The queries (`case_space.rs`), one per key set times one per aggregate list: +// +// SELECT FROM t -- no GROUP BY +// SELECT k1, k2, FROM t GROUP BY k1, k2 -- GroupValuesColumn +// SELECT b, FROM t GROUP BY b -- GroupValuesBoolean +// SELECT s, FROM t GROUP BY s -- GroupValuesBytes +// SELECT sv, FROM t GROUP BY sv -- GroupValuesBytesView +// SELECT p, FROM t GROUP BY p -- GroupValuesPrimitive +// SELECT b, s, sv, p, FROM t GROUP BY b, s, sv, p +// SELECT st, FROM t GROUP BY st -- row format fallback +// +// is one of +// count(v), count(DISTINCT v), sum(v), avg(v), min(v), max(v) +// nothing, as in SELECT DISTINCT keys +// max(v) alone, the one aggregate the TopK stream supports +// +// TopK chains add LIMIT 64K, above any group count, so all groups survive. +// +// The input (`data.rs`), for each query: +// +// order unordered (shuffled, round-robin over partitions), sorted by +// the first key, or sorted by all keys, each partition sorted +// partitions 1 or 4, in batches of 64 rows +// memory unlimited, or a 4 MB pool too small for a 32K-group table +// skip-partial on or off, where a `Partial` stage on unordered input runs it +// +// Not every combination can be planned, see `shapes` and `Shape::orders`. +// +// Per case, besides the rows matching `SINGLE`, the test asserts +// (`assertions.rs`) that the plan was built as intended (modes, input order +// modes, partition counts), that nothing fails or hangs, in particular not +// with out of memory, that only spill-capable stages spill, and that the +// skip-partial probe fires exactly when it may. Per chain, at least one case +// must spill when the chain has a spill-capable stage on unordered input. + +/// The reference chain every other chain is compared against. +const SINGLE: Chain = chain("single", &[Aggregate(Single)], 1); + +/// `Single` on one partition. Also the reference every other chain is compared +/// against. +/// +/// ```text +/// AggregateExec: mode=Single +/// DataSourceExec: partitions=1 +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn single() { + assert_chain_matches_single_aggregate(SINGLE).await; } -/// Every chain the planner can produce. Each runs over every query it can be -/// planned for, see [`all_shapes`]. -const CHAINS: &[Chain] = &[ - chain("single", &[Aggregate(Single)], 1), - chain( +/// Each partition aggregates its own keys in one pass. +/// +/// ```text +/// AggregateExec: mode=SinglePartitioned +/// RepartitionExec: partitioning=Hash(keys) +/// DataSourceExec: partitions=PARTITIONS +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn single_partitioned() { + assert_chain_matches_single_aggregate(chain( "single_partitioned", &[HashRepartition, Aggregate(SinglePartitioned)], PARTITIONS, - ), - chain( + )) + .await; +} + +/// The shuffle keeps the source ordering, so the single stage still sees sorted +/// input. +/// +/// ```text +/// AggregateExec: mode=SinglePartitioned +/// RepartitionExec: partitioning=Hash(keys), preserve_order=true +/// DataSourceExec: partitions=PARTITIONS +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn single_partitioned_order_preserving() { + assert_chain_matches_single_aggregate(chain( "single_partitioned_order_preserving", &[OrderPreservingHashRepartition, Aggregate(SinglePartitioned)], PARTITIONS, - ), - chain( + )) + .await; +} + +/// The planner's default two-stage plan. +/// +/// ```text +/// AggregateExec: mode=FinalPartitioned +/// RepartitionExec: partitioning=Hash(keys) +/// AggregateExec: mode=Partial +/// DataSourceExec: partitions=PARTITIONS +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn partial_repartition_final() { + assert_chain_matches_single_aggregate(chain( "partial_repartition_final", &[ Aggregate(Partial), @@ -298,13 +207,40 @@ const CHAINS: &[Chain] = &[ Aggregate(FinalPartitioned), ], PARTITIONS, - ), - chain( + )) + .await; +} + +/// Two stages merged into one output partition. +/// +/// ```text +/// AggregateExec: mode=Final +/// CoalescePartitionsExec +/// AggregateExec: mode=Partial +/// DataSourceExec: partitions=PARTITIONS +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn partial_coalesce_final() { + assert_chain_matches_single_aggregate(chain( "partial_coalesce_final", &[Aggregate(Partial), CoalescePartitions, Aggregate(Final)], PARTITIONS, - ), - chain( + )) + .await; +} + +/// Two stages whose shuffle keeps the source ordering, so the final stage sees +/// sorted input. +/// +/// ```text +/// AggregateExec: mode=FinalPartitioned +/// RepartitionExec: partitioning=Hash(keys), preserve_order=true +/// AggregateExec: mode=Partial +/// DataSourceExec: partitions=PARTITIONS +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn partial_order_preserving_repartition_final() { + assert_chain_matches_single_aggregate(chain( "partial_order_preserving_repartition_final", &[ Aggregate(Partial), @@ -312,18 +248,59 @@ const CHAINS: &[Chain] = &[ Aggregate(FinalPartitioned), ], PARTITIONS, - ), - chain( + )) + .await; +} + +/// Two stages merged by a sort-preserving merge, so the final stage sees sorted +/// input. +/// +/// ```text +/// AggregateExec: mode=Final +/// SortPreservingMergeExec: [keys] +/// AggregateExec: mode=Partial +/// DataSourceExec: partitions=PARTITIONS +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn partial_sort_preserving_merge_final() { + assert_chain_matches_single_aggregate(chain( "partial_sort_preserving_merge_final", &[Aggregate(Partial), SortPreservingMerge, Aggregate(Final)], PARTITIONS, - ), - chain( + )) + .await; +} + +/// Two stages back to back on one partition, no shuffle between. +/// +/// ```text +/// AggregateExec: mode=Final +/// AggregateExec: mode=Partial +/// DataSourceExec: partitions=1 +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn partial_final_single_partition() { + assert_chain_matches_single_aggregate(chain( "partial_final_single_partition", &[Aggregate(Partial), Aggregate(Final)], 1, - ), - chain( + )) + .await; +} + +/// Three stages with a `PartialReduce` between two shuffles. +/// +/// ```text +/// AggregateExec: mode=FinalPartitioned +/// RepartitionExec: partitioning=Hash(keys) +/// AggregateExec: mode=PartialReduce +/// RepartitionExec: partitioning=Hash(keys) +/// AggregateExec: mode=Partial +/// DataSourceExec: partitions=PARTITIONS +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn partial_repartition_reduce_repartition_final() { + assert_chain_matches_single_aggregate(chain( "partial_repartition_reduce_repartition_final", &[ Aggregate(Partial), @@ -333,8 +310,23 @@ const CHAINS: &[Chain] = &[ Aggregate(FinalPartitioned), ], PARTITIONS, - ), - chain( + )) + .await; +} + +/// Three stages: a shuffled `PartialReduce` merged into one final partition. +/// +/// ```text +/// AggregateExec: mode=Final +/// CoalescePartitionsExec +/// AggregateExec: mode=PartialReduce +/// RepartitionExec: partitioning=Hash(keys) +/// AggregateExec: mode=Partial +/// DataSourceExec: partitions=PARTITIONS +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn partial_repartition_reduce_coalesce_final() { + assert_chain_matches_single_aggregate(chain( "partial_repartition_reduce_coalesce_final", &[ Aggregate(Partial), @@ -344,8 +336,24 @@ const CHAINS: &[Chain] = &[ Aggregate(Final), ], PARTITIONS, - ), - chain( + )) + .await; +} + +/// Three stages where `PartialReduce` and `Final` each run on one coalesced +/// partition. +/// +/// ```text +/// AggregateExec: mode=Final +/// CoalescePartitionsExec +/// AggregateExec: mode=PartialReduce +/// CoalescePartitionsExec +/// AggregateExec: mode=Partial +/// DataSourceExec: partitions=PARTITIONS +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn partial_coalesce_reduce_coalesce_final() { + assert_chain_matches_single_aggregate(chain( "partial_coalesce_reduce_coalesce_final", &[ Aggregate(Partial), @@ -355,8 +363,22 @@ const CHAINS: &[Chain] = &[ Aggregate(Final), ], PARTITIONS, - ), - chain( + )) + .await; +} + +/// `PartialReduce` directly on top of `Partial`, before the shuffle. +/// +/// ```text +/// AggregateExec: mode=FinalPartitioned +/// RepartitionExec: partitioning=Hash(keys) +/// AggregateExec: mode=PartialReduce +/// AggregateExec: mode=Partial +/// DataSourceExec: partitions=PARTITIONS +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn partial_local_reduce_repartition_final() { + assert_chain_matches_single_aggregate(chain( "partial_local_reduce_repartition_final", &[ Aggregate(Partial), @@ -365,9 +387,24 @@ const CHAINS: &[Chain] = &[ Aggregate(FinalPartitioned), ], PARTITIONS, - ), - // ordered PartialReduce has no dedicated stream, lands on the fallback - chain( + )) + .await; +} + +/// Three stages joined by order-preserving shuffles. Ordered `PartialReduce` +/// has no dedicated stream and lands on the fallback. +/// +/// ```text +/// AggregateExec: mode=FinalPartitioned +/// RepartitionExec: partitioning=Hash(keys), preserve_order=true +/// AggregateExec: mode=PartialReduce +/// RepartitionExec: partitioning=Hash(keys), preserve_order=true +/// AggregateExec: mode=Partial +/// DataSourceExec: partitions=PARTITIONS +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn partial_reduce_final_order_preserving() { + assert_chain_matches_single_aggregate(chain( "partial_reduce_final_order_preserving", &[ Aggregate(Partial), @@ -377,748 +414,76 @@ const CHAINS: &[Chain] = &[ Aggregate(FinalPartitioned), ], PARTITIONS, - ), - chain("top_k_single", &[TopK(Single)], 1), - // planner shape: the limit lands on the aggregate under the sort - chain( - "top_k_partial_repartition_final", - &[Aggregate(Partial), HashRepartition, TopK(FinalPartitioned)], - PARTITIONS, - ), - chain( - "top_k_partial_coalesce_final", - &[Aggregate(Partial), CoalescePartitions, TopK(Final)], - PARTITIONS, - ), - chain( - "top_k_both_stages", - &[TopK(Partial), HashRepartition, TopK(FinalPartitioned)], - PARTITIONS, - ), -]; - -fn chain_by_name(name: &str) -> &'static Chain { - CHAINS.iter().find(|chain| chain.name == name).unwrap() -} - -impl Chain { - /// Whether the chain hashes or sorts on the group keys, so it cannot run - /// without any. - fn needs_keys(&self) -> bool { - self.operators.iter().any(|operator| { - matches!( - operator, - HashRepartition - | OrderPreservingHashRepartition - | SortPreservingMerge - | TopK(_) - ) - }) - } - - fn is_top_k(&self) -> bool { - self.operators - .iter() - .any(|operator| matches!(operator, TopK(_))) - } -} - -/// A plan shape: a chain computing a query. -#[derive(Clone, Copy, Debug)] -struct Shape { - chain: &'static Chain, - query: Query, + )) + .await; } -impl Shape { - fn name(&self) -> String { - format!( - "{} {:?} {:?}", - self.chain.name, self.query.keys, self.query.aggregates - ) - } - - /// Source orders that make sense for this shape. Order-preserving shuffles - /// need an ordering to preserve; no-grouping chains ignore ordering. - fn orders(&self) -> Vec { - let needs_ordered_input = self.chain.operators.iter().any(|operator| { - matches!( - operator, - OrderPreservingHashRepartition | SortPreservingMerge - ) - }); - let keys = self.query.keys.columns(); - let mut orders = vec![]; - if !needs_ordered_input { - orders.push(Order::Unordered); - } - if self.query.keys.sortable() && !keys.is_empty() { - // With a single key, sorting by the first key is already sorting - // by all keys. - if keys.len() > 1 { - orders.push(Order::SortedByFirstKey); - } - orders.push(Order::SortedByAllKeys); - } - orders - } - - /// Whether some `Partial` stage of this shape runs the skip-partial probe - /// for the given source order: grouped, not TopK, and Linear input. - fn has_skip_partial_candidate(&self, order: Order) -> bool { - if self.query.keys == Keys::None { - return false; - } - let mut current = order; - for operator in self.chain.operators { - match operator { - HashRepartition | CoalescePartitions => current = Order::Unordered, - Aggregate(Partial) if current == Order::Unordered => return true, - _ => {} - } - } - false - } -} - -/// Every chain over every query it can be planned for. +/// `GroupedTopKAggregateStream` alone. The limit is above the group count, so +/// every group survives. /// -/// - Without keys a chain can neither hash nor sort, and `max` alone is a -/// subset of the full aggregate list, so only that list runs. -/// - The TopK stream needs one primitive or string key with either a single -/// `max` or no aggregates at all (the `DISTINCT ... LIMIT` form). -/// - Every other chain runs the full aggregate list and the accumulator-free -/// form; `max` alone adds nothing there. -fn all_shapes() -> Vec { - let mut shapes = vec![]; - for chain in CHAINS { - for keys in Keys::ALL { - for aggregates in Aggregates::ALL { - let valid = if keys == Keys::None { - !chain.needs_keys() && aggregates == Aggregates::All - } else if chain.is_top_k() { - keys.top_k_supported() && aggregates != Aggregates::All - } else { - aggregates != Aggregates::Max - }; - if valid { - shapes.push(Shape { - chain, - query: Query { keys, aggregates }, - }); - } - } - } - } - shapes -} - -#[derive(Clone, Debug)] -struct Case { - shape: Shape, - params: CaseParams, -} - -fn all_cases() -> Vec { - // `AGGREGATE_CHAIN_SHAPES=a,b` restricts the run to shapes whose name - // contains one of the given substrings, to reproduce or bisect quickly. - let shape_filter: Vec = std::env::var("AGGREGATE_CHAIN_SHAPES") - .map(|value| value.split(',').map(str::to_string).collect()) - .unwrap_or_default(); - let mut cases = vec![]; - for shape in all_shapes().into_iter().filter(|shape| { - shape_filter.is_empty() - || shape_filter - .iter() - .any(|needle| shape.name().contains(needle)) - }) { - for order in shape.orders() { - let skip_partial_variants: &[bool] = - if shape.has_skip_partial_candidate(order) { - &[true, false] - } else { - &[true] - }; - for cardinality in Cardinality::ALL { - for memory in [Memory::Unlimited, Memory::Limited] { - for &skip_partial_enabled in skip_partial_variants { - cases.push(Case { - shape, - params: CaseParams { - order, - cardinality, - memory, - skip_partial_enabled, - }, - }); - } - } - } - } - } - cases -} - -// --------------------------------------------------------------------------- -// Data generation -// --------------------------------------------------------------------------- - -/// `k1 Int64 nullable, k2 Int64 nullable, v Int64` -/// `k1, k2 Int64` (two-key query), `v Int64` (aggregated), and one column per -/// key type: `b Boolean`, `s Utf8`, `sv Utf8View`, `p Int64`, and -/// `st Struct, num: Int64>`. -/// Every key column is nullable. -fn schema() -> SchemaRef { - Arc::new(Schema::new(vec![ - Field::new("k1", DataType::Int64, true), - Field::new("k2", DataType::Int64, true), - Field::new("v", DataType::Int64, false), - Field::new("b", DataType::Boolean, true), - Field::new("s", DataType::Utf8, true), - Field::new("sv", DataType::Utf8View, true), - Field::new("p", DataType::Int64, true), - Field::new_struct("st", struct_fields(), true), - ])) -} - -/// About 3% nulls. -fn not_null(rng: &mut StdRng) -> bool { - rng.random_range(0..100) >= 3 -} - -fn struct_fields() -> Fields { - Fields::from(vec![ - Field::new("list", DataType::new_list(DataType::Int64, true), true), - Field::new("num", DataType::Int64, true), - ]) +/// ```text +/// AggregateExec: mode=Single, lim=[TOP_K_LIMIT] +/// DataSourceExec: partitions=1 +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn top_k_single() { + assert_chain_matches_single_aggregate(chain("top_k_single", &[TopK(Single)], 1)) + .await; } -/// The raw rows for one cardinality, deterministic per seed. The same multiset -/// is used for every `Order` and `Shape` so results are comparable. +/// Planner shape for `GROUP BY ... ORDER BY max(v) LIMIT n`: the limit lands on +/// the final stage. /// -/// Requirements: -/// - exactly `ROWS` rows -/// - `cardinality.groups()` distinct `(k1, k2)` pairs, spread so that `k1` -/// alone has fewer distinct values than `(k1, k2)`. Otherwise -/// `SortedByFirstKey` degenerates into `SortedByAllKeys`. -/// - some nulls in `k1` and `k2` -fn generate_rows(cardinality: Cardinality, seed: u64) -> RecordBatch { - let mut rng = StdRng::seed_from_u64(seed); - let groups = cardinality.groups(); - // `k2` cycles through at most sqrt(groups) values, so `k1` alone has fewer - // distinct values than the `(k1, k2)` pair. - let k2_values = (groups as f64).sqrt().ceil().max(2.0) as i64; - - let mut k1 = Vec::with_capacity(ROWS); - let mut k2 = Vec::with_capacity(ROWS); - let mut v = Vec::with_capacity(ROWS); - let mut b = Vec::with_capacity(ROWS); - let mut s = Vec::with_capacity(ROWS); - let mut sv = Vec::with_capacity(ROWS); - let mut p = Vec::with_capacity(ROWS); - let mut st_list = ListBuilder::new(Int64Builder::new()); - let mut st_num = Vec::with_capacity(ROWS); - let mut st_valid = Vec::with_capacity(ROWS); - for row in 0..ROWS { - let group = (row % groups) as i64; - k1.push(not_null(&mut rng).then_some(group / k2_values)); - k2.push(not_null(&mut rng).then_some(group % k2_values)); - v.push(rng.random_range(-1_000i64..1_000)); - // every key-type column has `groups` distinct values (boolean: two) - b.push(not_null(&mut rng).then_some(group % 2 == 0)); - s.push(not_null(&mut rng).then(|| format!("s{group:06}"))); - sv.push(not_null(&mut rng).then(|| format!("sv{group:06}"))); - p.push(not_null(&mut rng).then_some(group)); - // struct { list: [group, group + 1], [] or null; num: group or null } - st_valid.push(not_null(&mut rng)); - match rng.random_range(0..100) { - 0..3 => st_list.append_null(), - 3..6 => st_list.append(true), - _ => { - st_list.values().append_value(group); - st_list.values().append_value(group + 1); - st_list.append(true); - } - } - st_num.push(not_null(&mut rng).then_some(group)); - } - let st = StructArray::try_new( - struct_fields(), - vec![ - Arc::new(st_list.finish()), - Arc::new(Int64Array::from(st_num)), - ], - Some(NullBuffer::from(st_valid)), - ) - .unwrap(); - - RecordBatch::try_new( - schema(), - vec![ - Arc::new(Int64Array::from(k1)), - Arc::new(Int64Array::from(k2)), - Arc::new(Int64Array::from(v)), - Arc::new(BooleanArray::from(b)), - Arc::new(StringArray::from(s)), - Arc::new(StringViewArray::from(sv)), - Arc::new(Int64Array::from(p)), - Arc::new(st), - ], - ) - .unwrap() +/// ```text +/// AggregateExec: mode=FinalPartitioned, lim=[TOP_K_LIMIT] +/// RepartitionExec: partitioning=Hash(keys) +/// AggregateExec: mode=Partial +/// DataSourceExec: partitions=PARTITIONS +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn top_k_partial_repartition_final() { + assert_chain_matches_single_aggregate(chain( + "top_k_partial_repartition_final", + &[Aggregate(Partial), HashRepartition, TopK(FinalPartitioned)], + PARTITIONS, + )) + .await; } -/// Arrange `rows` for the given `order` and split into `partitions` partitions -/// of `BATCH_SIZE` batches. +/// TopK final stage on one coalesced partition. /// -/// - `Unordered`: shuffle rows, round-robin into partitions -/// - `SortedByFirstKey`: sort by `k1` (nulls first), contiguous slice per partition -/// - `SortedByAllKeys`: sort by `k1, k2` (nulls first), contiguous slice per partition -/// -/// Every partition individually satisfies the ordering. -fn arrange( - rows: &RecordBatch, - keys: Keys, - order: Order, - partitions: usize, -) -> Vec> { - let schema = rows.schema(); - let per_partition: Vec = match source_ordering(&schema, keys, order) { - None => { - let mut permutation: Vec = (0..rows.num_rows() as u32).collect(); - permutation.shuffle(&mut StdRng::seed_from_u64(0)); - let shuffled = - take_record_batch(rows, &UInt32Array::from(permutation)).unwrap(); - (0..partitions) - .map(|partition| { - let indices: UInt32Array = (partition as u32 - ..shuffled.num_rows() as u32) - .step_by(partitions) - .collect(); - take_record_batch(&shuffled, &indices).unwrap() - }) - .collect() - } - Some(ordering) => { - let sort_columns: Vec = ordering - .iter() - .map(|sort_expr| SortColumn { - values: sort_expr - .expr - .evaluate(rows) - .unwrap() - .into_array(rows.num_rows()) - .unwrap(), - options: Some(sort_expr.options), - }) - .collect(); - let indices = lexsort_to_indices(&sort_columns, None).unwrap(); - let sorted = take_record_batch(rows, &indices).unwrap(); - let per_partition = sorted.num_rows().div_ceil(partitions); - (0..partitions) - .map(|partition| { - let start = (partition * per_partition).min(sorted.num_rows()); - let length = per_partition.min(sorted.num_rows() - start); - copy_rows(&sorted, start, length) - }) - .collect() - } - }; - - // Copy every batch into its own buffers, as a real scan would produce. - // A slice shares the whole partition's buffers, and operators that - // account batches by `get_array_memory_size` (RepartitionExec, the merge) - // would charge every 64-row batch the size of the entire partition. - per_partition - .iter() - .map(|partition| { - (0..partition.num_rows()) - .step_by(BATCH_SIZE) - .map(|start| { - copy_rows( - partition, - start, - BATCH_SIZE.min(partition.num_rows() - start), - ) - }) - .collect() - }) - .collect() +/// ```text +/// AggregateExec: mode=Final, lim=[TOP_K_LIMIT] +/// CoalescePartitionsExec +/// AggregateExec: mode=Partial +/// DataSourceExec: partitions=PARTITIONS +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn top_k_partial_coalesce_final() { + assert_chain_matches_single_aggregate(chain( + "top_k_partial_coalesce_final", + &[Aggregate(Partial), CoalescePartitions, TopK(Final)], + PARTITIONS, + )) + .await; } -/// `batch[start..start + length]` in its own buffers. `take` copies where -/// `slice` shares and `concat_batches` of one batch only slices. +/// TopK on both stages. /// -/// Take from the unsliced batch: `take` on a list sizes the new values buffer -/// as child length / list length * taken rows, so taking 64 rows out of a -/// 64-row slice of a 32k-row list allocates a values buffer for the whole -/// child, and `get_array_memory_size` reports capacity. That charged every -/// batch about 500 KB instead of 5 KB. -fn copy_rows(batch: &RecordBatch, start: usize, length: usize) -> RecordBatch { - let indices = UInt32Array::from_iter_values(start as u32..(start + length) as u32); - take_record_batch(batch, &indices).unwrap() -} - -// --------------------------------------------------------------------------- -// Plan construction -// --------------------------------------------------------------------------- - -fn sort_expr(schema: &Schema, column: &str) -> PhysicalSortExpr { - PhysicalSortExpr::new( - col(column, schema).unwrap(), - SortOptions { - descending: false, - nulls_first: true, - }, - ) -} - -/// The ordering the source declares for `order`. -fn source_ordering(schema: &Schema, keys: Keys, order: Order) -> Option { - let keys = keys.columns(); - let sort_columns: &[&str] = match order { - Order::Unordered => return None, - Order::SortedByFirstKey => &keys[..1], - Order::SortedByAllKeys => keys, - }; - LexOrdering::new(sort_columns.iter().map(|column| sort_expr(schema, column))) -} - -fn source( - partitions: &[Vec], - keys: Keys, - order: Order, -) -> Arc { - let schema = schema(); - let mut memory_source = - MemorySourceConfig::try_new(partitions, Arc::clone(&schema), None).unwrap(); - if let Some(ordering) = source_ordering(&schema, keys, order) { - memory_source = memory_source - .try_with_sort_information(vec![ordering]) - .unwrap(); - } - DataSourceExec::from_data_source(memory_source) -} - -fn group_by(schema: &Schema, keys: Keys) -> PhysicalGroupBy { - PhysicalGroupBy::new_single( - keys.columns() - .iter() - .map(|key| (col(key, schema).unwrap(), key.to_string())) - .collect(), - ) -} - -fn aggregates(schema: &SchemaRef, query: Query) -> Vec> { - let value_column = || vec![col("v", schema).unwrap()]; - let build = |builder: AggregateExprBuilder, alias: &str| { - Arc::new( - builder - .schema(Arc::clone(schema)) - .alias(alias) - .build() - .unwrap(), - ) - }; - if query.aggregates == Aggregates::None { - return vec![]; - } - if query.aggregates == Aggregates::Max { - // TopK supports exactly one min/max aggregate over a non-nullable input - return vec![build( - AggregateExprBuilder::new(max_udaf(), value_column()), - "max", - )]; - } - vec![ - build( - AggregateExprBuilder::new(count_udaf(), value_column()), - "count", - ), - build( - AggregateExprBuilder::new(count_udaf(), value_column()).distinct(), - "count_distinct", - ), - build(AggregateExprBuilder::new(sum_udaf(), value_column()), "sum"), - // avg has no Int64 groups accumulator; the values are small integers so - // the Float64 sum stays exact and the result is order-independent. - build( - AggregateExprBuilder::new( - avg_udaf(), - vec![cast(col("v", schema).unwrap(), schema, DataType::Float64).unwrap()], - ), - "avg", - ), - build(AggregateExprBuilder::new(min_udaf(), value_column()), "min"), - build(AggregateExprBuilder::new(max_udaf(), value_column()), "max"), - ] -} - -/// Folds `shape.operators` bottom-up into a plan. The group-by, aggregate -/// expressions and hash keys are rewritten after every aggregate stage so the -/// next stage consumes that stage's output. -fn build_plan(shape: &Shape, input: Arc) -> Arc { - let input_schema = schema(); - let mut plan = input; - let mut group_by = group_by(&input_schema, shape.query.keys); - let mut aggregates = aggregates(&input_schema, shape.query); - let mut hash_keys: Vec> = group_by.input_exprs(); - - for operator in shape.chain.operators { - plan = match operator { - Aggregate(mode) | TopK(mode) => { - let limit_options = matches!(operator, TopK(_)) - .then(|| LimitOptions::new_with_order(TOP_K_LIMIT, true)); - let aggregate = Arc::new( - AggregateExec::try_new( - *mode, - group_by.clone(), - aggregates.clone(), - vec![None; aggregates.len()], - plan, - Arc::clone(&input_schema), - ) - .unwrap() - .with_limit_options(limit_options), - ); - group_by = aggregate.group_expr().as_final(); - aggregates = aggregate.aggr_expr().to_vec(); - hash_keys = aggregate.output_group_expr(); - aggregate - } - HashRepartition => Arc::new( - RepartitionExec::try_new( - plan, - Partitioning::Hash(hash_keys.clone(), PARTITIONS), - ) - .unwrap(), - ), - OrderPreservingHashRepartition => Arc::new( - RepartitionExec::try_new( - plan, - Partitioning::Hash(hash_keys.clone(), PARTITIONS), - ) - .unwrap() - .with_preserve_order(), - ), - CoalescePartitions => Arc::new(CoalescePartitionsExec::new(plan)), - SortPreservingMerge => { - let ordering = plan.properties().output_ordering().cloned().unwrap(); - Arc::new(SortPreservingMergeExec::new(ordering, plan)) - } - }; - } - plan -} - -// --------------------------------------------------------------------------- -// Execution context -// --------------------------------------------------------------------------- - -fn task_context(case: &Case) -> Arc { - let config = SessionConfig::new() - .with_batch_size(BATCH_SIZE) - .with_target_partitions(PARTITIONS) - // The default is 100k rows. Lower it so the skip-partial probe can - // fire on our per-partition row counts. A ratio threshold of 1.0 - // disables the probe entirely. - .set_usize( - "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", - 1024, - ); - let mut config = config; - config - .options_mut() - .execution - .skip_partial_aggregation_probe_ratio_threshold = - if case.params.skip_partial_enabled { - 0.8 - } else { - 1.0 - }; - - let runtime = match case.params.memory { - // Not using UnboundedMemoryPool, so users would still think that we have a valid pool, but just with enough memory - Memory::Unlimited => RuntimeEnvBuilder::new().with_memory_limit(usize::MAX, 1.0), - // Small enough that a very-high-cardinality final table spills, large - // enough that the legacy stream can still reserve its sort headroom - // and that RepartitionExec / SortPreservingMergeExec succeed. The - // fair pool keeps one stage from starving the others. - Memory::Limited => { - RuntimeEnvBuilder::new().with_memory_pool(Arc::new(TrackConsumersPool::new( - FairSpillPool::new(LIMITED_POOL_BYTES), - NonZeroUsize::new(5).unwrap(), - ))) - } - } - .build_arc() - .unwrap(); - - Arc::new( - TaskContext::default() - .with_session_config(config) - .with_runtime(runtime), - ) -} - -// --------------------------------------------------------------------------- -// Assertions -// --------------------------------------------------------------------------- - -/// All `AggregateExec` nodes in the plan, bottom-up. -fn aggregate_nodes(plan: &Arc) -> Vec> { - let mut nodes = vec![]; - let mut node = Arc::clone(plan); - loop { - if node.downcast_ref::().is_some() { - nodes.push(Arc::clone(&node)); - } - match node.children().first() { - Some(child) => node = Arc::clone(child), - None => break, - } - } - nodes.reverse(); - nodes -} - -fn as_aggregate(node: &Arc) -> &AggregateExec { - node.downcast_ref::().unwrap() -} - -/// Expected source order seen by each aggregate stage, bottom-up. Ordering is -/// lost at `HashRepartition` and `CoalescePartitions`, and kept by the -/// order-preserving shuffles and by aggregate stages themselves. -fn expected_orders(shape: &Shape, source_order: Order) -> Vec { - let mut current = source_order; - let mut expected = vec![]; - for operator in shape.chain.operators { - match operator { - HashRepartition | CoalescePartitions => current = Order::Unordered, - // `AggregateExec::try_new` forces `InputOrderMode::Linear` for - // partial reduce, since it emits its groups in hash table order, - // and it advertises no output ordering either. Everything above it - // is unordered until something sorts again. - Aggregate(PartialReduce) => { - expected.push(Order::Unordered); - current = Order::Unordered; - } - Aggregate(_) | TopK(_) => expected.push(current), - OrderPreservingHashRepartition | SortPreservingMerge => {} - } - } - expected -} - -fn order_matches(query: Query, expected: Order, actual: &InputOrderMode) -> bool { - // With a single group key, sorting by the first key already covers every - // group key. - let single_key = query.keys.columns().len() == 1; - match (expected, actual) { - (Order::Unordered, InputOrderMode::Linear) => true, - (Order::SortedByFirstKey, InputOrderMode::PartiallySorted(indices)) => { - !single_key && indices == &[0] - } - (Order::SortedByFirstKey, InputOrderMode::Sorted) => single_key, - (Order::SortedByAllKeys, InputOrderMode::Sorted) => true, - _ => false, - } -} - -/// Whether this stage's stream is allowed to spill. -fn can_spill(aggregate: &AggregateExec) -> bool { - if aggregate.limit_options().is_some() { - // GroupedTopKAggregateStream keeps a bounded heap and never spills - return false; - } - let spilling_mode = match aggregate.mode() { - Final | FinalPartitioned | Single | SinglePartitioned => true, - // Both partial streams emit their state early instead of spilling. - PartialReduce | Partial => false, - }; - let has_groups = !aggregate.group_expr().is_empty(); - spilling_mode && has_groups && *aggregate.input_order_mode() != InputOrderMode::Sorted -} - -/// Whether this stage runs the skip-partial probe. -fn runs_skip_partial_probe(aggregate: &AggregateExec) -> bool { - *aggregate.mode() == Partial - && aggregate.limit_options().is_none() - && !aggregate.group_expr().is_empty() - && *aggregate.input_order_mode() == InputOrderMode::Linear -} - -fn check_plan_shape(case: &Case, plan: &Arc) { - if case.shape.query.keys == Keys::None { - return; - } - let nodes = aggregate_nodes(plan); - let expected = expected_orders(&case.shape, case.params.order); - assert_eq!(nodes.len(), expected.len(), "{case:?}"); - for (node, expected_order) in nodes.iter().zip(expected) { - let aggregate = as_aggregate(node); - assert!( - order_matches( - case.shape.query, - expected_order, - aggregate.input_order_mode() - ), - "{case:?}: expected {expected_order:?} got {:?}\n{}", - aggregate.input_order_mode(), - displayable(plan.as_ref()).indent(true) - ); - } -} - -/// Returns a description of every stage that spilled, bottom-up, such as -/// `Final(Linear)`. -fn check_metrics(case: &Case, plan: &Arc) -> Vec { - let mut spilled = vec![]; - for node in aggregate_nodes(plan) { - let aggregate = as_aggregate(&node); - let mode = aggregate.mode(); - let metrics = node.metrics().unwrap(); - let spill_count = metrics.spill_count().unwrap_or(0); - if spill_count > 0 { - spilled.push(format!("{mode:?}({:?})", aggregate.input_order_mode())); - } - let skipped_rows = metrics - .sum_by_name("skipped_aggregation_rows") - .map(|metric| metric.as_usize()) - .unwrap_or(0); - - match case.params.memory { - Memory::Unlimited => { - assert_eq!(spill_count, 0, "{case:?}: unexpected spill in {mode:?}"); - } - Memory::Limited => { - // Whether a spilling-capable stage actually spills depends on - // the pool geometry, so only the run-wide coverage check in the - // driver requires it. Streams that cannot spill must not. - if !can_spill(aggregate) { - assert_eq!(spill_count, 0, "{case:?}: {mode:?} must never spill"); - } - } - } - - // Boolean keys have two groups whatever `cardinality` says, far - // below the ratio. - if case.params.memory == Memory::Unlimited - && case.params.cardinality == Cardinality::VeryHigh - && case.shape.query.keys.tracks_cardinality() - && case.params.skip_partial_enabled - && runs_skip_partial_probe(aggregate) - { - assert!( - skipped_rows > 0, - "{case:?}: skip-partial probe did not fire" - ); - } - if !case.params.skip_partial_enabled || !runs_skip_partial_probe(aggregate) { - assert_eq!(skipped_rows, 0, "{case:?}: skip-partial fired in {mode:?}"); - } - } - spilled +/// ```text +/// AggregateExec: mode=FinalPartitioned, lim=[TOP_K_LIMIT] +/// RepartitionExec: partitioning=Hash(keys) +/// AggregateExec: mode=Partial, lim=[TOP_K_LIMIT] +/// DataSourceExec: partitions=PARTITIONS +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn top_k_both_stages() { + assert_chain_matches_single_aggregate(chain( + "top_k_both_stages", + &[TopK(Partial), HashRepartition, TopK(FinalPartitioned)], + PARTITIONS, + )) + .await; } // --------------------------------------------------------------------------- @@ -1210,7 +575,7 @@ async fn run_case_inner(case: &Case, partitions: &[Vec]) -> Outcome fn reference_case(query: Query, cardinality: Cardinality) -> Case { Case { shape: Shape { - chain: chain_by_name("single"), + chain: SINGLE, query, }, params: CaseParams { @@ -1222,8 +587,11 @@ fn reference_case(query: Query, cardinality: Cardinality) -> Case { } } -#[tokio::test(flavor = "multi_thread")] -async fn aggregate_chain_fuzz() { +/// Runs every case of `chain` and asserts each returns the rows of the +/// `SINGLE` chain for its query, see the test preamble for the full list of +/// checks. A failure does not stop the run, so one run reports every failing +/// case. +async fn assert_chain_matches_single_aggregate(chain: Chain) { const SEED: u64 = 42; let mut total_spilled = 0; let mut failures: Vec = vec![]; @@ -1234,7 +602,7 @@ async fn aggregate_chain_fuzz() { for cardinality in Cardinality::ALL { let rows = generate_rows(cardinality, SEED); - let cases: Vec = all_cases() + let cases: Vec = cases(chain) .into_iter() .filter(|case| case.params.cardinality == cardinality) .collect(); @@ -1287,12 +655,17 @@ async fn aggregate_chain_fuzz() { total_spilled += spilled.len(); } // A shape filter may select only shapes that cannot spill - if std::env::var("AGGREGATE_CHAIN_SHAPES").is_err() { - assert!(total_spilled > 0, "no case exercised the spill path"); + if chain.expects_spill() && std::env::var("AGGREGATE_CHAIN_SHAPES").is_err() { + assert!( + total_spilled > 0, + "{}: no case exercised the spill path", + chain.name + ); } assert!( failures.is_empty(), - "{} cases failed:\n\n{}", + "{}: {} cases failed:\n\n{}", + chain.name, failures.len(), failures.join("\n\n") ); @@ -1303,8 +676,7 @@ async fn aggregate_chain_fuzz() { /// case there. Generous, so only a real hang fires it. const CASE_TIMEOUT_SECS: u64 = 600; -/// Waits for one case and files it under spilled, finished or failed. A -/// failure does not stop the run, so one run reports every failing case. +/// Waits for one case and files it under spilled, finished or failed. async fn collect_finished( join_set: &mut JoinSet<(Case, Vec)>, spilled: &mut Vec<(Case, Vec)>, @@ -1356,36 +728,3 @@ fn print_cases(cardinality: Cardinality, outcome: &str, cases: &[(Case, Vec Result<()> { - let shape = Shape { - chain: chain_by_name(chain_name), - query: Query { keys, aggregates }, - }; - let rows = generate_rows(cardinality, seed); - let reference = reference_case(shape.query, cardinality); - let case = Case { - shape, - params: CaseParams { - order, - cardinality, - memory, - skip_partial_enabled: true, - }, - }; - let inputs = Arc::new(arrange_all(&rows, [&case, &reference].into_iter())); - let expected = run_case(reference, Arc::clone(&inputs)).await; - let actual = run_case(case, inputs).await; - assert_eq!(actual.output, expected.output); - Ok(()) -} diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/assertions.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/assertions.rs new file mode 100644 index 0000000000000..f65e58ef3feb7 --- /dev/null +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/assertions.rs @@ -0,0 +1,181 @@ +// 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. +//! Assertions on plan shape and metrics. + +use super::*; + +/// All `AggregateExec` nodes in the plan, bottom-up. +pub(super) fn aggregate_nodes( + plan: &Arc, +) -> Vec> { + let mut nodes = vec![]; + let mut node = Arc::clone(plan); + loop { + if node.downcast_ref::().is_some() { + nodes.push(Arc::clone(&node)); + } + match node.children().first() { + Some(child) => node = Arc::clone(child), + None => break, + } + } + nodes.reverse(); + nodes +} + +pub(super) fn as_aggregate(node: &Arc) -> &AggregateExec { + node.downcast_ref::().unwrap() +} + +/// Expected source order seen by each aggregate stage, bottom-up. Ordering is +/// lost at `HashRepartition` and `CoalescePartitions`, and kept by the +/// order-preserving shuffles and by aggregate stages themselves. +pub(super) fn expected_orders(shape: &Shape, source_order: Order) -> Vec { + let mut current = source_order; + let mut expected = vec![]; + for operator in shape.chain.operators { + match operator { + HashRepartition | CoalescePartitions => current = Order::Unordered, + // `AggregateExec::try_new` forces `InputOrderMode::Linear` for + // partial reduce, since it emits its groups in hash table order, + // and it advertises no output ordering either. Everything above it + // is unordered until something sorts again. + Aggregate(PartialReduce) => { + expected.push(Order::Unordered); + current = Order::Unordered; + } + Aggregate(_) | TopK(_) => expected.push(current), + OrderPreservingHashRepartition | SortPreservingMerge => {} + } + } + expected +} + +pub(super) fn order_matches( + query: Query, + expected: Order, + actual: &InputOrderMode, +) -> bool { + // With a single group key, sorting by the first key already covers every + // group key. + let single_key = query.keys.columns().len() == 1; + match (expected, actual) { + (Order::Unordered, InputOrderMode::Linear) => true, + (Order::SortedByFirstKey, InputOrderMode::PartiallySorted(indices)) => { + !single_key && indices == &[0] + } + (Order::SortedByFirstKey, InputOrderMode::Sorted) => single_key, + (Order::SortedByAllKeys, InputOrderMode::Sorted) => true, + _ => false, + } +} + +/// Whether this stage's stream is allowed to spill. +pub(super) fn can_spill(aggregate: &AggregateExec) -> bool { + if aggregate.limit_options().is_some() { + // GroupedTopKAggregateStream keeps a bounded heap and never spills + return false; + } + let spilling_mode = match aggregate.mode() { + Final | FinalPartitioned | Single | SinglePartitioned => true, + // Both partial streams emit their state early instead of spilling. + PartialReduce | Partial => false, + }; + let has_groups = !aggregate.group_expr().is_empty(); + spilling_mode && has_groups && *aggregate.input_order_mode() != InputOrderMode::Sorted +} + +/// Whether this stage runs the skip-partial probe. +pub(super) fn runs_skip_partial_probe(aggregate: &AggregateExec) -> bool { + *aggregate.mode() == Partial + && aggregate.limit_options().is_none() + && !aggregate.group_expr().is_empty() + && *aggregate.input_order_mode() == InputOrderMode::Linear +} + +pub(super) fn check_plan_shape(case: &Case, plan: &Arc) { + if case.shape.query.keys == Keys::None { + return; + } + let nodes = aggregate_nodes(plan); + let expected = expected_orders(&case.shape, case.params.order); + assert_eq!(nodes.len(), expected.len(), "{case:?}"); + for (node, expected_order) in nodes.iter().zip(expected) { + let aggregate = as_aggregate(node); + assert!( + order_matches( + case.shape.query, + expected_order, + aggregate.input_order_mode() + ), + "{case:?}: expected {expected_order:?} got {:?}\n{}", + aggregate.input_order_mode(), + displayable(plan.as_ref()).indent(true) + ); + } +} + +/// Returns a description of every stage that spilled, bottom-up, such as +/// `Final(Linear)`. +pub(super) fn check_metrics(case: &Case, plan: &Arc) -> Vec { + let mut spilled = vec![]; + for node in aggregate_nodes(plan) { + let aggregate = as_aggregate(&node); + let mode = aggregate.mode(); + let metrics = node.metrics().unwrap(); + let spill_count = metrics.spill_count().unwrap_or(0); + if spill_count > 0 { + spilled.push(format!("{mode:?}({:?})", aggregate.input_order_mode())); + } + let skipped_rows = metrics + .sum_by_name("skipped_aggregation_rows") + .map(|metric| metric.as_usize()) + .unwrap_or(0); + + match case.params.memory { + Memory::Unlimited => { + assert_eq!(spill_count, 0, "{case:?}: unexpected spill in {mode:?}"); + } + Memory::Limited => { + // Whether a spilling-capable stage actually spills depends on + // the pool geometry, so only the run-wide coverage check in the + // driver requires it. Streams that cannot spill must not. + if !can_spill(aggregate) { + assert_eq!(spill_count, 0, "{case:?}: {mode:?} must never spill"); + } + } + } + + // Boolean keys have two groups whatever `cardinality` says, far + // below the ratio. + if case.params.memory == Memory::Unlimited + && case.params.cardinality == Cardinality::VeryHigh + && case.shape.query.keys.tracks_cardinality() + && case.params.skip_partial_enabled + && runs_skip_partial_probe(aggregate) + { + assert!( + skipped_rows > 0, + "{case:?}: skip-partial probe did not fire" + ); + } + if !case.params.skip_partial_enabled || !runs_skip_partial_probe(aggregate) { + assert_eq!(skipped_rows, 0, "{case:?}: skip-partial fired in {mode:?}"); + } + } + spilled +} diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/case_space.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/case_space.rs new file mode 100644 index 0000000000000..5b52814462662 --- /dev/null +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/case_space.rs @@ -0,0 +1,400 @@ +// 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. +//! The case space: every axis a case varies along and the valid combinations. + +use super::*; + +pub(super) const ROWS: usize = 32 * 1024; +pub(super) const PARTITIONS: usize = 4; +pub(super) const BATCH_SIZE: usize = 64; +/// The fair pool caps every spillable consumer at `pool / consumers`, and a +/// chain registers up to twenty consumers (aggregate streams plus one per +/// repartition channel). The cap has to clear a small table's legitimate +/// footprint, which at very low cardinality is dominated by the `count +/// distinct` sets and grows in steps of roughly 100 KB, while a final table at +/// very high cardinality must still exceed it. +pub(super) const LIMITED_POOL_BYTES: usize = 4 * 1024 * 1024; + +/// How the source data is ordered relative to the group keys. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub(super) enum Order { + /// Not ordered. Aggregates see `InputOrderMode::Linear`. + Unordered, + /// Sorted by the first key only. Aggregates see + /// `InputOrderMode::PartiallySorted([0])`. + SortedByFirstKey, + /// Sorted by all keys. Aggregates see `InputOrderMode::Sorted`. + SortedByAllKeys, +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub(super) enum Cardinality { + VeryHigh, + Medium, + Low, + VeryLow, +} + +impl Cardinality { + pub(super) const ALL: [Self; 4] = + [Self::VeryHigh, Self::Medium, Self::Low, Self::VeryLow]; + + /// Number of distinct `(k1, k2)` groups. + pub(super) fn groups(self) -> usize { + match self { + Self::VeryHigh => ROWS, + Self::Medium => ROWS / 32, + Self::Low => 16, + Self::VeryLow => 2, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub(super) enum Memory { + /// Unlimited pool. Nothing spills or emits early. + Unlimited, + /// Pool sized so final and single hash tables cannot fit. + Limited, +} + +/// One operator in a chain, listed bottom to top. +#[derive(Clone, Copy, Debug)] +pub(super) enum Operator { + Aggregate(AggregateMode), + /// `AggregateExec` with `limit_options` set, which selects + /// `GroupedTopKAggregateStream` regardless of mode. + TopK(AggregateMode), + /// `RepartitionExec` hashed on the group keys. Destroys ordering. + HashRepartition, + /// `RepartitionExec` hashed on the group keys with `preserve_order`. + OrderPreservingHashRepartition, + /// `CoalescePartitionsExec`. Destroys ordering. + CoalescePartitions, + /// `SortPreservingMergeExec` on the current ordering. + SortPreservingMerge, +} + +/// The `GROUP BY` keys. Every key type has its own `GroupValues` +/// implementation, so each is a value of this axis. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub(super) enum Keys { + /// No `GROUP BY`. + None, + /// `k1, k2` (two Int64), handled by `GroupValuesColumn`. + TwoInts, + /// `b` (Boolean), handled by `GroupValuesBoolean`. + Boolean, + /// `s` (Utf8), handled by `GroupValuesBytes`. + Bytes, + /// `sv` (Utf8View), handled by `GroupValuesBytesView`. + BytesView, + /// `p` (Int64 with as many distinct values as groups), handled by + /// `GroupValuesPrimitive`. + Primitive, + /// `b, s, sv, p`, handled by `GroupValuesColumn` with mixed column types. + Mixed, + /// `st` (Struct of a List and an Int64), which no specialized + /// implementation supports, so it falls back to the row format + /// `GroupValuesRows`. + Struct, +} + +impl Keys { + pub(super) const ALL: [Self; 8] = [ + Self::None, + Self::TwoInts, + Self::Boolean, + Self::Bytes, + Self::BytesView, + Self::Primitive, + Self::Mixed, + Self::Struct, + ]; + + /// Key columns, in `GROUP BY` order. + pub(super) fn columns(self) -> &'static [&'static str] { + match self { + Keys::None => &[], + Keys::TwoInts => &["k1", "k2"], + Keys::Boolean => &["b"], + Keys::Bytes => &["s"], + Keys::BytesView => &["sv"], + Keys::Primitive => &["p"], + Keys::Mixed => &["b", "s", "sv", "p"], + Keys::Struct => &["st"], + } + } + + /// Whether the source can be sorted by the keys. Struct columns cannot be + /// sorted by the arrow sort kernels, so those keys only run unordered. + pub(super) fn sortable(self) -> bool { + self != Keys::Struct + } + + /// Whether the number of groups is `Cardinality::groups()`. Every key + /// column has one distinct value per group except the Boolean one. + pub(super) fn tracks_cardinality(self) -> bool { + !matches!(self, Keys::None | Keys::Boolean) + } + + /// Whether `GroupedTopKAggregateStream` supports these keys: exactly one + /// primitive or string column. + pub(super) fn top_k_supported(self) -> bool { + matches!(self, Keys::Bytes | Keys::BytesView | Keys::Primitive) + } +} + +/// The aggregate expressions. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum Aggregates { + /// count, count distinct, sum, avg, min, max: non-trivial partial state so + /// the Partial, PartialReduce and Final stages are actually exercised. + /// `avg` (two-field state) and `count distinct` (set state) matter most. + All, + /// No aggregate expressions, as `SELECT DISTINCT` plans: the + /// accumulator-free path of every stream. + None, + /// `max(v)` only, the one aggregate the TopK stream supports. Chains using + /// `Operator::TopK` set a limit larger than any possible group count, so + /// the result must still be the complete aggregate. + Max, +} + +impl Aggregates { + pub(super) const ALL: [Self; 3] = [Self::All, Self::None, Self::Max]; +} + +/// The logical query a chain computes. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) struct Query { + pub(super) keys: Keys, + pub(super) aggregates: Aggregates, +} + +/// Larger than any possible number of groups, so TopK keeps every group. +pub(super) const TOP_K_LIMIT: usize = 2 * ROWS; + +/// Everything that varies for a case apart from the shape itself. Passed as a +/// struct so a new dimension does not change every shape predicate. +#[derive(Clone, Copy, Debug, PartialEq)] +pub(super) struct CaseParams { + pub(super) order: Order, + pub(super) cardinality: Cardinality, + pub(super) memory: Memory, + /// Whether the skip-partial probe may fire. Only varied for shapes with a + /// grouped `Partial` stage on Linear input, since nothing else runs it. + pub(super) skip_partial_enabled: bool, +} + +/// An operator chain, independent of the query it computes. +#[derive(Clone, Copy, Debug)] +pub(super) struct Chain { + pub(super) name: &'static str, + pub(super) operators: &'static [Operator], + /// Source partition count. + pub(super) source_partitions: usize, +} + +pub(super) const fn chain( + name: &'static str, + operators: &'static [Operator], + source_partitions: usize, +) -> Chain { + Chain { + name, + operators, + source_partitions, + } +} + +impl Chain { + /// Whether the chain hashes or sorts on the group keys, so it cannot run + /// without any. + pub(super) fn needs_keys(&self) -> bool { + self.operators.iter().any(|operator| { + matches!( + operator, + HashRepartition + | OrderPreservingHashRepartition + | SortPreservingMerge + | TopK(_) + ) + }) + } + + pub(super) fn is_top_k(&self) -> bool { + self.operators + .iter() + .any(|operator| matches!(operator, TopK(_))) + } + + /// Whether the chain keeps the source ordering through its shuffles, so + /// it only makes sense on ordered input. + pub(super) fn preserves_order(&self) -> bool { + self.operators.iter().any(|operator| { + matches!( + operator, + OrderPreservingHashRepartition | SortPreservingMerge + ) + }) + } + + /// Whether some case of this chain must spill: a final or single hash + /// stage on unordered input, whose very-high-cardinality table cannot fit + /// the limited pool. Ordered stages emit early or are bounded, and TopK + /// keeps a bounded heap, so those chains never spill. + pub(super) fn expects_spill(&self) -> bool { + !self.preserves_order() + && self.operators.iter().any(|operator| { + matches!( + operator, + Aggregate(Final | FinalPartitioned | Single | SinglePartitioned) + ) + }) + } +} + +/// A plan shape: a chain computing a query. +#[derive(Clone, Copy, Debug)] +pub(super) struct Shape { + pub(super) chain: Chain, + pub(super) query: Query, +} + +impl Shape { + pub(super) fn name(&self) -> String { + format!( + "{} {:?} {:?}", + self.chain.name, self.query.keys, self.query.aggregates + ) + } + + /// Source orders that make sense for this shape. Order-preserving shuffles + /// need an ordering to preserve; no-grouping chains ignore ordering. + pub(super) fn orders(&self) -> Vec { + let keys = self.query.keys.columns(); + let mut orders = vec![]; + if !self.chain.preserves_order() { + orders.push(Order::Unordered); + } + if self.query.keys.sortable() && !keys.is_empty() { + // With a single key, sorting by the first key is already sorting + // by all keys. + if keys.len() > 1 { + orders.push(Order::SortedByFirstKey); + } + orders.push(Order::SortedByAllKeys); + } + orders + } + + /// Whether some `Partial` stage of this shape runs the skip-partial probe + /// for the given source order: grouped, not TopK, and Linear input. + pub(super) fn has_skip_partial_candidate(&self, order: Order) -> bool { + if self.query.keys == Keys::None { + return false; + } + let mut current = order; + for operator in self.chain.operators { + match operator { + HashRepartition | CoalescePartitions => current = Order::Unordered, + Aggregate(Partial) if current == Order::Unordered => return true, + _ => {} + } + } + false + } +} + +/// Every query the chain can be planned for. +/// +/// - Without keys a chain can neither hash nor sort, and `max` alone is a +/// subset of the full aggregate list, so only that list runs. +/// - The TopK stream needs one primitive or string key with either a single +/// `max` or no aggregates at all (the `DISTINCT ... LIMIT` form). +/// - Every other chain runs the full aggregate list and the accumulator-free +/// form; `max` alone adds nothing there. +pub(super) fn shapes(chain: Chain) -> Vec { + let mut shapes = vec![]; + for keys in Keys::ALL { + for aggregates in Aggregates::ALL { + let valid = if keys == Keys::None { + !chain.needs_keys() && aggregates == Aggregates::All + } else if chain.is_top_k() { + keys.top_k_supported() && aggregates != Aggregates::All + } else { + aggregates != Aggregates::Max + }; + if valid { + shapes.push(Shape { + chain, + query: Query { keys, aggregates }, + }); + } + } + } + shapes +} + +#[derive(Clone, Debug)] +pub(super) struct Case { + pub(super) shape: Shape, + pub(super) params: CaseParams, +} + +/// Every case of the chain: each valid query over every source order, +/// cardinality, memory budget and skip-partial setting. +pub(super) fn cases(chain: Chain) -> Vec { + // `AGGREGATE_CHAIN_SHAPES=a,b` restricts the run to shapes whose name + // contains one of the given substrings, to reproduce or bisect quickly. + let shape_filter: Vec = std::env::var("AGGREGATE_CHAIN_SHAPES") + .map(|value| value.split(',').map(str::to_string).collect()) + .unwrap_or_default(); + let mut cases = vec![]; + for shape in shapes(chain).into_iter().filter(|shape| { + shape_filter.is_empty() + || shape_filter + .iter() + .any(|needle| shape.name().contains(needle)) + }) { + for order in shape.orders() { + let skip_partial_variants: &[bool] = + if shape.has_skip_partial_candidate(order) { + &[true, false] + } else { + &[true] + }; + for cardinality in Cardinality::ALL { + for memory in [Memory::Unlimited, Memory::Limited] { + for &skip_partial_enabled in skip_partial_variants { + cases.push(Case { + shape, + params: CaseParams { + order, + cardinality, + memory, + skip_partial_enabled, + }, + }); + } + } + } + } + } + cases +} diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/context.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/context.rs new file mode 100644 index 0000000000000..faf4bc2c94c81 --- /dev/null +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/context.rs @@ -0,0 +1,65 @@ +// 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. +//! Execution context: session config and memory pool per case. + +use super::*; + +pub(super) fn task_context(case: &Case) -> Arc { + let config = SessionConfig::new() + .with_batch_size(BATCH_SIZE) + .with_target_partitions(PARTITIONS) + // The default is 100k rows. Lower it so the skip-partial probe can + // fire on our per-partition row counts. A ratio threshold of 1.0 + // disables the probe entirely. + .set_usize( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + 1024, + ); + let mut config = config; + config + .options_mut() + .execution + .skip_partial_aggregation_probe_ratio_threshold = + if case.params.skip_partial_enabled { + 0.8 + } else { + 1.0 + }; + + let runtime = match case.params.memory { + // Not using UnboundedMemoryPool, so users would still think that we have a valid pool, but just with enough memory + Memory::Unlimited => RuntimeEnvBuilder::new().with_memory_limit(usize::MAX, 1.0), + // Small enough that a very-high-cardinality final table spills, large + // enough that the legacy stream can still reserve its sort headroom + // and that RepartitionExec / SortPreservingMergeExec succeed. The + // fair pool keeps one stage from starving the others. + Memory::Limited => { + RuntimeEnvBuilder::new().with_memory_pool(Arc::new(TrackConsumersPool::new( + FairSpillPool::new(LIMITED_POOL_BYTES), + NonZeroUsize::new(5).unwrap(), + ))) + } + } + .build_arc() + .unwrap(); + + Arc::new( + TaskContext::default() + .with_session_config(config) + .with_runtime(runtime), + ) +} diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/data.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/data.rs new file mode 100644 index 0000000000000..60f687ff06f63 --- /dev/null +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/data.rs @@ -0,0 +1,215 @@ +// 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. +//! Data generation: deterministic rows per seed, arranged per order and partition count. + +use super::*; + +/// `k1 Int64 nullable, k2 Int64 nullable, v Int64` +/// `k1, k2 Int64` (two-key query), `v Int64` (aggregated), and one column per +/// key type: `b Boolean`, `s Utf8`, `sv Utf8View`, `p Int64`, and +/// `st Struct, num: Int64>`. +/// Every key column is nullable. +pub(super) fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("k1", DataType::Int64, true), + Field::new("k2", DataType::Int64, true), + Field::new("v", DataType::Int64, false), + Field::new("b", DataType::Boolean, true), + Field::new("s", DataType::Utf8, true), + Field::new("sv", DataType::Utf8View, true), + Field::new("p", DataType::Int64, true), + Field::new_struct("st", struct_fields(), true), + ])) +} + +/// About 3% nulls. +pub(super) fn not_null(rng: &mut StdRng) -> bool { + rng.random_range(0..100) >= 3 +} + +pub(super) fn struct_fields() -> Fields { + Fields::from(vec![ + Field::new("list", DataType::new_list(DataType::Int64, true), true), + Field::new("num", DataType::Int64, true), + ]) +} + +/// The raw rows for one cardinality, deterministic per seed. The same multiset +/// is used for every `Order` and `Shape` so results are comparable. +/// +/// Requirements: +/// - exactly `ROWS` rows +/// - `cardinality.groups()` distinct `(k1, k2)` pairs, spread so that `k1` +/// alone has fewer distinct values than `(k1, k2)`. Otherwise +/// `SortedByFirstKey` degenerates into `SortedByAllKeys`. +/// - some nulls in `k1` and `k2` +pub(super) fn generate_rows(cardinality: Cardinality, seed: u64) -> RecordBatch { + let mut rng = StdRng::seed_from_u64(seed); + let groups = cardinality.groups(); + // `k2` cycles through at most sqrt(groups) values, so `k1` alone has fewer + // distinct values than the `(k1, k2)` pair. + let k2_values = (groups as f64).sqrt().ceil().max(2.0) as i64; + + let mut k1 = Vec::with_capacity(ROWS); + let mut k2 = Vec::with_capacity(ROWS); + let mut v = Vec::with_capacity(ROWS); + let mut b = Vec::with_capacity(ROWS); + let mut s = Vec::with_capacity(ROWS); + let mut sv = Vec::with_capacity(ROWS); + let mut p = Vec::with_capacity(ROWS); + let mut st_list = ListBuilder::new(Int64Builder::new()); + let mut st_num = Vec::with_capacity(ROWS); + let mut st_valid = Vec::with_capacity(ROWS); + for row in 0..ROWS { + let group = (row % groups) as i64; + k1.push(not_null(&mut rng).then_some(group / k2_values)); + k2.push(not_null(&mut rng).then_some(group % k2_values)); + v.push(rng.random_range(-1_000i64..1_000)); + // every key-type column has `groups` distinct values (boolean: two) + b.push(not_null(&mut rng).then_some(group % 2 == 0)); + s.push(not_null(&mut rng).then(|| format!("s{group:06}"))); + sv.push(not_null(&mut rng).then(|| format!("sv{group:06}"))); + p.push(not_null(&mut rng).then_some(group)); + // struct { list: [group, group + 1], [] or null; num: group or null } + st_valid.push(not_null(&mut rng)); + match rng.random_range(0..100) { + 0..3 => st_list.append_null(), + 3..6 => st_list.append(true), + _ => { + st_list.values().append_value(group); + st_list.values().append_value(group + 1); + st_list.append(true); + } + } + st_num.push(not_null(&mut rng).then_some(group)); + } + let st = StructArray::try_new( + struct_fields(), + vec![ + Arc::new(st_list.finish()), + Arc::new(Int64Array::from(st_num)), + ], + Some(NullBuffer::from(st_valid)), + ) + .unwrap(); + + RecordBatch::try_new( + schema(), + vec![ + Arc::new(Int64Array::from(k1)), + Arc::new(Int64Array::from(k2)), + Arc::new(Int64Array::from(v)), + Arc::new(BooleanArray::from(b)), + Arc::new(StringArray::from(s)), + Arc::new(StringViewArray::from(sv)), + Arc::new(Int64Array::from(p)), + Arc::new(st), + ], + ) + .unwrap() +} + +/// Arrange `rows` for the given `order` and split into `partitions` partitions +/// of `BATCH_SIZE` batches. +/// +/// - `Unordered`: shuffle rows, round-robin into partitions +/// - `SortedByFirstKey`: sort by `k1` (nulls first), contiguous slice per partition +/// - `SortedByAllKeys`: sort by `k1, k2` (nulls first), contiguous slice per partition +/// +/// Every partition individually satisfies the ordering. +pub(super) fn arrange( + rows: &RecordBatch, + keys: Keys, + order: Order, + partitions: usize, +) -> Vec> { + let schema = rows.schema(); + let per_partition: Vec = match source_ordering(&schema, keys, order) { + None => { + let mut permutation: Vec = (0..rows.num_rows() as u32).collect(); + permutation.shuffle(&mut StdRng::seed_from_u64(0)); + let shuffled = + take_record_batch(rows, &UInt32Array::from(permutation)).unwrap(); + (0..partitions) + .map(|partition| { + let indices: UInt32Array = (partition as u32 + ..shuffled.num_rows() as u32) + .step_by(partitions) + .collect(); + take_record_batch(&shuffled, &indices).unwrap() + }) + .collect() + } + Some(ordering) => { + let sort_columns: Vec = ordering + .iter() + .map(|sort_expr| SortColumn { + values: sort_expr + .expr + .evaluate(rows) + .unwrap() + .into_array(rows.num_rows()) + .unwrap(), + options: Some(sort_expr.options), + }) + .collect(); + let indices = lexsort_to_indices(&sort_columns, None).unwrap(); + let sorted = take_record_batch(rows, &indices).unwrap(); + let per_partition = sorted.num_rows().div_ceil(partitions); + (0..partitions) + .map(|partition| { + let start = (partition * per_partition).min(sorted.num_rows()); + let length = per_partition.min(sorted.num_rows() - start); + copy_rows(&sorted, start, length) + }) + .collect() + } + }; + + // Copy every batch into its own buffers, as a real scan would produce. + // A slice shares the whole partition's buffers, and operators that + // account batches by `get_array_memory_size` (RepartitionExec, the merge) + // would charge every 64-row batch the size of the entire partition. + per_partition + .iter() + .map(|partition| { + (0..partition.num_rows()) + .step_by(BATCH_SIZE) + .map(|start| { + copy_rows( + partition, + start, + BATCH_SIZE.min(partition.num_rows() - start), + ) + }) + .collect() + }) + .collect() +} + +/// `batch[start..start + length]` in its own buffers. `take` copies where +/// `slice` shares and `concat_batches` of one batch only slices. +/// +/// Take from the unsliced batch: `take` on a list sizes the new values buffer +/// as child length / list length * taken rows, so taking 64 rows out of a +/// 64-row slice of a 32k-row list allocates a values buffer for the whole +/// child, and `get_array_memory_size` reports capacity. That charged every +/// batch about 500 KB instead of 5 KB. +pub(super) fn copy_rows(batch: &RecordBatch, start: usize, length: usize) -> RecordBatch { + let indices = UInt32Array::from_iter_values(start as u32..(start + length) as u32); + take_record_batch(batch, &indices).unwrap() +} diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/plan.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/plan.rs new file mode 100644 index 0000000000000..dc12785793078 --- /dev/null +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/plan.rs @@ -0,0 +1,177 @@ +// 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. +//! Plan construction for a [`Shape`] over a source. + +use super::*; + +pub(super) fn sort_expr(schema: &Schema, column: &str) -> PhysicalSortExpr { + PhysicalSortExpr::new( + col(column, schema).unwrap(), + SortOptions { + descending: false, + nulls_first: true, + }, + ) +} + +/// The ordering the source declares for `order`. +pub(super) fn source_ordering( + schema: &Schema, + keys: Keys, + order: Order, +) -> Option { + let keys = keys.columns(); + let sort_columns: &[&str] = match order { + Order::Unordered => return None, + Order::SortedByFirstKey => &keys[..1], + Order::SortedByAllKeys => keys, + }; + LexOrdering::new(sort_columns.iter().map(|column| sort_expr(schema, column))) +} + +pub(super) fn source( + partitions: &[Vec], + keys: Keys, + order: Order, +) -> Arc { + let schema = schema(); + let mut memory_source = + MemorySourceConfig::try_new(partitions, Arc::clone(&schema), None).unwrap(); + if let Some(ordering) = source_ordering(&schema, keys, order) { + memory_source = memory_source + .try_with_sort_information(vec![ordering]) + .unwrap(); + } + DataSourceExec::from_data_source(memory_source) +} + +pub(super) fn group_by(schema: &Schema, keys: Keys) -> PhysicalGroupBy { + PhysicalGroupBy::new_single( + keys.columns() + .iter() + .map(|key| (col(key, schema).unwrap(), key.to_string())) + .collect(), + ) +} + +pub(super) fn aggregates( + schema: &SchemaRef, + query: Query, +) -> Vec> { + let value_column = || vec![col("v", schema).unwrap()]; + let build = |builder: AggregateExprBuilder, alias: &str| { + Arc::new( + builder + .schema(Arc::clone(schema)) + .alias(alias) + .build() + .unwrap(), + ) + }; + if query.aggregates == Aggregates::None { + return vec![]; + } + if query.aggregates == Aggregates::Max { + // TopK supports exactly one min/max aggregate over a non-nullable input + return vec![build( + AggregateExprBuilder::new(max_udaf(), value_column()), + "max", + )]; + } + vec![ + build( + AggregateExprBuilder::new(count_udaf(), value_column()), + "count", + ), + build( + AggregateExprBuilder::new(count_udaf(), value_column()).distinct(), + "count_distinct", + ), + build(AggregateExprBuilder::new(sum_udaf(), value_column()), "sum"), + // avg has no Int64 groups accumulator; the values are small integers so + // the Float64 sum stays exact and the result is order-independent. + build( + AggregateExprBuilder::new( + avg_udaf(), + vec![cast(col("v", schema).unwrap(), schema, DataType::Float64).unwrap()], + ), + "avg", + ), + build(AggregateExprBuilder::new(min_udaf(), value_column()), "min"), + build(AggregateExprBuilder::new(max_udaf(), value_column()), "max"), + ] +} + +/// Folds `shape.operators` bottom-up into a plan. The group-by, aggregate +/// expressions and hash keys are rewritten after every aggregate stage so the +/// next stage consumes that stage's output. +pub(super) fn build_plan( + shape: &Shape, + input: Arc, +) -> Arc { + let input_schema = schema(); + let mut plan = input; + let mut group_by = group_by(&input_schema, shape.query.keys); + let mut aggregates = aggregates(&input_schema, shape.query); + let mut hash_keys: Vec> = group_by.input_exprs(); + + for operator in shape.chain.operators { + plan = match operator { + Aggregate(mode) | TopK(mode) => { + let limit_options = matches!(operator, TopK(_)) + .then(|| LimitOptions::new_with_order(TOP_K_LIMIT, true)); + let aggregate = Arc::new( + AggregateExec::try_new( + *mode, + group_by.clone(), + aggregates.clone(), + vec![None; aggregates.len()], + plan, + Arc::clone(&input_schema), + ) + .unwrap() + .with_limit_options(limit_options), + ); + group_by = aggregate.group_expr().as_final(); + aggregates = aggregate.aggr_expr().to_vec(); + hash_keys = aggregate.output_group_expr(); + aggregate + } + HashRepartition => Arc::new( + RepartitionExec::try_new( + plan, + Partitioning::Hash(hash_keys.clone(), PARTITIONS), + ) + .unwrap(), + ), + OrderPreservingHashRepartition => Arc::new( + RepartitionExec::try_new( + plan, + Partitioning::Hash(hash_keys.clone(), PARTITIONS), + ) + .unwrap() + .with_preserve_order(), + ), + CoalescePartitions => Arc::new(CoalescePartitionsExec::new(plan)), + SortPreservingMerge => { + let ordering = plan.properties().output_ordering().cloned().unwrap(); + Arc::new(SortPreservingMergeExec::new(ordering, plan)) + } + }; + } + plan +} diff --git a/datafusion/core/tests/fuzz_cases/mod.rs b/datafusion/core/tests/fuzz_cases/mod.rs index 14dbacfa1a9a8..92b68361d007b 100644 --- a/datafusion/core/tests/fuzz_cases/mod.rs +++ b/datafusion/core/tests/fuzz_cases/mod.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -mod aggregate_chain_fuzz; #[expect(clippy::needless_pass_by_value)] mod aggregate_fuzz; mod distinct_count_string_fuzz; @@ -28,6 +27,7 @@ mod sort_fuzz; mod sort_query_fuzz; mod topk_filter_pushdown; +mod aggregate_chain_fuzz; mod aggregation_fuzzer; #[expect(clippy::needless_pass_by_value)] mod equivalence; From fddce09de6a981a2ce2f032b2254c44fc1797d28 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 17 Sep 2026 16:55:18 +0300 Subject: [PATCH 19/21] cleanup --- .../tests/fuzz_cases/aggregate_chain_fuzz.rs | 586 +++++++++++------- .../aggregate_chain_fuzz/case_space.rs | 214 ++++--- 2 files changed, 508 insertions(+), 292 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs index 1d733d9f21d14..ebf50154345de 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -84,11 +84,11 @@ use plan::*; // What is tested // ============== // -// Every test below is one physical plan shape, an `AggregateExec` chain. -// `chain` lists its operators bottom-up, source first; the doc comment shows -// the same plan as DataFusion prints it. The test runs that plan for every -// query and input below and asserts it returns the same rows as the plain -// single-stage aggregate `SINGLE`. +// Every test below is one physical plan shape, an `AggregateExec` chain, run +// over every combination of the axes listed in its `ChainTest`. Each run must +// return the same rows as the plain single-stage aggregate `SINGLE`. To test +// one case, narrow the lists in place. `chain` lists the operators bottom-up, +// source first; the doc comment shows the same plan as DataFusion prints it. // // The table (`data.rs`), 32K rows, generated from a fixed seed: // @@ -101,43 +101,41 @@ use plan::*; // p Int64 // st Struct, num: Int64> // -// Every key column has about 3% nulls. The number of groups is the -// cardinality axis: 32K (one row per group), 1K, 16 or 2. +// Every key column has about 3% nulls. `cardinalities` is the number of +// groups: 32K (one row per group), 1K, 16 or 2. // -// The queries (`case_space.rs`), one per key set times one per aggregate list: +// `group_by` picks the query, one per `GroupValues` implementation: // -// SELECT FROM t -- no GROUP BY -// SELECT k1, k2, FROM t GROUP BY k1, k2 -- GroupValuesColumn -// SELECT b, FROM t GROUP BY b -- GroupValuesBoolean -// SELECT s, FROM t GROUP BY s -- GroupValuesBytes -// SELECT sv, FROM t GROUP BY sv -- GroupValuesBytesView -// SELECT p, FROM t GROUP BY p -- GroupValuesPrimitive -// SELECT b, s, sv, p, FROM t GROUP BY b, s, sv, p -// SELECT st, FROM t GROUP BY st -- row format fallback +// Keys::None SELECT FROM t +// Keys::TwoInts SELECT k1, k2, FROM t GROUP BY k1, k2 -- Column +// Keys::Boolean SELECT b, FROM t GROUP BY b -- Boolean +// Keys::Bytes SELECT s, FROM t GROUP BY s -- Bytes +// Keys::BytesView SELECT sv, FROM t GROUP BY sv -- BytesView +// Keys::Primitive SELECT p, FROM t GROUP BY p -- Primitive +// Keys::Mixed SELECT b, s, sv, p, FROM t GROUP BY b, s, sv, p +// Keys::Struct SELECT st, FROM t GROUP BY st -- row fallback // -// is one of -// count(v), count(DISTINCT v), sum(v), avg(v), min(v), max(v) -// nothing, as in SELECT DISTINCT keys -// max(v) alone, the one aggregate the TopK stream supports +// `aggregates` picks : // -// TopK chains add LIMIT 64K, above any group count, so all groups survive. +// Aggregates::All count(v), count(DISTINCT v), sum(v), avg(v), min(v), max(v) +// Aggregates::None nothing, as in SELECT DISTINCT keys +// Aggregates::Max max(v) alone, the one aggregate the TopK stream supports // -// The input (`data.rs`), for each query: +// TopK chains add LIMIT 64K, above any group count, so all groups survive. // -// order unordered (shuffled, round-robin over partitions), sorted by -// the first key, or sorted by all keys, each partition sorted -// partitions 1 or 4, in batches of 64 rows -// memory unlimited, or a 4 MB pool too small for a 32K-group table -// skip-partial on or off, where a `Partial` stage on unordered input runs it -// -// Not every combination can be planned, see `shapes` and `Shape::orders`. +// `orders` arranges the source: shuffled and round-robin over the partitions, +// or sorted by the first key or by all keys with each partition sorted. +// `memory` is an unlimited pool or a 4 MB one too small for a 32K-group +// table. Combinations a key set cannot take, such as sorting struct keys, +// are skipped in `ChainTest::cases`; combinations a chain cannot run panic +// there. // // Per case, besides the rows matching `SINGLE`, the test asserts // (`assertions.rs`) that the plan was built as intended (modes, input order // modes, partition counts), that nothing fails or hangs, in particular not // with out of memory, that only spill-capable stages spill, and that the -// skip-partial probe fires exactly when it may. Per chain, at least one case -// must spill when the chain has a spill-capable stage on unordered input. +// skip-partial probe fires exactly when it may. Per test, at least one case +// must spill when the axes allow it, see `ChainTest::expects_spill`. /// The reference chain every other chain is compared against. const SINGLE: Chain = chain("single", &[Aggregate(Single)], 1); @@ -151,7 +149,17 @@ const SINGLE: Chain = chain("single", &[Aggregate(Single)], 1); /// ``` #[tokio::test(flavor = "multi_thread")] async fn single() { - assert_chain_matches_single_aggregate(SINGLE).await; + ChainTest { + chain: SINGLE, + group_by: &Keys::ALL, + aggregates: &Aggregates::HASH, + orders: &Order::ALL, + cardinalities: &Cardinality::ALL, + memory: &Memory::ALL, + skip_partial_config: &[true, false], + } + .assert_matches_single_aggregate() + .await; } /// Each partition aggregates its own keys in one pass. @@ -163,11 +171,20 @@ async fn single() { /// ``` #[tokio::test(flavor = "multi_thread")] async fn single_partitioned() { - assert_chain_matches_single_aggregate(chain( - "single_partitioned", - &[HashRepartition, Aggregate(SinglePartitioned)], - PARTITIONS, - )) + ChainTest { + chain: chain( + "single_partitioned", + &[HashRepartition, Aggregate(SinglePartitioned)], + PARTITIONS, + ), + group_by: &Keys::GROUPED, + aggregates: &Aggregates::HASH, + orders: &Order::ALL, + cardinalities: &Cardinality::ALL, + memory: &Memory::ALL, + skip_partial_config: &[true, false], + } + .assert_matches_single_aggregate() .await; } @@ -181,11 +198,20 @@ async fn single_partitioned() { /// ``` #[tokio::test(flavor = "multi_thread")] async fn single_partitioned_order_preserving() { - assert_chain_matches_single_aggregate(chain( - "single_partitioned_order_preserving", - &[OrderPreservingHashRepartition, Aggregate(SinglePartitioned)], - PARTITIONS, - )) + ChainTest { + chain: chain( + "single_partitioned_order_preserving", + &[OrderPreservingHashRepartition, Aggregate(SinglePartitioned)], + PARTITIONS, + ), + group_by: &Keys::GROUPED, + aggregates: &Aggregates::HASH, + orders: &Order::SORTED, + cardinalities: &Cardinality::ALL, + memory: &Memory::ALL, + skip_partial_config: &[true, false], + } + .assert_matches_single_aggregate() .await; } @@ -199,15 +225,24 @@ async fn single_partitioned_order_preserving() { /// ``` #[tokio::test(flavor = "multi_thread")] async fn partial_repartition_final() { - assert_chain_matches_single_aggregate(chain( - "partial_repartition_final", - &[ - Aggregate(Partial), - HashRepartition, - Aggregate(FinalPartitioned), - ], - PARTITIONS, - )) + ChainTest { + chain: chain( + "partial_repartition_final", + &[ + Aggregate(Partial), + HashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + ), + group_by: &Keys::GROUPED, + aggregates: &Aggregates::HASH, + orders: &Order::ALL, + cardinalities: &Cardinality::ALL, + memory: &Memory::ALL, + skip_partial_config: &[true, false], + } + .assert_matches_single_aggregate() .await; } @@ -221,11 +256,20 @@ async fn partial_repartition_final() { /// ``` #[tokio::test(flavor = "multi_thread")] async fn partial_coalesce_final() { - assert_chain_matches_single_aggregate(chain( - "partial_coalesce_final", - &[Aggregate(Partial), CoalescePartitions, Aggregate(Final)], - PARTITIONS, - )) + ChainTest { + chain: chain( + "partial_coalesce_final", + &[Aggregate(Partial), CoalescePartitions, Aggregate(Final)], + PARTITIONS, + ), + group_by: &Keys::ALL, + aggregates: &Aggregates::HASH, + orders: &Order::ALL, + cardinalities: &Cardinality::ALL, + memory: &Memory::ALL, + skip_partial_config: &[true, false], + } + .assert_matches_single_aggregate() .await; } @@ -240,15 +284,24 @@ async fn partial_coalesce_final() { /// ``` #[tokio::test(flavor = "multi_thread")] async fn partial_order_preserving_repartition_final() { - assert_chain_matches_single_aggregate(chain( - "partial_order_preserving_repartition_final", - &[ - Aggregate(Partial), - OrderPreservingHashRepartition, - Aggregate(FinalPartitioned), - ], - PARTITIONS, - )) + ChainTest { + chain: chain( + "partial_order_preserving_repartition_final", + &[ + Aggregate(Partial), + OrderPreservingHashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + ), + group_by: &Keys::GROUPED, + aggregates: &Aggregates::HASH, + orders: &Order::SORTED, + cardinalities: &Cardinality::ALL, + memory: &Memory::ALL, + skip_partial_config: &[true, false], + } + .assert_matches_single_aggregate() .await; } @@ -263,11 +316,20 @@ async fn partial_order_preserving_repartition_final() { /// ``` #[tokio::test(flavor = "multi_thread")] async fn partial_sort_preserving_merge_final() { - assert_chain_matches_single_aggregate(chain( - "partial_sort_preserving_merge_final", - &[Aggregate(Partial), SortPreservingMerge, Aggregate(Final)], - PARTITIONS, - )) + ChainTest { + chain: chain( + "partial_sort_preserving_merge_final", + &[Aggregate(Partial), SortPreservingMerge, Aggregate(Final)], + PARTITIONS, + ), + group_by: &Keys::GROUPED, + aggregates: &Aggregates::HASH, + orders: &Order::SORTED, + cardinalities: &Cardinality::ALL, + memory: &Memory::ALL, + skip_partial_config: &[true, false], + } + .assert_matches_single_aggregate() .await; } @@ -280,11 +342,20 @@ async fn partial_sort_preserving_merge_final() { /// ``` #[tokio::test(flavor = "multi_thread")] async fn partial_final_single_partition() { - assert_chain_matches_single_aggregate(chain( - "partial_final_single_partition", - &[Aggregate(Partial), Aggregate(Final)], - 1, - )) + ChainTest { + chain: chain( + "partial_final_single_partition", + &[Aggregate(Partial), Aggregate(Final)], + 1, + ), + group_by: &Keys::ALL, + aggregates: &Aggregates::HASH, + orders: &Order::ALL, + cardinalities: &Cardinality::ALL, + memory: &Memory::ALL, + skip_partial_config: &[true, false], + } + .assert_matches_single_aggregate() .await; } @@ -300,17 +371,26 @@ async fn partial_final_single_partition() { /// ``` #[tokio::test(flavor = "multi_thread")] async fn partial_repartition_reduce_repartition_final() { - assert_chain_matches_single_aggregate(chain( - "partial_repartition_reduce_repartition_final", - &[ - Aggregate(Partial), - HashRepartition, - Aggregate(PartialReduce), - HashRepartition, - Aggregate(FinalPartitioned), - ], - PARTITIONS, - )) + ChainTest { + chain: chain( + "partial_repartition_reduce_repartition_final", + &[ + Aggregate(Partial), + HashRepartition, + Aggregate(PartialReduce), + HashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + ), + group_by: &Keys::GROUPED, + aggregates: &Aggregates::HASH, + orders: &Order::ALL, + cardinalities: &Cardinality::ALL, + memory: &Memory::ALL, + skip_partial_config: &[true, false], + } + .assert_matches_single_aggregate() .await; } @@ -326,17 +406,26 @@ async fn partial_repartition_reduce_repartition_final() { /// ``` #[tokio::test(flavor = "multi_thread")] async fn partial_repartition_reduce_coalesce_final() { - assert_chain_matches_single_aggregate(chain( - "partial_repartition_reduce_coalesce_final", - &[ - Aggregate(Partial), - HashRepartition, - Aggregate(PartialReduce), - CoalescePartitions, - Aggregate(Final), - ], - PARTITIONS, - )) + ChainTest { + chain: chain( + "partial_repartition_reduce_coalesce_final", + &[ + Aggregate(Partial), + HashRepartition, + Aggregate(PartialReduce), + CoalescePartitions, + Aggregate(Final), + ], + PARTITIONS, + ), + group_by: &Keys::GROUPED, + aggregates: &Aggregates::HASH, + orders: &Order::ALL, + cardinalities: &Cardinality::ALL, + memory: &Memory::ALL, + skip_partial_config: &[true, false], + } + .assert_matches_single_aggregate() .await; } @@ -353,17 +442,26 @@ async fn partial_repartition_reduce_coalesce_final() { /// ``` #[tokio::test(flavor = "multi_thread")] async fn partial_coalesce_reduce_coalesce_final() { - assert_chain_matches_single_aggregate(chain( - "partial_coalesce_reduce_coalesce_final", - &[ - Aggregate(Partial), - CoalescePartitions, - Aggregate(PartialReduce), - CoalescePartitions, - Aggregate(Final), - ], - PARTITIONS, - )) + ChainTest { + chain: chain( + "partial_coalesce_reduce_coalesce_final", + &[ + Aggregate(Partial), + CoalescePartitions, + Aggregate(PartialReduce), + CoalescePartitions, + Aggregate(Final), + ], + PARTITIONS, + ), + group_by: &Keys::ALL, + aggregates: &Aggregates::HASH, + orders: &Order::ALL, + cardinalities: &Cardinality::ALL, + memory: &Memory::ALL, + skip_partial_config: &[true, false], + } + .assert_matches_single_aggregate() .await; } @@ -378,16 +476,25 @@ async fn partial_coalesce_reduce_coalesce_final() { /// ``` #[tokio::test(flavor = "multi_thread")] async fn partial_local_reduce_repartition_final() { - assert_chain_matches_single_aggregate(chain( - "partial_local_reduce_repartition_final", - &[ - Aggregate(Partial), - Aggregate(PartialReduce), - HashRepartition, - Aggregate(FinalPartitioned), - ], - PARTITIONS, - )) + ChainTest { + chain: chain( + "partial_local_reduce_repartition_final", + &[ + Aggregate(Partial), + Aggregate(PartialReduce), + HashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + ), + group_by: &Keys::GROUPED, + aggregates: &Aggregates::HASH, + orders: &Order::ALL, + cardinalities: &Cardinality::ALL, + memory: &Memory::ALL, + skip_partial_config: &[true, false], + } + .assert_matches_single_aggregate() .await; } @@ -404,17 +511,26 @@ async fn partial_local_reduce_repartition_final() { /// ``` #[tokio::test(flavor = "multi_thread")] async fn partial_reduce_final_order_preserving() { - assert_chain_matches_single_aggregate(chain( - "partial_reduce_final_order_preserving", - &[ - Aggregate(Partial), - OrderPreservingHashRepartition, - Aggregate(PartialReduce), - OrderPreservingHashRepartition, - Aggregate(FinalPartitioned), - ], - PARTITIONS, - )) + ChainTest { + chain: chain( + "partial_reduce_final_order_preserving", + &[ + Aggregate(Partial), + OrderPreservingHashRepartition, + Aggregate(PartialReduce), + OrderPreservingHashRepartition, + Aggregate(FinalPartitioned), + ], + PARTITIONS, + ), + group_by: &Keys::GROUPED, + aggregates: &Aggregates::HASH, + orders: &Order::SORTED, + cardinalities: &Cardinality::ALL, + memory: &Memory::ALL, + skip_partial_config: &[true, false], + } + .assert_matches_single_aggregate() .await; } @@ -427,8 +543,17 @@ async fn partial_reduce_final_order_preserving() { /// ``` #[tokio::test(flavor = "multi_thread")] async fn top_k_single() { - assert_chain_matches_single_aggregate(chain("top_k_single", &[TopK(Single)], 1)) - .await; + ChainTest { + chain: chain("top_k_single", &[TopK(Single)], 1), + group_by: &Keys::TOP_K, + aggregates: &Aggregates::TOP_K, + orders: &Order::ALL, + cardinalities: &Cardinality::ALL, + memory: &Memory::ALL, + skip_partial_config: &[true, false], + } + .assert_matches_single_aggregate() + .await; } /// Planner shape for `GROUP BY ... ORDER BY max(v) LIMIT n`: the limit lands on @@ -442,11 +567,20 @@ async fn top_k_single() { /// ``` #[tokio::test(flavor = "multi_thread")] async fn top_k_partial_repartition_final() { - assert_chain_matches_single_aggregate(chain( - "top_k_partial_repartition_final", - &[Aggregate(Partial), HashRepartition, TopK(FinalPartitioned)], - PARTITIONS, - )) + ChainTest { + chain: chain( + "top_k_partial_repartition_final", + &[Aggregate(Partial), HashRepartition, TopK(FinalPartitioned)], + PARTITIONS, + ), + group_by: &Keys::TOP_K, + aggregates: &Aggregates::TOP_K, + orders: &Order::ALL, + cardinalities: &Cardinality::ALL, + memory: &Memory::ALL, + skip_partial_config: &[true, false], + } + .assert_matches_single_aggregate() .await; } @@ -460,11 +594,20 @@ async fn top_k_partial_repartition_final() { /// ``` #[tokio::test(flavor = "multi_thread")] async fn top_k_partial_coalesce_final() { - assert_chain_matches_single_aggregate(chain( - "top_k_partial_coalesce_final", - &[Aggregate(Partial), CoalescePartitions, TopK(Final)], - PARTITIONS, - )) + ChainTest { + chain: chain( + "top_k_partial_coalesce_final", + &[Aggregate(Partial), CoalescePartitions, TopK(Final)], + PARTITIONS, + ), + group_by: &Keys::TOP_K, + aggregates: &Aggregates::TOP_K, + orders: &Order::ALL, + cardinalities: &Cardinality::ALL, + memory: &Memory::ALL, + skip_partial_config: &[true, false], + } + .assert_matches_single_aggregate() .await; } @@ -478,11 +621,20 @@ async fn top_k_partial_coalesce_final() { /// ``` #[tokio::test(flavor = "multi_thread")] async fn top_k_both_stages() { - assert_chain_matches_single_aggregate(chain( - "top_k_both_stages", - &[TopK(Partial), HashRepartition, TopK(FinalPartitioned)], - PARTITIONS, - )) + ChainTest { + chain: chain( + "top_k_both_stages", + &[TopK(Partial), HashRepartition, TopK(FinalPartitioned)], + PARTITIONS, + ), + group_by: &Keys::TOP_K, + aggregates: &Aggregates::TOP_K, + orders: &Order::ALL, + cardinalities: &Cardinality::ALL, + memory: &Memory::ALL, + skip_partial_config: &[true, false], + } + .assert_matches_single_aggregate() .await; } @@ -587,51 +739,69 @@ fn reference_case(query: Query, cardinality: Cardinality) -> Case { } } -/// Runs every case of `chain` and asserts each returns the rows of the -/// `SINGLE` chain for its query, see the test preamble for the full list of -/// checks. A failure does not stop the run, so one run reports every failing -/// case. -async fn assert_chain_matches_single_aggregate(chain: Chain) { - const SEED: u64 = 42; - let mut total_spilled = 0; - let mut failures: Vec = vec![]; - // Every in-flight case holds several copies of the dataset and its own - // partitioned streams, so bound the concurrency by the cores at hand - // instead of spawning the whole matrix. - let max_concurrent_cases = get_available_parallelism(); - - for cardinality in Cardinality::ALL { - let rows = generate_rows(cardinality, SEED); - let cases: Vec = cases(chain) - .into_iter() - .filter(|case| case.params.cardinality == cardinality) - .collect(); - let mut reference_cases: Vec = vec![]; - for case in &cases { - let query = case.shape.query; - if !reference_cases.iter().any(|case| case.shape.query == query) { - reference_cases.push(reference_case(query, cardinality)); +impl ChainTest { + /// Runs every case and asserts each returns the rows of the `SINGLE` + /// chain for its query, see the preamble for the full list of checks. A + /// failure does not stop the run, so one run reports every failing case. + async fn assert_matches_single_aggregate(self) { + const SEED: u64 = 42; + let chain = self.chain; + let mut total_spilled = 0; + let mut failures: Vec = vec![]; + // Every in-flight case holds several copies of the dataset and its own + // partitioned streams, so bound the concurrency by the cores at hand + // instead of spawning the whole matrix. + let max_concurrent_cases = get_available_parallelism(); + + for cardinality in Cardinality::ALL { + let rows = generate_rows(cardinality, SEED); + let cases: Vec = self + .cases() + .into_iter() + .filter(|case| case.params.cardinality == cardinality) + .collect(); + let mut reference_cases: Vec = vec![]; + for case in &cases { + let query = case.shape.query; + if !reference_cases.iter().any(|case| case.shape.query == query) { + reference_cases.push(reference_case(query, cardinality)); + } + } + let inputs = + Arc::new(arrange_all(&rows, cases.iter().chain(&reference_cases))); + + let mut expected_by_query: Vec<(Query, String)> = Vec::new(); + for case in reference_cases { + let query = case.shape.query; + let outcome = run_case(case, Arc::clone(&inputs)).await; + expected_by_query.push((query, outcome.output)); } - } - let inputs = Arc::new(arrange_all(&rows, cases.iter().chain(&reference_cases))); - - let mut expected_by_query: Vec<(Query, String)> = Vec::new(); - for case in reference_cases { - let query = case.shape.query; - let outcome = run_case(case, Arc::clone(&inputs)).await; - expected_by_query.push((query, outcome.output)); - } - let mut join_set = JoinSet::new(); - let (mut spilled, mut finished) = (vec![], vec![]); - for case in cases { - let inputs = Arc::clone(&inputs); - let expected = expected_by_query - .iter() - .find(|(query, _)| *query == case.shape.query) - .map(|(_, expected)| expected.clone()) - .unwrap(); - while join_set.len() >= max_concurrent_cases { + let mut join_set = JoinSet::new(); + let (mut spilled, mut finished) = (vec![], vec![]); + for case in cases { + let inputs = Arc::clone(&inputs); + let expected = expected_by_query + .iter() + .find(|(query, _)| *query == case.shape.query) + .map(|(_, expected)| expected.clone()) + .unwrap(); + while join_set.len() >= max_concurrent_cases { + collect_finished( + &mut join_set, + &mut spilled, + &mut finished, + &mut failures, + ) + .await; + } + join_set.spawn(async move { + let outcome = run_case(case.clone(), inputs).await; + assert_eq!(outcome.output, expected, "{case:?}"); + (case, outcome.spilled) + }); + } + while !join_set.is_empty() { collect_finished( &mut join_set, &mut spilled, @@ -640,35 +810,25 @@ async fn assert_chain_matches_single_aggregate(chain: Chain) { ) .await; } - join_set.spawn(async move { - let outcome = run_case(case.clone(), inputs).await; - assert_eq!(outcome.output, expected, "{case:?}"); - (case, outcome.spilled) - }); + print_cases(cardinality, "spilled", &spilled); + print_cases(cardinality, "finished without spilling", &finished); + total_spilled += spilled.len(); } - while !join_set.is_empty() { - collect_finished(&mut join_set, &mut spilled, &mut finished, &mut failures) - .await; + if self.expects_spill() { + assert!( + total_spilled > 0, + "{}: no case exercised the spill path", + chain.name + ); } - print_cases(cardinality, "spilled", &spilled); - print_cases(cardinality, "finished without spilling", &finished); - total_spilled += spilled.len(); - } - // A shape filter may select only shapes that cannot spill - if chain.expects_spill() && std::env::var("AGGREGATE_CHAIN_SHAPES").is_err() { assert!( - total_spilled > 0, - "{}: no case exercised the spill path", - chain.name + failures.is_empty(), + "{}: {} cases failed:\n\n{}", + chain.name, + failures.len(), + failures.join("\n\n") ); } - assert!( - failures.is_empty(), - "{}: {} cases failed:\n\n{}", - chain.name, - failures.len(), - failures.join("\n\n") - ); } /// A case takes about two seconds alone in a debug build, but CI runs the diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/case_space.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/case_space.rs index 5b52814462662..571b178574672 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/case_space.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/case_space.rs @@ -41,6 +41,16 @@ pub(super) enum Order { SortedByAllKeys, } +impl Order { + pub(super) const ALL: [Self; 3] = [ + Self::Unordered, + Self::SortedByFirstKey, + Self::SortedByAllKeys, + ]; + /// For chains that preserve ordering, which need an ordering to preserve. + pub(super) const SORTED: [Self; 2] = [Self::SortedByFirstKey, Self::SortedByAllKeys]; +} + #[derive(Clone, Copy, Debug, PartialEq)] pub(super) enum Cardinality { VeryHigh, @@ -72,6 +82,10 @@ pub(super) enum Memory { Limited, } +impl Memory { + pub(super) const ALL: [Self; 2] = [Self::Unlimited, Self::Limited]; +} + /// One operator in a chain, listed bottom to top. #[derive(Clone, Copy, Debug)] pub(super) enum Operator { @@ -125,6 +139,19 @@ impl Keys { Self::Mixed, Self::Struct, ]; + /// Every `GROUP BY`, for chains that hash or sort on the keys. + pub(super) const GROUPED: [Self; 7] = [ + Self::TwoInts, + Self::Boolean, + Self::Bytes, + Self::BytesView, + Self::Primitive, + Self::Mixed, + Self::Struct, + ]; + /// The keys `GroupedTopKAggregateStream` supports: one primitive or + /// string column. + pub(super) const TOP_K: [Self; 3] = [Self::Bytes, Self::BytesView, Self::Primitive]; /// Key columns, in `GROUP BY` order. pub(super) fn columns(self) -> &'static [&'static str] { @@ -176,7 +203,10 @@ pub(super) enum Aggregates { } impl Aggregates { - pub(super) const ALL: [Self; 3] = [Self::All, Self::None, Self::Max]; + /// For hash chains: `max` alone is a subset of `All` and adds nothing. + pub(super) const HASH: [Self; 2] = [Self::All, Self::None]; + /// For TopK chains, which support a single `max` or no aggregates. + pub(super) const TOP_K: [Self; 2] = [Self::Max, Self::None]; } /// The logical query a chain computes. @@ -284,23 +314,27 @@ impl Shape { ) } - /// Source orders that make sense for this shape. Order-preserving shuffles - /// need an ordering to preserve; no-grouping chains ignore ordering. - pub(super) fn orders(&self) -> Vec { - let keys = self.query.keys.columns(); - let mut orders = vec![]; - if !self.chain.preserves_order() { - orders.push(Order::Unordered); - } - if self.query.keys.sortable() && !keys.is_empty() { - // With a single key, sorting by the first key is already sorting - // by all keys. - if keys.len() > 1 { - orders.push(Order::SortedByFirstKey); - } - orders.push(Order::SortedByAllKeys); - } - orders + /// The subset of `requested` the source can be arranged in for these keys. + /// Struct keys cannot be sorted, and with a single key sorting by the + /// first key is already sorting by all keys. + pub(super) fn orders(&self, requested: &[Order]) -> Vec { + let keys = self.query.keys; + requested + .iter() + .copied() + .filter(|order| match order { + Order::Unordered => { + assert!( + !self.chain.preserves_order(), + "{}: an order-preserving chain needs sorted input", + self.chain.name + ); + true + } + Order::SortedByFirstKey => keys.sortable() && keys.columns().len() > 1, + Order::SortedByAllKeys => keys.sortable() && !keys.columns().is_empty(), + }) + .collect() } /// Whether some `Partial` stage of this shape runs the skip-partial probe @@ -321,34 +355,19 @@ impl Shape { } } -/// Every query the chain can be planned for. -/// -/// - Without keys a chain can neither hash nor sort, and `max` alone is a -/// subset of the full aggregate list, so only that list runs. -/// - The TopK stream needs one primitive or string key with either a single -/// `max` or no aggregates at all (the `DISTINCT ... LIMIT` form). -/// - Every other chain runs the full aggregate list and the accumulator-free -/// form; `max` alone adds nothing there. -pub(super) fn shapes(chain: Chain) -> Vec { - let mut shapes = vec![]; - for keys in Keys::ALL { - for aggregates in Aggregates::ALL { - let valid = if keys == Keys::None { - !chain.needs_keys() && aggregates == Aggregates::All - } else if chain.is_top_k() { - keys.top_k_supported() && aggregates != Aggregates::All - } else { - aggregates != Aggregates::Max - }; - if valid { - shapes.push(Shape { - chain, - query: Query { keys, aggregates }, - }); - } - } - } - shapes +/// One test: a chain and the axes it runs over. Every field is a list so a +/// test reads as the cases it covers, and narrowing a list runs just those. +pub(super) struct ChainTest { + pub(super) chain: Chain, + pub(super) group_by: &'static [Keys], + pub(super) aggregates: &'static [Aggregates], + pub(super) orders: &'static [Order], + pub(super) cardinalities: &'static [Cardinality], + pub(super) memory: &'static [Memory], + /// Whether the skip-partial probe may fire. Only a grouped `Partial` + /// stage on unordered input runs it; elsewhere the setting changes + /// nothing and only the first value runs. + pub(super) skip_partial_config: &'static [bool], } #[derive(Clone, Debug)] @@ -357,44 +376,81 @@ pub(super) struct Case { pub(super) params: CaseParams, } -/// Every case of the chain: each valid query over every source order, -/// cardinality, memory budget and skip-partial setting. -pub(super) fn cases(chain: Chain) -> Vec { - // `AGGREGATE_CHAIN_SHAPES=a,b` restricts the run to shapes whose name - // contains one of the given substrings, to reproduce or bisect quickly. - let shape_filter: Vec = std::env::var("AGGREGATE_CHAIN_SHAPES") - .map(|value| value.split(',').map(str::to_string).collect()) - .unwrap_or_default(); - let mut cases = vec![]; - for shape in shapes(chain).into_iter().filter(|shape| { - shape_filter.is_empty() - || shape_filter - .iter() - .any(|needle| shape.name().contains(needle)) - }) { - for order in shape.orders() { - let skip_partial_variants: &[bool] = - if shape.has_skip_partial_candidate(order) { - &[true, false] +impl ChainTest { + /// Every query of the test: each key set with each aggregate list. + /// Without keys only the full aggregate list runs: `max` alone is a + /// subset of it and no aggregates at all is not a query. + fn shapes(&self) -> Vec { + let chain = self.chain; + let mut shapes = vec![]; + for &keys in self.group_by { + assert!( + keys != Keys::None || !chain.needs_keys(), + "{}: the chain hashes or sorts on group keys, so it needs some", + chain.name + ); + assert!( + !chain.is_top_k() || keys.top_k_supported(), + "{}: TopK needs one primitive or string key, not {keys:?}", + chain.name + ); + for &aggregates in self.aggregates { + assert!( + !chain.is_top_k() || aggregates != Aggregates::All, + "{}: TopK supports a single max or no aggregates, not {aggregates:?}", + chain.name + ); + if keys == Keys::None && aggregates != Aggregates::All { + continue; + } + shapes.push(Shape { + chain, + query: Query { keys, aggregates }, + }); + } + } + shapes + } + + /// Every case of the test: each query over each source order, + /// cardinality, memory budget and skip-partial setting. + pub(super) fn cases(&self) -> Vec { + let mut cases = vec![]; + for shape in self.shapes() { + for order in shape.orders(self.orders) { + let skip_partial = if shape.has_skip_partial_candidate(order) { + self.skip_partial_config } else { - &[true] + &self.skip_partial_config[..1] }; - for cardinality in Cardinality::ALL { - for memory in [Memory::Unlimited, Memory::Limited] { - for &skip_partial_enabled in skip_partial_variants { - cases.push(Case { - shape, - params: CaseParams { - order, - cardinality, - memory, - skip_partial_enabled, - }, - }); + for &cardinality in self.cardinalities { + for &memory in self.memory { + for &skip_partial_enabled in skip_partial { + cases.push(Case { + shape, + params: CaseParams { + order, + cardinality, + memory, + skip_partial_enabled, + }, + }); + } } } } } + cases + } + + /// Whether some case must spill: the chain has a spill-capable stage on + /// unordered input and the axes include the very-high-cardinality table + /// under the limited pool that cannot fit. + pub(super) fn expects_spill(&self) -> bool { + self.chain.expects_spill() + && self.orders.contains(&Order::Unordered) + && self.cardinalities.contains(&Cardinality::VeryHigh) + && self.memory.contains(&Memory::Limited) + && self.group_by.iter().any(|keys| keys.tracks_cardinality()) } - cases } From 356adb00a5cc62d8dc6644b0f2e6d3865d3d63c3 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:02:59 +0300 Subject: [PATCH 20/21] cleanup --- .../tests/fuzz_cases/aggregate_chain_fuzz.rs | 76 +++++-------------- 1 file changed, 18 insertions(+), 58 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs index ebf50154345de..32f9ee24fcf8f 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -19,11 +19,12 @@ //! and assert identical results. //! //! One test per operator [`Chain`], each running the chain over the cross -//! product of the other axes: the group [`Keys`] (one per `GroupValues` -//! implementation), the [`Aggregates`], the source [`Order`], the group -//! [`Cardinality`], the [`Memory`] budget and whether the skip-partial probe -//! may fire. Only the combinations that cannot be planned are left out, see -//! [`shapes`]. +//! product of the other axes: +//! - The group [`Keys`] (single column with specific types, multiple columns, or fallback implementation) +//! - The [`Aggregates`] can be no aggregates or multiple aggregates that both need to track uniqueness and not (so more memory will be used), +//! - The source whether it is ordered by all keys, subset, or none [`Order`]. +//! - The group [`Cardinality`] too test different memory and spill behavior +//! - The [`Memory`] budget - whether it is limited or not use std::collections::HashMap; use std::num::NonZeroUsize; @@ -81,61 +82,20 @@ use context::*; use data::*; use plan::*; -// What is tested -// ============== -// // Every test below is one physical plan shape, an `AggregateExec` chain, run -// over every combination of the axes listed in its `ChainTest`. Each run must -// return the same rows as the plain single-stage aggregate `SINGLE`. To test -// one case, narrow the lists in place. `chain` lists the operators bottom-up, -// source first; the doc comment shows the same plan as DataFusion prints it. -// -// The table (`data.rs`), 32K rows, generated from a fixed seed: -// -// k1 Int64, k2 Int64 two-column key, k1 alone has fewer distinct -// values than (k1, k2) -// v Int64 the aggregated value, -1000..1000 -// b Boolean one column per `GroupValues` implementation, -// s Utf8 each with one distinct value per (k1, k2) -// sv Utf8View group (Boolean: two) -// p Int64 -// st Struct, num: Int64> -// -// Every key column has about 3% nulls. `cardinalities` is the number of -// groups: 32K (one row per group), 1K, 16 or 2. -// -// `group_by` picks the query, one per `GroupValues` implementation: -// -// Keys::None SELECT FROM t -// Keys::TwoInts SELECT k1, k2, FROM t GROUP BY k1, k2 -- Column -// Keys::Boolean SELECT b, FROM t GROUP BY b -- Boolean -// Keys::Bytes SELECT s, FROM t GROUP BY s -- Bytes -// Keys::BytesView SELECT sv, FROM t GROUP BY sv -- BytesView -// Keys::Primitive SELECT p, FROM t GROUP BY p -- Primitive -// Keys::Mixed SELECT b, s, sv, p, FROM t GROUP BY b, s, sv, p -// Keys::Struct SELECT st, FROM t GROUP BY st -- row fallback -// -// `aggregates` picks : -// -// Aggregates::All count(v), count(DISTINCT v), sum(v), avg(v), min(v), max(v) -// Aggregates::None nothing, as in SELECT DISTINCT keys -// Aggregates::Max max(v) alone, the one aggregate the TopK stream supports -// -// TopK chains add LIMIT 64K, above any group count, so all groups survive. -// -// `orders` arranges the source: shuffled and round-robin over the partitions, -// or sorted by the first key or by all keys with each partition sorted. -// `memory` is an unlimited pool or a 4 MB one too small for a 32K-group -// table. Combinations a key set cannot take, such as sorting struct keys, -// are skipped in `ChainTest::cases`; combinations a chain cannot run panic -// there. +// over the same table for every combination of the axes in its `ChainTest`: +// which `GROUP BY` (one per `GroupValues` implementation), which aggregates, +// how the source is ordered and how many groups it has, how much memory the +// pool allows, and whether the skip-partial probe may fire. Every run must +// return the same rows as the plain single-stage aggregate `SINGLE`, and +// `assertions.rs` checks along the way that the plan was built as intended, +// that nothing runs out of memory or hangs, and that spilling and the +// skip-partial probe happen exactly where they may. // -// Per case, besides the rows matching `SINGLE`, the test asserts -// (`assertions.rs`) that the plan was built as intended (modes, input order -// modes, partition counts), that nothing fails or hangs, in particular not -// with out of memory, that only spill-capable stages spill, and that the -// skip-partial probe fires exactly when it may. Per test, at least one case -// must spill when the axes allow it, see `ChainTest::expects_spill`. +// `chain` lists the operators bottom-up, source first; the doc comment shows +// the same plan as DataFusion prints it. To run one case, narrow the lists +// in place. The axes and their values are documented in `case_space.rs`, the +// table in `data.rs`. /// The reference chain every other chain is compared against. const SINGLE: Chain = chain("single", &[Aggregate(Single)], 1); From de3939a3af8ebfbbcc32e186d8b6db48c830240e Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:21:44 +0300 Subject: [PATCH 21/21] fix ci --- datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs | 4 +++- .../core/tests/fuzz_cases/aggregate_chain_fuzz/assertions.rs | 1 + .../core/tests/fuzz_cases/aggregate_chain_fuzz/case_space.rs | 1 + .../core/tests/fuzz_cases/aggregate_chain_fuzz/context.rs | 1 + datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/data.rs | 1 + datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/plan.rs | 1 + 6 files changed, 8 insertions(+), 1 deletion(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs index 32f9ee24fcf8f..a46204d56a46e 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -100,6 +100,9 @@ use plan::*; /// The reference chain every other chain is compared against. const SINGLE: Chain = chain("single", &[Aggregate(Single)], 1); +/// Fixed so every test sees the same table. +const SEED: u64 = 42; + /// `Single` on one partition. Also the reference every other chain is compared /// against. /// @@ -704,7 +707,6 @@ impl ChainTest { /// chain for its query, see the preamble for the full list of checks. A /// failure does not stop the run, so one run reports every failing case. async fn assert_matches_single_aggregate(self) { - const SEED: u64 = 42; let chain = self.chain; let mut total_spilled = 0; let mut failures: Vec = vec![]; diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/assertions.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/assertions.rs index f65e58ef3feb7..8b026a8c1230c 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/assertions.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/assertions.rs @@ -14,6 +14,7 @@ // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. + //! Assertions on plan shape and metrics. use super::*; diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/case_space.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/case_space.rs index 571b178574672..6b12cac0e3c26 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/case_space.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/case_space.rs @@ -14,6 +14,7 @@ // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. + //! The case space: every axis a case varies along and the valid combinations. use super::*; diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/context.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/context.rs index faf4bc2c94c81..a918c3186c2ff 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/context.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/context.rs @@ -14,6 +14,7 @@ // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. + //! Execution context: session config and memory pool per case. use super::*; diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/data.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/data.rs index 60f687ff06f63..c897319f1a170 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/data.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/data.rs @@ -14,6 +14,7 @@ // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. + //! Data generation: deterministic rows per seed, arranged per order and partition count. use super::*; diff --git a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/plan.rs b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/plan.rs index dc12785793078..3dc9a21531a31 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/plan.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/plan.rs @@ -14,6 +14,7 @@ // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. + //! Plan construction for a [`Shape`] over a source. use super::*;