test: catch operators that hold input-sized memory without reserving it

.github/workflows/extended.yml

@@ -84,6 +84,8 @@ jobs:
           df -h
       - name: Run tests (excluding doctests)
         run: cargo xtask ci step test extended
+      - name: Run spill pool fuzzer with more iterations
+        run: cargo xtask ci step test spill-pool-fuzz
       - parallel:
           - name: Verify Working Directory Clean
             run: git diff --exit-code

datafusion/core/Cargo.toml

@@ -198,6 +198,16 @@ ignored = ["datafusion-doc", "datafusion-macros", "dashmap"]
 [target.'cfg(not(target_os = "windows"))'.dev-dependencies]
 nix = { version = "0.31.3", features = ["fs"] }
 
+[[test]]
+name = "fuzz"
+path = "tests/fuzz.rs"
+required-features = ["extended_tests"]
+
+[[test]]
+name = "memory_limit_validation"
+path = "tests/memory_limit_validation/main.rs"
+required-features = ["extended_tests"]
+
 [[bin]]
 name = "print_config_docs"
 path = "src/bin/print_config_docs.rs"

datafusion/core/tests/fuzz.rs

@@ -16,9 +16,8 @@
 // under the License.
 
 /// Run all tests that are found in the `fuzz_cases` directory.
-/// Fuzz tests are slow and gated behind the `extended_tests` feature.
-/// Run with: cargo test --features extended_tests
-#[cfg(feature = "extended_tests")]
+/// Fuzz tests are slow, so this target requires the `extended_tests` feature.
+/// Run with: cargo test -p datafusion --test fuzz --features extended_tests
 mod fuzz_cases;
 
 #[cfg(test)]

datafusion/core/tests/fuzz_cases/leaf_pushdown_fuzz.rs

@@ -76,7 +76,7 @@ const INCLUDE_FALSE_BRANCH_UNION_SHAPES: bool = false;
 /// Cases of the short run. Keep this low: the whole test must stay in a few
 /// seconds.
 const DEFAULT_CASES: usize = 250;
-/// Cases of the long run, gated behind the `extended_tests` feature.
+/// Cases of the long run.
 const EXTENDED_CASES: usize = 5000;
 /// Highest share of cases that may fail to plan on both sides. A higher share
 /// means the generator makes junk.
@@ -777,9 +777,6 @@ async fn leaf_pushdown_parquet_fuzz() {
     run_parquet(env_usize("LEAF_PUSHDOWN_FUZZ_CASES", DEFAULT_CASES / 2)).await;
 }
 
-/// The whole `fuzz_cases` tree is behind the `extended_tests` feature today.
-/// The attribute keeps the long run gated if that changes.
-#[cfg(feature = "extended_tests")]
 #[tokio::test(flavor = "multi_thread")]
 async fn leaf_pushdown_fuzz_extended() {
     run_memory(env_usize("LEAF_PUSHDOWN_FUZZ_CASES", EXTENDED_CASES)).await;

datafusion/core/tests/memory_limit/mod.rs

@@ -20,8 +20,6 @@
 use std::num::NonZeroUsize;
 use std::sync::{Arc, LazyLock};
 
-#[cfg(feature = "extended_tests")]
-mod memory_limit_validation;
 mod nlj_spill_unmatched;
 mod repartition_mem_limit;
 mod union_nullable_spill;

datafusion/core/tests/memory_limit_validation/aggregate_mem_validation.rs

@@ -0,0 +1,117 @@
+// 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.
+
+//! Memory validation for aggregates.
+//!
+//! These tests run in separate processes to accurately measure memory usage.
+
+use datafusion::prelude::SessionConfig;
+
+use crate::utils;
+
+#[test]
+fn group_by_few_groups_runner() {
+    utils::spawn_test_process("aggregate_mem_validation", "group_by_few_groups");
+}
+
+#[test]
+fn distinct_runner() {
+    utils::spawn_test_process("aggregate_mem_validation", "distinct");
+}
+
+#[test]
+fn median_runner() {
+    utils::spawn_test_process("aggregate_mem_validation", "median");
+}
+
+#[test]
+fn array_agg_few_groups_runner() {
+    utils::spawn_test_process("aggregate_mem_validation", "array_agg_few_groups");
+}
+
+#[test]
+fn array_agg_few_groups_one_partition_runner() {
+    utils::spawn_test_process(
+        "aggregate_mem_validation",
+        "array_agg_few_groups_one_partition",
+    );
+}
+
+/// Few groups: the aggregate holds almost nothing, whatever the input size.
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn group_by_few_groups() {
+    utils::validate_query_over_rows_with_strings(
+        20_000_000,
+        SessionConfig::new().with_target_partitions(4),
+        "AggregateExec",
+        |t| format!("SELECT v % 10 AS k, count(*), max(s) FROM {t} GROUP BY k"),
+    )
+    .await;
+}
+
+/// Every row is its own group.
+#[ignore = "hash table grows before it is reserved, see https://github.com/apache/datafusion/issues/25606"]
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn distinct() {
+    utils::validate_query_over_rows_with_strings(
+        1_800_000_000,
+        SessionConfig::new().with_target_partitions(4),
+        "AggregateExec",
+        |t| format!("SELECT DISTINCT s FROM {t}"),
+    )
+    .await;
+}
+
+/// The state of `median` holds every input value.
+#[ignore = "partial state reaches the final aggregate unreserved, see https://github.com/apache/datafusion/issues/26705"]
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn median() {
+    utils::validate_query_over_rows_with_strings(
+        80_000_000 * 4,
+        SessionConfig::new().with_target_partitions(4),
+        "AggregateExec",
+        |t| format!("SELECT median(v) FROM {t}"),
+    )
+    .await;
+}
+
+/// Few groups, each with a state that holds every input value of the group.
+#[ignore = "partial state reaches the final aggregate unreserved, see https://github.com/apache/datafusion/issues/26705"]
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn array_agg_few_groups() {
+    utils::validate_query_over_rows_with_strings(
+        2_500_000_000,
+        SessionConfig::new().with_target_partitions(4),
+        "AggregateExec",
+        |t| format!("SELECT v % 10 AS k, array_agg(s) FROM {t} GROUP BY k"),
+    )
+    .await;
+}
+
+/// Like `array_agg_few_groups`, in one partition, so there is a single
+/// aggregate whose output is as large as its input.
+#[ignore = "aggregate output is materialized unreserved, see https://github.com/apache/datafusion/issues/26701"]
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn array_agg_few_groups_one_partition() {
+    utils::validate_query_over_rows_with_strings(
+        1_200_000_000,
+        SessionConfig::new().with_target_partitions(1),
+        "AggregateExec",
+        |t| format!("SELECT v % 10 AS k, array_agg(s) FROM {t} GROUP BY k"),
+    )
+    .await;
+}

datafusion/core/tests/memory_limit_validation/hash_join_mem_validation.rs

@@ -0,0 +1,60 @@
+// 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.
+
+//! Memory validation for hash joins.
+//!
+//! These tests run in separate processes to accurately measure memory usage.
+
+use datafusion::prelude::SessionConfig;
+
+use crate::utils;
+
+#[test]
+fn hash_join_partitioned_runner() {
+    utils::spawn_test_process("hash_join_mem_validation", "hash_join_partitioned");
+}
+
+#[test]
+fn hash_join_collect_left_runner() {
+    utils::spawn_test_process("hash_join_mem_validation", "hash_join_collect_left");
+}
+
+/// Both inputs are repartitioned on the join key.
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn hash_join_partitioned() {
+    utils::validate_query_over_rows_with_strings(
+        80_000_000 * 5,
+        SessionConfig::new().with_target_partitions(4),
+        "HashJoinExec: mode=Partitioned",
+        |t| format!("SELECT count(*), max(b.s) FROM {t} a JOIN {t} b ON a.v = b.v"),
+    )
+    .await;
+}
+
+/// Every partition probes one build side.
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn hash_join_collect_left() {
+    utils::validate_query_over_rows_with_strings(
+        80_000_000 * 4,
+        SessionConfig::new()
+            .with_target_partitions(4)
+            .set_bool("datafusion.optimizer.repartition_joins", false),
+        "HashJoinExec: mode=CollectLeft",
+        |t| format!("SELECT count(*), max(b.s) FROM {t} a JOIN {t} b ON a.v = b.v"),
+    )
+    .await;
+}

datafusion/core/tests/memory_limit/memory_limit_validation/mod.rs → datafusion/core/tests/memory_limit_validation/main.rs

@@ -16,9 +16,25 @@
 // under the License.
 
 //! Validates query's actual memory usage is consistent with the specified memory
-//! limit.
+//! limit. This target requires the `extended_tests` feature.
+//!
+//! Run with: cargo test -p datafusion --test memory_limit_validation --features extended_tests
 
+mod aggregate_mem_validation;
+mod hash_join_mem_validation;
 mod nlj_mem_validation;
 mod smj_mem_validation;
 mod sort_mem_validation;
 mod utils;
+mod window_mem_validation;
+
+#[global_allocator]
+static ALLOC: test_utils::counting_allocator::CountingAllocator =
+    test_utils::counting_allocator::CountingAllocator;
+
+#[cfg(test)]
+#[ctor::ctor(unsafe)]
+fn init() {
+    // Enable RUST_LOG logging configuration for test
+    let _ = env_logger::try_init();
+}

datafusion/core/tests/memory_limit/memory_limit_validation/nlj_mem_validation.rs → datafusion/core/tests/memory_limit_validation/nlj_mem_validation.rs

@@ -19,7 +19,7 @@
 
 use datafusion::prelude::SessionConfig;
 
-use crate::memory_limit::memory_limit_validation::utils;
+use crate::utils;
 
 #[test]
 fn nlj_no_mem_limit_runner() {

datafusion/core/tests/memory_limit/memory_limit_validation/smj_mem_validation.rs → datafusion/core/tests/memory_limit_validation/smj_mem_validation.rs

@@ -21,7 +21,7 @@
 
 use datafusion::prelude::SessionConfig;
 
-use crate::memory_limit::memory_limit_validation::utils;
+use crate::utils;
 
 /// Ensures the planner selected a sort-merge join.
 const SMJ_OPERATOR_NAME: &str = "SortMergeJoinExec";

datafusion/core/tests/memory_limit/memory_limit_validation/sort_mem_validation.rs → datafusion/core/tests/memory_limit_validation/sort_mem_validation.rs

@@ -22,7 +22,9 @@
 //! - Test runners that spawn individual test processes
 //! - Test cases that contain the actual validation logic
 
-use crate::memory_limit::memory_limit_validation::utils;
+use datafusion::prelude::SessionConfig;
+
+use crate::utils;
 
 // ===========================================================================
 // Test runners:
@@ -67,6 +69,16 @@ fn sort_with_mem_limit_2_cols_2_runner() {
     utils::spawn_test_process("sort_mem_validation", "sort_with_mem_limit_2_cols_2");
 }
 
+#[test]
+fn sort_strings_no_mem_limit_runner() {
+    utils::spawn_test_process("sort_mem_validation", "sort_strings_no_mem_limit");
+}
+
+#[test]
+fn topk_no_mem_limit_runner() {
+    utils::spawn_test_process("sort_mem_validation", "topk_no_mem_limit");
+}
+
 // ===========================================================================
 // Test cases:
 // All following tests need to be run through their individual test wrapper.
@@ -178,3 +190,28 @@ async fn sort_with_mem_limit_2_cols_2() {
     )
     .await;
 }
+
+/// Sorts on a column the input is not ordered by, so the sort is not removed.
+#[ignore = "sorted runs wait in the merge unreserved, see https://github.com/apache/datafusion/issues/26411"]
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn sort_strings_no_mem_limit() {
+    utils::validate_query_over_rows_with_strings(
+        700_000_000,
+        SessionConfig::new().with_target_partitions(4),
+        "SortExec",
+        |t| format!("SELECT v, s FROM {t} ORDER BY s DESC"),
+    )
+    .await;
+}
+
+/// TopK holds 10 rows, whatever the input size.
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn topk_no_mem_limit() {
+    utils::validate_query_over_rows_with_strings(
+        40_000_000,
+        SessionConfig::new().with_target_partitions(4),
+        "SortExec: TopK",
+        |t| format!("SELECT v, s FROM {t} ORDER BY s DESC LIMIT 10"),
+    )
+    .await;
+}