From 1be9411e1199f9a256cf22c945bb87adaf59e4a6 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sat, 24 Oct 2026 10:42:43 +0800 Subject: [PATCH 1/4] ci: gate extended test targets as a whole instead of filtering by name --- .github/workflows/extended.yml | 2 + datafusion/core/Cargo.toml | 10 ++++ datafusion/core/tests/fuzz.rs | 5 +- .../tests/fuzz_cases/leaf_pushdown_fuzz.rs | 5 +- datafusion/core/tests/memory_limit/mod.rs | 2 - .../main.rs} | 11 +++- .../nlj_mem_validation.rs | 2 +- .../smj_mem_validation.rs | 2 +- .../sort_mem_validation.rs | 2 +- .../memory_limit_validation/utils.rs | 2 +- xtask/src/ci_steps.rs | 59 ++++++++++++------- 11 files changed, 66 insertions(+), 36 deletions(-) rename datafusion/core/tests/{memory_limit/memory_limit_validation/mod.rs => memory_limit_validation/main.rs} (76%) rename datafusion/core/tests/{memory_limit => }/memory_limit_validation/nlj_mem_validation.rs (97%) rename datafusion/core/tests/{memory_limit => }/memory_limit_validation/smj_mem_validation.rs (98%) rename datafusion/core/tests/{memory_limit => }/memory_limit_validation/sort_mem_validation.rs (99%) rename datafusion/core/tests/{memory_limit => }/memory_limit_validation/utils.rs (99%) diff --git a/.github/workflows/extended.yml b/.github/workflows/extended.yml index 0b8434206e282..28262d34b7010 100644 --- a/.github/workflows/extended.yml +++ b/.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 diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 89f42761535e3..11fe91cdf00de 100644 --- a/datafusion/core/Cargo.toml +++ b/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" diff --git a/datafusion/core/tests/fuzz.rs b/datafusion/core/tests/fuzz.rs index 9987721318f83..4ee25da374bf0 100644 --- a/datafusion/core/tests/fuzz.rs +++ b/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)] diff --git a/datafusion/core/tests/fuzz_cases/leaf_pushdown_fuzz.rs b/datafusion/core/tests/fuzz_cases/leaf_pushdown_fuzz.rs index 299423edc1edd..2308abb1afa74 100644 --- a/datafusion/core/tests/fuzz_cases/leaf_pushdown_fuzz.rs +++ b/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; diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index d8814eaa86307..7239bcf850255 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/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; diff --git a/datafusion/core/tests/memory_limit/memory_limit_validation/mod.rs b/datafusion/core/tests/memory_limit_validation/main.rs similarity index 76% rename from datafusion/core/tests/memory_limit/memory_limit_validation/mod.rs rename to datafusion/core/tests/memory_limit_validation/main.rs index aeaffb50af979..27eddc8ca5420 100644 --- a/datafusion/core/tests/memory_limit/memory_limit_validation/mod.rs +++ b/datafusion/core/tests/memory_limit_validation/main.rs @@ -16,9 +16,18 @@ // 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 nlj_mem_validation; mod smj_mem_validation; mod sort_mem_validation; mod utils; + +#[cfg(test)] +#[ctor::ctor(unsafe)] +fn init() { + // Enable RUST_LOG logging configuration for test + let _ = env_logger::try_init(); +} diff --git a/datafusion/core/tests/memory_limit/memory_limit_validation/nlj_mem_validation.rs b/datafusion/core/tests/memory_limit_validation/nlj_mem_validation.rs similarity index 97% rename from datafusion/core/tests/memory_limit/memory_limit_validation/nlj_mem_validation.rs rename to datafusion/core/tests/memory_limit_validation/nlj_mem_validation.rs index 2b38de04d5cda..e4ec1a131d015 100644 --- a/datafusion/core/tests/memory_limit/memory_limit_validation/nlj_mem_validation.rs +++ b/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() { diff --git a/datafusion/core/tests/memory_limit/memory_limit_validation/smj_mem_validation.rs b/datafusion/core/tests/memory_limit_validation/smj_mem_validation.rs similarity index 98% rename from datafusion/core/tests/memory_limit/memory_limit_validation/smj_mem_validation.rs rename to datafusion/core/tests/memory_limit_validation/smj_mem_validation.rs index 3af642fffe101..4dd7627e99d73 100644 --- a/datafusion/core/tests/memory_limit/memory_limit_validation/smj_mem_validation.rs +++ b/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"; diff --git a/datafusion/core/tests/memory_limit/memory_limit_validation/sort_mem_validation.rs b/datafusion/core/tests/memory_limit_validation/sort_mem_validation.rs similarity index 99% rename from datafusion/core/tests/memory_limit/memory_limit_validation/sort_mem_validation.rs rename to datafusion/core/tests/memory_limit_validation/sort_mem_validation.rs index 689915f4aa9db..703deeeb4f5e2 100644 --- a/datafusion/core/tests/memory_limit/memory_limit_validation/sort_mem_validation.rs +++ b/datafusion/core/tests/memory_limit_validation/sort_mem_validation.rs @@ -22,7 +22,7 @@ //! - Test runners that spawn individual test processes //! - Test cases that contain the actual validation logic -use crate::memory_limit::memory_limit_validation::utils; +use crate::utils; // =========================================================================== // Test runners: diff --git a/datafusion/core/tests/memory_limit/memory_limit_validation/utils.rs b/datafusion/core/tests/memory_limit_validation/utils.rs similarity index 99% rename from datafusion/core/tests/memory_limit/memory_limit_validation/utils.rs rename to datafusion/core/tests/memory_limit_validation/utils.rs index aae30418b25f9..ee898c791aaa0 100644 --- a/datafusion/core/tests/memory_limit/memory_limit_validation/utils.rs +++ b/datafusion/core/tests/memory_limit_validation/utils.rs @@ -97,7 +97,7 @@ where /// environment variable set. Re-invokes the current test binary directly, /// avoiding cargo overhead and recompilation. pub fn spawn_test_process(module: &str, test: &str) { - let test_path = format!("memory_limit::memory_limit_validation::{module}::{test}"); + let test_path = format!("{module}::{test}"); let exe = std::env::current_exe().expect("Failed to get test binary path"); let output = Command::new(exe) .arg(&test_path) diff --git a/xtask/src/ci_steps.rs b/xtask/src/ci_steps.rs index 6ee4d73edfb20..6f98dcdf595fa 100644 --- a/xtask/src/ci_steps.rs +++ b/xtask/src/ci_steps.rs @@ -55,7 +55,7 @@ static CI_STEPS: &[StepInfo] = &[ }, StepInfo { command: "test", - help_usage: "test [--explain]", + help_usage: "test [--explain]", help_examples: &["test cli", "test workspace"], help_description: "Run a CI test suite", error_message: "Cargo test step failed", @@ -397,8 +397,8 @@ impl StepContext { } } "extended" => { - // Run only the tests that the `workspace` step does not. The - // merge queue runs that step on the same commit. + // Tests behind the `extended_tests` feature. The `workspace` + // step runs everything else. command .args([ "test", @@ -406,27 +406,31 @@ impl StepContext { "ci", "-p", "datafusion", - "-p", - "datafusion-physical-plan", - "--lib", "--test", "fuzz", "--test", - "core_integration", + "memory_limit_validation", "--features", "avro,backtrace,extended_tests,parquet_encryption", + ]) + .env("RUST_BACKTRACE", "1"); + } + "spill-pool-fuzz" => { + // Run more random scenarios of the spill pool fuzzer than the + // `workspace` step does (about a minute). + command + .args([ + "test", + "--profile", + "ci", + "-p", + "datafusion-physical-plan", + "--lib", "--", - // Tests behind the `extended_tests` feature. A new test - // behind it needs a filter here too, otherwise it never - // runs in CI. - "fuzz_cases::", - "memory_limit_validation::", - // Runs with more iterations, see below. - "spill_pool_scenario_fuzz", + "--exact", + "spill::spill_pool::tests::spill_pool_scenario_fuzz", ]) .env("RUST_BACKTRACE", "1") - // Run more random scenarios of the spill pool fuzzer than the - // default suite does (about a minute). .env("DATAFUSION_SPILL_POOL_FUZZ_ITERATIONS", "1000"); } "hash-collisions" => { @@ -657,7 +661,7 @@ fn required_env(name: &str) -> Result { mod tests { use super::*; - const TEST_USAGE: &str = "usage: cargo xtask ci step test [--explain]"; + const TEST_USAGE: &str = "usage: cargo xtask ci step test [--explain]"; fn args(args: &[&str]) -> Vec { args.iter().map(|arg| (*arg).to_string()).collect() @@ -936,17 +940,28 @@ cargo test \ expected: |actual| { insta::assert_snapshot!(actual, @r" cd /workspace && \ -DATAFUSION_SPILL_POOL_FUZZ_ITERATIONS=1000 \ RUST_BACKTRACE=1 \ cargo test \ --profile ci \ -p datafusion \ +--test fuzz \ +--test memory_limit_validation \ +--features avro,backtrace,extended_tests,parquet_encryption +") + }, + }, + Cmd { + command: &["test", "spill-pool-fuzz"], + expected: |actual| { + insta::assert_snapshot!(actual, @r" +cd /workspace && \ +DATAFUSION_SPILL_POOL_FUZZ_ITERATIONS=1000 \ +RUST_BACKTRACE=1 \ +cargo test \ +--profile ci \ -p datafusion-physical-plan \ --lib \ ---test fuzz \ ---test core_integration \ ---features avro,backtrace,extended_tests,parquet_encryption \ --- fuzz_cases:: memory_limit_validation:: spill_pool_scenario_fuzz +-- --exact spill::spill_pool::tests::spill_pool_scenario_fuzz ") }, }, From 522a13897255b668eb5c3d8cb0d99fb563a3a420 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sat, 24 Oct 2026 10:42:43 +0800 Subject: [PATCH 2/4] ci: fail tests when an extended CI test would silently stop running --- .../tests/memory_limit_validation/utils.rs | 5 ++ xtask/src/ci_steps.rs | 54 ++++++++++++++++++- 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/datafusion/core/tests/memory_limit_validation/utils.rs b/datafusion/core/tests/memory_limit_validation/utils.rs index ee898c791aaa0..7f4d5b15aba98 100644 --- a/datafusion/core/tests/memory_limit_validation/utils.rs +++ b/datafusion/core/tests/memory_limit_validation/utils.rs @@ -114,6 +114,11 @@ pub fn spawn_test_process(module: &str, test: &str) { "Test '{test}' failed with status: {}\nstdout:\n{stdout}\nstderr:\n{stderr}", output.status, ); + // A path that matches no test passes without running anything. + assert!( + stdout.contains("running 1 test"), + "Test path '{test_path}' did not select exactly one test\nstdout:\n{stdout}", + ); } fn operator_spill_count(plan: &dyn ExecutionPlan, operator_name: &str) -> usize { diff --git a/xtask/src/ci_steps.rs b/xtask/src/ci_steps.rs index 6f98dcdf595fa..7793cb8d0c0f9 100644 --- a/xtask/src/ci_steps.rs +++ b/xtask/src/ci_steps.rs @@ -30,6 +30,8 @@ type StepRunner = fn(&StepContext, &[String]) -> StepResult; const CI_COMMAND: &str = "cargo xtask ci step"; const CI_SHORTCUT: &str = "cargo ci-step"; const TPCH_DATA_DIR: &str = "datafusion/sqllogictest/test_files/tpch/data"; +/// The `datafusion-physical-plan` unit test that the `spill-pool-fuzz` step runs. +const SPILL_POOL_FUZZ_TEST: &str = "spill::spill_pool::tests::spill_pool_scenario_fuzz"; /// Metadata and implementation for one CI step. /// @@ -428,7 +430,7 @@ impl StepContext { "--lib", "--", "--exact", - "spill::spill_pool::tests::spill_pool_scenario_fuzz", + SPILL_POOL_FUZZ_TEST, ]) .env("RUST_BACKTRACE", "1") .env("DATAFUSION_SPILL_POOL_FUZZ_ITERATIONS", "1000"); @@ -710,6 +712,56 @@ mod tests { assert_eq!(error, TEST_USAGE); } + /// A filter that matches no test passes without running anything, so a + /// renamed or moved spill pool fuzzer would silently stop running in CI. + #[test] + fn spill_pool_fuzz_test_exists() { + let (module, test) = SPILL_POOL_FUZZ_TEST + .rsplit_once("::tests::") + .expect("test path has a `tests` module"); + let file = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../datafusion/physical-plan/src") + .join(format!("{}.rs", module.replace("::", "/"))); + let source = std::fs::read_to_string(&file) + .unwrap_or_else(|e| panic!("cannot read {}: {e}", file.display())); + assert!( + source.contains("mod tests {") && source.contains(&format!("fn {test}(")), + "{} has no test `{test}` in `mod tests`; update SPILL_POOL_FUZZ_TEST", + file.display() + ); + } + + /// The `extended` step runs only the test targets that require the + /// `extended_tests` feature. A test gated on the feature with a `cfg` in any + /// other target would never run in CI. + #[test] + fn extended_tests_feature_only_gates_whole_targets() { + fn visit(dir: &Path, found: &mut Vec) { + for entry in std::fs::read_dir(dir).unwrap() { + let path = entry.unwrap().path(); + if path.is_dir() { + visit(&path, found); + } else if path.extension() == Some(OsStr::new("rs")) + && std::fs::read_to_string(&path) + .unwrap() + .contains("feature = \"extended_tests\"") + { + found.push(path); + } + } + } + let mut found = vec![]; + visit( + &Path::new(env!("CARGO_MANIFEST_DIR")).join("../datafusion"), + &mut found, + ); + assert!( + found.is_empty(), + "put tests behind `extended_tests` in a test target that requires the \ + feature (see datafusion/core/Cargo.toml), not behind a `cfg` in {found:?}" + ); + } + #[test] fn unknown_test_variant_is_rejected() { let error = context() From 42699c39ba1460c6f98ff04046fd84dd4f14d820 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sun, 25 Oct 2026 09:53:09 +0800 Subject: [PATCH 3/4] test: move the counting allocator from sqllogictest to test-utils --- Cargo.lock | 1 + datafusion/sqllogictest/Cargo.toml | 1 + datafusion/sqllogictest/bin/sqllogictests.rs | 13 ++- datafusion/sqllogictest/src/lib.rs | 5 +- datafusion/sqllogictest/src/memory_drift.rs | 108 ++---------------- test-utils/src/counting_allocator.rs | 113 +++++++++++++++++++ test-utils/src/lib.rs | 1 + 7 files changed, 135 insertions(+), 107 deletions(-) create mode 100644 test-utils/src/counting_allocator.rs diff --git a/Cargo.lock b/Cargo.lock index be2c79d42d304..90570f3720670 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2768,6 +2768,7 @@ dependencies = [ "sqllogictest", "sqlparser", "tempfile", + "test-utils", "testcontainers-modules", "thiserror", "tokio", diff --git a/datafusion/sqllogictest/Cargo.toml b/datafusion/sqllogictest/Cargo.toml index c43f78f846092..4ff3b561914f7 100644 --- a/datafusion/sqllogictest/Cargo.toml +++ b/datafusion/sqllogictest/Cargo.toml @@ -88,6 +88,7 @@ regex = { workspace = true } # Required to make sure tests for pg display behaves consistently # regardless of feature unification with dependencies serde_json = { workspace = true, features = ["preserve_order"] } +test-utils = { path = "../../test-utils" } tokio = { workspace = true, features = ["rt-multi-thread"] } [[test]] diff --git a/datafusion/sqllogictest/bin/sqllogictests.rs b/datafusion/sqllogictest/bin/sqllogictests.rs index 1fd5d9d469c71..541c8568ede26 100644 --- a/datafusion/sqllogictest/bin/sqllogictests.rs +++ b/datafusion/sqllogictest/bin/sqllogictests.rs @@ -23,15 +23,12 @@ use datafusion::execution::memory_pool::DEFAULT_DRIFT_LOG_THRESHOLD; #[cfg(feature = "substrait")] use datafusion_sqllogictest::DataFusionSubstraitRoundTrip; use datafusion_sqllogictest::TestFile; -use datafusion_sqllogictest::{ - CountingAllocator, enable_memory_drift_logging, flush_thread_allocations, - memory_drift_tracker, -}; use datafusion_sqllogictest::{ CurrentlyExecutingSqlTracker, DFColumnType, DataFusion, Filter, TestContext, df_value_validator, read_dir_recursive, run_each_configuration, setup_scratch_dir, should_skip_file, should_skip_record, test_configurations, value_normalizer, }; +use datafusion_sqllogictest::{enable_memory_drift_logging, memory_drift_tracker}; use futures::stream::StreamExt; use indicatif::{ HumanDuration, MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle, @@ -43,6 +40,9 @@ use sqllogictest::{ AsyncDB, Condition, MakeConnection, Normalizer, Record, Validator, parse_file, strict_column_validator, }; +use test_utils::counting_allocator::{ + CountingAllocator, allocated_bytes, flush_thread_allocations, +}; #[cfg(feature = "postgres")] use crate::postgres_container::{ @@ -149,7 +149,10 @@ async fn run_tests() -> Result<()> { options.warn_on_ignored(); if options.memory_drift { - enable_memory_drift_logging(options.memory_drift_log_threshold); + enable_memory_drift_logging( + options.memory_drift_log_threshold, + Arc::new(allocated_bytes), + ); } // Print parallelism info for debugging CI performance diff --git a/datafusion/sqllogictest/src/lib.rs b/datafusion/sqllogictest/src/lib.rs index a397763933a81..1c01c1807e966 100644 --- a/datafusion/sqllogictest/src/lib.rs +++ b/datafusion/sqllogictest/src/lib.rs @@ -50,10 +50,7 @@ mod util; pub use config_matrix::{TestConfiguration, run_each_configuration, test_configurations}; pub use filters::*; -pub use memory_drift::{ - CountingAllocator, enable_memory_drift_logging, flush_thread_allocations, - memory_drift_tracker, -}; +pub use memory_drift::{enable_memory_drift_logging, memory_drift_tracker}; pub use test_context::TestContext; pub use test_file::TestFile; pub use util::*; diff --git a/datafusion/sqllogictest/src/memory_drift.rs b/datafusion/sqllogictest/src/memory_drift.rs index 2aff6499bfc71..aee07c2f4c0d2 100644 --- a/datafusion/sqllogictest/src/memory_drift.rs +++ b/datafusion/sqllogictest/src/memory_drift.rs @@ -21,122 +21,34 @@ //! Test files run concurrently in one process, so allocations cannot be split //! per file. Instead every file's pool reports to one process-wide //! [`MemoryDriftTracker`], which compares the sum of all reservations with the -//! bytes counted by [`CountingAllocator`]. +//! allocated bytes passed to [`enable_memory_drift_logging`]. //! //! Files that `SET datafusion.runtime.memory_limit` replace their pool; the //! runner wraps the replacement too (see [`rewrap_replaced_pool`]). //! //! This only logs. It never fails a test. -use std::{ - alloc::{GlobalAlloc, Layout, System}, - cell::Cell, - sync::{ - Arc, OnceLock, - atomic::{AtomicIsize, Ordering}, - }, -}; +use std::sync::{Arc, OnceLock}; use datafusion::execution::SessionStateBuilder; use datafusion::execution::memory_pool::{ - MemoryDriftTracker, MemoryPool, PeakRecordingPool, + AllocatedBytesFn, MemoryDriftTracker, MemoryPool, PeakRecordingPool, }; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::prelude::SessionContext; -static ALLOCATED: AtomicIsize = AtomicIsize::new(0); static TRACKER: OnceLock> = OnceLock::new(); -/// A [`GlobalAlloc`] that counts the bytes currently allocated through it, -/// delegating the allocation itself to [`System`]. -/// -/// Counts requested sizes, so allocator overhead and memory retained by the -/// allocator are not included. Counting starts with the process, so memory -/// freed later was always counted when it was allocated. -pub struct CountingAllocator; - -/// Per-thread count is flushed to [`ALLOCATED`] once it moves this far, so -/// threads do not contend on one atomic for every allocation. The global count -/// is therefore off by up to `FLUSH_BYTES` per live thread. Tokio threads -/// flush the rest when they stop (see [`flush_thread_allocations`]); the -/// unflushed count of any other thread that exits is lost. -const FLUSH_BYTES: isize = 256 * 1024; - -thread_local! { - // `const` with no destructor, so accessing it never allocates. - static UNFLUSHED: Cell = const { Cell::new(0) }; -} - -fn count(delta: isize) { - // `try_with` because this can run while the thread is being torn down. - let _ = UNFLUSHED.try_with(|unflushed| { - let pending = unflushed.get() + delta; - if pending.abs() >= FLUSH_BYTES { - ALLOCATED.fetch_add(pending, Ordering::Relaxed); - unflushed.set(0); - } else { - unflushed.set(pending); - } - }); -} - -/// Flush this thread's unflushed count into the global count. -/// -/// Registered as the Tokio runtime's `on_thread_stop` hook, which runs just -/// before a worker or blocking thread exits, outside the allocator. -pub fn flush_thread_allocations() { - let _ = UNFLUSHED.try_with(|unflushed| { - ALLOCATED.fetch_add(unflushed.replace(0), Ordering::Relaxed); - }); -} - -// SAFETY: every call is forwarded unchanged to `System`, which upholds the -// `GlobalAlloc` contract. Counting has no effect on the returned memory. -unsafe impl GlobalAlloc for CountingAllocator { - unsafe fn alloc(&self, layout: Layout) -> *mut u8 { - let ptr = unsafe { System.alloc(layout) }; - if !ptr.is_null() { - count(layout.size() as isize); - } - ptr - } - - unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { - let ptr = unsafe { System.alloc_zeroed(layout) }; - if !ptr.is_null() { - count(layout.size() as isize); - } - ptr - } - - unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { - unsafe { System.dealloc(ptr, layout) }; - count(-(layout.size() as isize)); - } - - unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { - let new_ptr = unsafe { System.realloc(ptr, layout, new_size) }; - if !new_ptr.is_null() { - count(new_size as isize - layout.size() as isize); - } - new_ptr - } -} - -/// Bytes currently allocated through [`CountingAllocator`], clamped at zero -/// because counts lost with exiting threads can bias it either way. -fn allocated_bytes() -> usize { - ALLOCATED.load(Ordering::Relaxed).max(0) as usize -} - /// Wrap the memory pool of every test file created from now on, logging each -/// rise in drift of `log_threshold` bytes. Only meaningful if -/// [`CountingAllocator`] is the global allocator. -pub fn enable_memory_drift_logging(log_threshold: usize) { +/// rise in drift of `log_threshold` bytes. `allocated_bytes` returns the bytes +/// the process has allocated, e.g. from a counting global allocator. +pub fn enable_memory_drift_logging( + log_threshold: usize, + allocated_bytes: AllocatedBytesFn, +) { TRACKER.get_or_init(|| { Arc::new( - MemoryDriftTracker::new(Arc::new(allocated_bytes)) - .with_log_threshold(log_threshold), + MemoryDriftTracker::new(allocated_bytes).with_log_threshold(log_threshold), ) }); } diff --git a/test-utils/src/counting_allocator.rs b/test-utils/src/counting_allocator.rs new file mode 100644 index 0000000000000..405ead93d7535 --- /dev/null +++ b/test-utils/src/counting_allocator.rs @@ -0,0 +1,113 @@ +// 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. + +//! A global allocator that counts the bytes currently allocated, so that tests +//! can compare them with what memory pools have reserved. + +use std::{ + alloc::{GlobalAlloc, Layout, System}, + cell::Cell, + sync::atomic::{AtomicIsize, Ordering}, +}; + +static ALLOCATED: AtomicIsize = AtomicIsize::new(0); + +/// A [`GlobalAlloc`] that counts the bytes currently allocated through it, +/// delegating the allocation itself to [`System`]. +/// +/// Counts requested sizes, so allocator overhead and memory retained by the +/// allocator are not included. Counting starts with the process, so memory +/// freed later was always counted when it was allocated. +/// +/// Install it with `#[global_allocator]` and read the count with +/// [`allocated_bytes`]. +pub struct CountingAllocator; + +/// Per-thread count is flushed to [`ALLOCATED`] once it moves this far, so +/// threads do not contend on one atomic for every allocation. The global count +/// is therefore off by up to `FLUSH_BYTES` per live thread. A Tokio runtime +/// that registers [`flush_thread_allocations`] as its `on_thread_stop` hook +/// flushes the rest when its threads stop; the unflushed count of any other +/// thread that exits is lost. +const FLUSH_BYTES: isize = 256 * 1024; + +thread_local! { + // `const` with no destructor, so accessing it never allocates. + static UNFLUSHED: Cell = const { Cell::new(0) }; +} + +fn count(delta: isize) { + // `try_with` because this can run while the thread is being torn down. + let _ = UNFLUSHED.try_with(|unflushed| { + let pending = unflushed.get() + delta; + if pending.abs() >= FLUSH_BYTES { + ALLOCATED.fetch_add(pending, Ordering::Relaxed); + unflushed.set(0); + } else { + unflushed.set(pending); + } + }); +} + +/// Flush this thread's unflushed count into the global count. +/// +/// Register it as a Tokio runtime's `on_thread_stop` hook, which runs just +/// before a worker or blocking thread exits, outside the allocator. +pub fn flush_thread_allocations() { + let _ = UNFLUSHED.try_with(|unflushed| { + ALLOCATED.fetch_add(unflushed.replace(0), Ordering::Relaxed); + }); +} + +// SAFETY: every call is forwarded unchanged to `System`, which upholds the +// `GlobalAlloc` contract. Counting has no effect on the returned memory. +unsafe impl GlobalAlloc for CountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { System.alloc(layout) }; + if !ptr.is_null() { + count(layout.size() as isize); + } + ptr + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { System.alloc_zeroed(layout) }; + if !ptr.is_null() { + count(layout.size() as isize); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { System.dealloc(ptr, layout) }; + count(-(layout.size() as isize)); + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let new_ptr = unsafe { System.realloc(ptr, layout, new_size) }; + if !new_ptr.is_null() { + count(new_size as isize - layout.size() as isize); + } + new_ptr + } +} + +/// Bytes currently allocated through [`CountingAllocator`], clamped at zero +/// because counts lost with exiting threads can bias it either way. +pub fn allocated_bytes() -> usize { + ALLOCATED.load(Ordering::Relaxed).max(0) as usize +} diff --git a/test-utils/src/lib.rs b/test-utils/src/lib.rs index 55717c717c4af..2eec76bc86b7f 100644 --- a/test-utils/src/lib.rs +++ b/test-utils/src/lib.rs @@ -24,6 +24,7 @@ use rand::prelude::StdRng; use rand::{Rng, SeedableRng}; pub mod array_gen; +pub mod counting_allocator; mod data_gen; mod string_gen; pub mod tpcds; From d3c7fb9e6e957ca090b3059f8368abcb52a1b20d Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sun, 25 Oct 2026 09:53:33 +0800 Subject: [PATCH 4/4] test: catch operators that hold input-sized memory without reserving it --- .../aggregate_mem_validation.rs | 117 +++++++++++++++ .../hash_join_mem_validation.rs | 60 ++++++++ .../tests/memory_limit_validation/main.rs | 7 + .../sort_mem_validation.rs | 37 +++++ .../tests/memory_limit_validation/utils.rs | 140 +++++++++++++++--- .../window_mem_validation.rs | 42 ++++++ 6 files changed, 381 insertions(+), 22 deletions(-) create mode 100644 datafusion/core/tests/memory_limit_validation/aggregate_mem_validation.rs create mode 100644 datafusion/core/tests/memory_limit_validation/hash_join_mem_validation.rs create mode 100644 datafusion/core/tests/memory_limit_validation/window_mem_validation.rs diff --git a/datafusion/core/tests/memory_limit_validation/aggregate_mem_validation.rs b/datafusion/core/tests/memory_limit_validation/aggregate_mem_validation.rs new file mode 100644 index 0000000000000..25a06cfbf527e --- /dev/null +++ b/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; +} diff --git a/datafusion/core/tests/memory_limit_validation/hash_join_mem_validation.rs b/datafusion/core/tests/memory_limit_validation/hash_join_mem_validation.rs new file mode 100644 index 0000000000000..e4deb48c44244 --- /dev/null +++ b/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; +} diff --git a/datafusion/core/tests/memory_limit_validation/main.rs b/datafusion/core/tests/memory_limit_validation/main.rs index 27eddc8ca5420..ca81e87f24442 100644 --- a/datafusion/core/tests/memory_limit_validation/main.rs +++ b/datafusion/core/tests/memory_limit_validation/main.rs @@ -20,10 +20,17 @@ //! //! 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)] diff --git a/datafusion/core/tests/memory_limit_validation/sort_mem_validation.rs b/datafusion/core/tests/memory_limit_validation/sort_mem_validation.rs index 703deeeb4f5e2..b86e44b1e808a 100644 --- a/datafusion/core/tests/memory_limit_validation/sort_mem_validation.rs +++ b/datafusion/core/tests/memory_limit_validation/sort_mem_validation.rs @@ -22,6 +22,8 @@ //! - Test runners that spawn individual test processes //! - Test cases that contain the actual validation logic +use datafusion::prelude::SessionConfig; + use crate::utils; // =========================================================================== @@ -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; +} diff --git a/datafusion/core/tests/memory_limit_validation/utils.rs b/datafusion/core/tests/memory_limit_validation/utils.rs index 7f4d5b15aba98..8b4dcf77144bb 100644 --- a/datafusion/core/tests/memory_limit_validation/utils.rs +++ b/datafusion/core/tests/memory_limit_validation/utils.rs @@ -16,6 +16,7 @@ // under the License. use datafusion_common_runtime::SpawnedTask; +use futures::StreamExt; use std::process::Command; use std::str; use std::sync::Arc; @@ -23,22 +24,30 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use sysinfo::{ProcessRefreshKind, ProcessesToUpdate, System}; use tokio::time::{Duration, interval}; -use datafusion::physical_plan::{ExecutionPlan, collect, displayable}; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::{ExecutionPlan, displayable, execute_stream}; use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_common::human_readable_size; -use datafusion_execution::{memory_pool::FairSpillPool, runtime_env::RuntimeEnvBuilder}; +use datafusion_execution::memory_pool::{ + FairSpillPool, MemoryDriftTracker, MemoryPool, PeakRecordingPool, UnboundedMemoryPool, +}; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; + +use test_utils::counting_allocator::allocated_bytes; /// Measures the maximum RSS (in bytes) during the execution of an async task. RSS -/// will be sampled every 7ms. +/// will be sampled every 7ms, and so will `drift`, in addition to the samples it +/// takes on every reservation change. /// /// # Arguments /// +/// * `drift` - The tracker of the memory pool the task uses. /// * `f` - A closure that returns the async task to be measured. /// /// # Returns /// /// A tuple containing the result of the async task and the maximum RSS observed. -async fn measure_max_rss(f: F) -> (T, usize) +async fn measure_max_rss(drift: &Arc, f: F) -> (T, usize) where F: FnOnce() -> Fut, Fut: Future, @@ -55,6 +64,7 @@ where // Clone for the monitoring task let max_rss_clone = Arc::clone(&max_rss); + let drift = Arc::clone(drift); // Spawn a monitoring task let monitor_handle = SpawnedTask::spawn(async move { @@ -63,6 +73,7 @@ where loop { interval.tick().await; + drift.sample("sampler"); sys.refresh_processes_specifics( ProcessesToUpdate::Some(&[pid]), true, @@ -182,6 +193,38 @@ pub async fn validate_query_with_memory_limits( .await; } +/// Validates `sql` with no memory limit, over [`rows_with_strings`] of 10M rows +/// and 1M rows for the baseline. `operator` must be in the plan. +/// +/// Callers run under a multi-threaded runtime: some operators only buffer +/// unreserved batches when their inputs run ahead on other threads. +pub async fn validate_query_over_rows_with_strings( + expected_mem_bytes: i64, + session_config: SessionConfig, + operator: &str, + sql: impl Fn(&str) -> String, +) { + validate_query_with_memory_limits_and_config( + expected_mem_bytes, + None, + &sql(&rows_with_strings(10_000_000)), + &sql(&rows_with_strings(1_000_000)), + session_config, + Some(operator), + None, + ) + .await; +} + +/// A table of `rows` rows: `v` is an Int64 from 1 to `rows`, in order, and `s` +/// is a string of about 25 bytes. +fn rows_with_strings(rows: usize) -> String { + format!( + "(SELECT v, concat('payload-', v, '-xxxxxxxx') AS s \ + FROM generate_series(1, {rows}) AS g(v))" + ) +} + /// Validate memory usage with a custom session configuration and optional /// operator and spill assertions. pub async fn validate_query_with_memory_limits_and_config( @@ -201,18 +244,7 @@ pub async fn validate_query_with_memory_limits_and_config( println!("Current process ID: {}", std::process::id()); - let runtime_builder = RuntimeEnvBuilder::new(); - - let runtime = match mem_limit_bytes { - Some(mem_limit_bytes) => runtime_builder - .with_memory_pool(Arc::new(FairSpillPool::new(mem_limit_bytes as usize))) - .build_arc() - .unwrap(), - None => runtime_builder.build_arc().unwrap(), - }; - - let ctx = SessionContext::new_with_config_rt(session_config, runtime); - + let (ctx, drift) = session(mem_limit_bytes, session_config.clone()); let df = ctx.sql(query).await.unwrap(); let physical_plan = df.create_physical_plan().await.unwrap(); @@ -225,23 +257,24 @@ pub async fn validate_query_with_memory_limits_and_config( } // Run a query with 10% data to estimate the constant overhead - let baseline_plan = ctx + let (baseline_ctx, baseline_drift) = session(mem_limit_bytes, session_config); + let baseline_plan = baseline_ctx .sql(baseline_query) .await .unwrap() .create_physical_plan() .await .unwrap(); - let baseline_task_ctx = ctx.task_ctx(); - let (_, baseline_max_rss) = measure_max_rss(|| async move { - collect(baseline_plan, baseline_task_ctx).await.unwrap() + let baseline_task_ctx = baseline_ctx.task_ctx(); + let (_, baseline_max_rss) = measure_max_rss(&baseline_drift, || { + run_to_completion(baseline_plan, baseline_task_ctx) }) .await; let execution_plan = Arc::clone(&physical_plan); let execution_task_ctx = ctx.task_ctx(); - let (_, max_rss) = measure_max_rss(|| async move { - collect(execution_plan, execution_task_ctx).await.unwrap() + let (_, max_rss) = measure_max_rss(&drift, || { + run_to_completion(execution_plan, execution_task_ctx) }) .await; @@ -271,10 +304,73 @@ pub async fn validate_query_with_memory_limits_and_config( human_readable_size(expected_mem_bytes as usize) ); + let baseline_peak = baseline_drift + .peak_drift() + .expect("the sampler samples drift"); + let peak = drift.peak_drift().expect("the sampler samples drift"); + let unreserved_growth = peak.drift - baseline_peak.drift; + println!( + "Unreserved memory at its peak: before: {baseline_peak}, after: {peak}, growth: {}", + signed_size(unreserved_growth), + ); + assert!( actual_mem_usage < expected_mem_bytes as f64, "Memory usage exceeded the theoretical limit. Actual: {}, Expected limit: {}", human_readable_size(actual_mem_usage as usize), human_readable_size(expected_mem_bytes as usize) ); + assert!( + unreserved_growth < MAX_UNRESERVED_GROWTH, + "Unreserved memory grew with the input by {}, more than {}. Memory proportional \ + to the input must be reserved, see the `MemoryPool` docs", + signed_size(unreserved_growth), + human_readable_size(MAX_UNRESERVED_GROWTH as usize), + ); +} + +/// Largest growth of unreserved memory (allocated minus reserved, at its peak) +/// allowed from the baseline query to the main query. +/// +/// Batches in flight between operators need not be reserved, but they are the +/// same size whatever the input size, so they cancel out. Memory that grows +/// with the input must be reserved. The allocated byte count is off by at most +/// 256 KB per thread, far below this. +const MAX_UNRESERVED_GROWTH: isize = 16 * 1024 * 1024; + +/// A session whose memory pool reports every reservation change to the +/// returned tracker, which compares the bytes reserved with the bytes allocated. +fn session( + mem_limit_bytes: Option, + session_config: SessionConfig, +) -> (SessionContext, Arc) { + let pool: Arc = match mem_limit_bytes { + Some(mem_limit_bytes) => Arc::new(FairSpillPool::new(mem_limit_bytes as usize)), + None => Arc::new(UnboundedMemoryPool::default()), + }; + let drift = Arc::new(MemoryDriftTracker::new(Arc::new(allocated_bytes))); + let pool = + PeakRecordingPool::new(pool).with_drift_tracker(Arc::clone(&drift), "query"); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(pool)) + .build_arc() + .unwrap(); + ( + SessionContext::new_with_config_rt(session_config, runtime), + drift, + ) +} + +/// Runs `plan` and drops each output batch as it arrives, so the output does not +/// count as memory the query holds. +async fn run_to_completion(plan: Arc, task_ctx: Arc) { + let mut stream = execute_stream(plan, task_ctx).unwrap(); + while let Some(batch) = stream.next().await { + batch.unwrap(); + } +} + +fn signed_size(bytes: isize) -> String { + let size = human_readable_size(bytes.unsigned_abs()); + if bytes < 0 { format!("-{size}") } else { size } } diff --git a/datafusion/core/tests/memory_limit_validation/window_mem_validation.rs b/datafusion/core/tests/memory_limit_validation/window_mem_validation.rs new file mode 100644 index 0000000000000..fa0efc7182f08 --- /dev/null +++ b/datafusion/core/tests/memory_limit_validation/window_mem_validation.rs @@ -0,0 +1,42 @@ +// 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 window functions. +//! +//! These tests run in separate processes to accurately measure memory usage. + +use datafusion::prelude::SessionConfig; + +use crate::utils; + +#[test] +fn window_over_all_rows_runner() { + utils::spawn_test_process("window_mem_validation", "window_over_all_rows"); +} + +/// The window spans the whole input, so all of it is buffered. +#[ignore = "window operators reserve no memory, see https://github.com/apache/datafusion/issues/26707"] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn window_over_all_rows() { + utils::validate_query_over_rows_with_strings( + 80_000_000 * 4, + SessionConfig::new().with_target_partitions(4), + "WindowAggExec", + |t| format!("SELECT v, sum(v) OVER () FROM {t}"), + ) + .await; +}