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..a46204d56a46e --- /dev/null +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz.rs @@ -0,0 +1,852 @@ +// 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 tests that run every valid `AggregateExec` chain over the same data +//! and assert identical results. +//! +//! One test per operator [`Chain`], each running the chain over the cross +//! 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; +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::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}; +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::*; + +mod assertions; +mod case_space; +mod context; +mod data; +mod plan; + +use assertions::*; +use case_space::*; +use context::*; +use data::*; +use plan::*; + +// Every test below is one physical plan shape, an `AggregateExec` chain, run +// 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. +// +// `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); + +/// 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. +/// +/// ```text +/// AggregateExec: mode=Single +/// DataSourceExec: partitions=1 +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn single() { + 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. +/// +/// ```text +/// AggregateExec: mode=SinglePartitioned +/// RepartitionExec: partitioning=Hash(keys) +/// DataSourceExec: partitions=PARTITIONS +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn single_partitioned() { + 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; +} + +/// 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() { + 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; +} + +/// 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() { + 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; +} + +/// 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() { + 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; +} + +/// 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() { + 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; +} + +/// 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() { + 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; +} + +/// 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() { + 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; +} + +/// 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() { + 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; +} + +/// 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() { + 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; +} + +/// 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() { + 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; +} + +/// `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() { + 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; +} + +/// 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() { + 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; +} + +/// `GroupedTopKAggregateStream` alone. The limit is above the group count, so +/// every group survives. +/// +/// ```text +/// AggregateExec: mode=Single, lim=[TOP_K_LIMIT] +/// DataSourceExec: partitions=1 +/// ``` +#[tokio::test(flavor = "multi_thread")] +async fn top_k_single() { + 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 +/// the final stage. +/// +/// ```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() { + 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; +} + +/// TopK final stage on one coalesced partition. +/// +/// ```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() { + 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; +} + +/// TopK on both stages. +/// +/// ```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() { + 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; +} + +// --------------------------------------------------------------------------- +// Driver +// --------------------------------------------------------------------------- + +/// Sorted output plus the stages that spilled, empty if none did. +struct Outcome { + output: String, + 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, inputs: Arc) -> Outcome { + log::debug!("start {case:?}"); + let outcome = run_case_inner(&case, &inputs[&input_key(&case)]).await; + log::debug!("done {case:?}"); + outcome +} + +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), + ); + 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, + } +} + +/// 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: SINGLE, + query, + }, + params: CaseParams { + order: Order::Unordered, + cardinality, + memory: Memory::Unlimited, + skip_partial_enabled: true, + }, + } +} + +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) { + 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 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, + &mut finished, + &mut failures, + ) + .await; + } + print_cases(cardinality, "spilled", &spilled); + print_cases(cardinality, "finished without spilling", &finished); + total_spilled += spilled.len(); + } + if self.expects_spill() { + assert!( + total_spilled > 0, + "{}: no case exercised the spill path", + chain.name + ); + } + 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 +/// 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. +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()), + } +} + +/// 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() + .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.params.order) + { + format!(" skip_partial={:<5}", case.params.skip_partial_enabled) + } else { + " ".repeat(19) + }; + format!( + " {:<45} {:<17} memory={:<9}{skip_partial}{spilled}", + case.shape.name(), + format!("{:?}", case.params.order), + format!("{:?}", case.params.memory), + ) + }) + .collect(); + lines.sort(); + // Enable with `RUST_LOG=debug` + log::debug!("{cardinality:?}: {} cases {outcome}", lines.len()); + for line in lines { + log::debug!("{line}"); + } +} 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..8b026a8c1230c --- /dev/null +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/assertions.rs @@ -0,0 +1,182 @@ +// 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..6b12cac0e3c26 --- /dev/null +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/case_space.rs @@ -0,0 +1,457 @@ +// 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, +} + +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, + 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, +} + +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 { + 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, + ]; + /// 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] { + 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 { + /// 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. +#[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 + ) + } + + /// 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 + /// 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 + } +} + +/// 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)] +pub(super) struct Case { + pub(super) shape: Shape, + pub(super) params: CaseParams, +} + +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 { + &self.skip_partial_config[..1] + }; + 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()) + } +} 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..a918c3186c2ff --- /dev/null +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/context.rs @@ -0,0 +1,66 @@ +// 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..c897319f1a170 --- /dev/null +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/data.rs @@ -0,0 +1,216 @@ +// 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..3dc9a21531a31 --- /dev/null +++ b/datafusion/core/tests/fuzz_cases/aggregate_chain_fuzz/plan.rs @@ -0,0 +1,178 @@ +// 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 3e425e48c7a0c..92b68361d007b 100644 --- a/datafusion/core/tests/fuzz_cases/mod.rs +++ b/datafusion/core/tests/fuzz_cases/mod.rs @@ -27,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; 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 {