diff --git a/datafusion/core/tests/memory_limit/memory_limit_validation/mod.rs b/datafusion/core/tests/memory_limit/memory_limit_validation/mod.rs index 32df6c5d62937..83ebb266c8257 100644 --- a/datafusion/core/tests/memory_limit/memory_limit_validation/mod.rs +++ b/datafusion/core/tests/memory_limit/memory_limit_validation/mod.rs @@ -18,5 +18,6 @@ //! Validates query's actual memory usage is consistent with the specified memory //! limit. +mod smj_mem_validation; mod sort_mem_validation; mod utils; diff --git a/datafusion/core/tests/memory_limit/memory_limit_validation/smj_mem_validation.rs b/datafusion/core/tests/memory_limit/memory_limit_validation/smj_mem_validation.rs new file mode 100644 index 0000000000000..3af642fffe101 --- /dev/null +++ b/datafusion/core/tests/memory_limit/memory_limit_validation/smj_mem_validation.rs @@ -0,0 +1,105 @@ +// 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-limit validation tests for sort-merge join queries. +//! +//! These tests run in separate processes to accurately measure memory usage. + +use datafusion::prelude::SessionConfig; + +use crate::memory_limit::memory_limit_validation::utils; + +/// Ensures the planner selected a sort-merge join. +const SMJ_OPERATOR_NAME: &str = "SortMergeJoinExec"; + +/// Configure a two-partition sort-merge join and reduce the sort reservation so +/// the join can spill under the tested memory limits. +fn smj_session_config() -> SessionConfig { + SessionConfig::new() + .with_target_partitions(2) + .with_sort_spill_reservation_bytes(1024 * 1024) + .set_bool("datafusion.optimizer.prefer_hash_join", false) +} + +/// Build a join with one large buffered key group and scalar output. +fn smj_sum_query(series_len: usize) -> String { + format!( + "SELECT sum(rr.v) FROM generate_series(0, 0) AS l(k) \ + JOIN (SELECT i % 1 AS k, i AS v FROM generate_series(1, {series_len}) AS r(i)) rr \ + ON l.k = rr.k" + ) +} + +#[test] +fn smj_with_mem_limit_1_runner() { + utils::spawn_test_process("smj_mem_validation", "smj_with_mem_limit_1"); +} + +#[test] +fn smj_with_mem_limit_2_runner() { + utils::spawn_test_process("smj_mem_validation", "smj_with_mem_limit_2"); +} + +#[test] +fn smj_no_mem_limit_runner() { + utils::spawn_test_process("smj_mem_validation", "smj_no_mem_limit"); +} + +/// Verify a 40 MB pool forces spilling within the RSS allowance. +#[tokio::test] +async fn smj_with_mem_limit_1() { + utils::validate_query_with_memory_limits_and_config( + 40_000_000 * 4, + Some(40_000_000), + &smj_sum_query(5_000_000), + &smj_sum_query(500_000), + smj_session_config(), + Some(SMJ_OPERATOR_NAME), + Some(true), + ) + .await; +} + +/// Verify a 16 MB pool forces spilling. The 5M join keys (~40 MB) stay resident +/// independently of the pool limit, so this case needs a larger RSS allowance. +#[tokio::test] +async fn smj_with_mem_limit_2() { + utils::validate_query_with_memory_limits_and_config( + 16_000_000 * 12, + Some(16_000_000), + &smj_sum_query(5_000_000), + &smj_sum_query(500_000), + smj_session_config(), + Some(SMJ_OPERATOR_NAME), + Some(true), + ) + .await; +} + +#[tokio::test] +async fn smj_no_mem_limit() { + utils::validate_query_with_memory_limits_and_config( + 40_000_000 * 5, + None, + &smj_sum_query(5_000_000), + &smj_sum_query(500_000), + smj_session_config(), + Some(SMJ_OPERATOR_NAME), + Some(false), + ) + .await; +} diff --git a/datafusion/core/tests/memory_limit/memory_limit_validation/sort_mem_validation.rs b/datafusion/core/tests/memory_limit/memory_limit_validation/sort_mem_validation.rs index bf04123fff7fa..b55a3039ec9d4 100644 --- a/datafusion/core/tests/memory_limit/memory_limit_validation/sort_mem_validation.rs +++ b/datafusion/core/tests/memory_limit/memory_limit_validation/sort_mem_validation.rs @@ -21,7 +21,6 @@ //! This file is organized as: //! - Test runners that spawn individual test processes //! - Test cases that contain the actual validation logic -use std::{process::Command, str}; use crate::memory_limit::memory_limit_validation::utils; @@ -32,67 +31,40 @@ use crate::memory_limit::memory_limit_validation::utils; #[test] fn memory_limit_validation_runner_works_runner() { - spawn_test_process("memory_limit_validation_runner_works"); + utils::spawn_test_process( + "sort_mem_validation", + "memory_limit_validation_runner_works", + ); } #[test] fn sort_no_mem_limit_runner() { - spawn_test_process("sort_no_mem_limit"); + utils::spawn_test_process("sort_mem_validation", "sort_no_mem_limit"); } #[test] fn sort_with_mem_limit_1_runner() { - spawn_test_process("sort_with_mem_limit_1"); + utils::spawn_test_process("sort_mem_validation", "sort_with_mem_limit_1"); } #[test] fn sort_with_mem_limit_2_runner() { - spawn_test_process("sort_with_mem_limit_2"); + utils::spawn_test_process("sort_mem_validation", "sort_with_mem_limit_2"); } #[test] fn sort_with_mem_limit_3_runner() { - spawn_test_process("sort_with_mem_limit_3"); + utils::spawn_test_process("sort_mem_validation", "sort_with_mem_limit_3"); } #[test] fn sort_with_mem_limit_2_cols_1_runner() { - spawn_test_process("sort_with_mem_limit_2_cols_1"); + utils::spawn_test_process("sort_mem_validation", "sort_with_mem_limit_2_cols_1"); } #[test] fn sort_with_mem_limit_2_cols_2_runner() { - spawn_test_process("sort_with_mem_limit_2_cols_2"); -} - -/// Helper function that executes a test in a separate process with the required -/// environment variable set. Re-invokes the current test binary directly, -/// avoiding cargo overhead and recompilation. -fn spawn_test_process(test: &str) { - let test_path = - format!("memory_limit::memory_limit_validation::sort_mem_validation::{test}"); - - let exe = std::env::current_exe().expect("Failed to get test binary path"); - - let output = Command::new(exe) - .arg(&test_path) - .arg("--exact") - .arg("--nocapture") - .env("DATAFUSION_TEST_MEM_LIMIT_VALIDATION", "1") - .output() - .expect("Failed to execute test command"); - - let stdout = str::from_utf8(&output.stdout).unwrap_or(""); - let stderr = str::from_utf8(&output.stderr).unwrap_or(""); - - assert!( - output.status.success(), - "Test '{}' failed with status: {}\nstdout:\n{}\nstderr:\n{}", - test, - output.status, - stdout, - stderr - ); + utils::spawn_test_process("sort_mem_validation", "sort_with_mem_limit_2_cols_2"); } // =========================================================================== diff --git a/datafusion/core/tests/memory_limit/memory_limit_validation/utils.rs b/datafusion/core/tests/memory_limit/memory_limit_validation/utils.rs index 2c9fae20c8606..788b8f4942ee4 100644 --- a/datafusion/core/tests/memory_limit/memory_limit_validation/utils.rs +++ b/datafusion/core/tests/memory_limit/memory_limit_validation/utils.rs @@ -16,11 +16,14 @@ // under the License. use datafusion_common_runtime::SpawnedTask; +use std::process::Command; +use std::str; use std::sync::Arc; 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::prelude::{SessionConfig, SessionContext}; use datafusion_common::human_readable_size; use datafusion_execution::{memory_pool::FairSpillPool, runtime_env::RuntimeEnvBuilder}; @@ -98,6 +101,42 @@ where (result, peak_rss) } +/// Helper function that executes a test in a separate process with the required +/// 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 exe = std::env::current_exe().expect("Failed to get test binary path"); + let output = Command::new(exe) + .arg(&test_path) + .arg("--exact") + .arg("--nocapture") + .env("DATAFUSION_TEST_MEM_LIMIT_VALIDATION", "1") + .output() + .expect("Failed to execute test command"); + + let stdout = str::from_utf8(&output.stdout).unwrap_or(""); + let stderr = str::from_utf8(&output.stderr).unwrap_or(""); + assert!( + output.status.success(), + "Test '{test}' failed with status: {}\nstdout:\n{stdout}\nstderr:\n{stderr}", + output.status, + ); +} + +fn operator_spill_count(plan: &dyn ExecutionPlan, operator_name: &str) -> usize { + let own = if plan.name() == operator_name { + plan.metrics().and_then(|m| m.spill_count()).unwrap_or(0) + } else { + 0 + }; + own + plan + .children() + .into_iter() + .map(|child| operator_spill_count(child.as_ref(), operator_name)) + .sum::() +} + /// Query runner that validates the memory usage of the query. /// /// Note this function is supposed to run in a separate process for accurate memory @@ -132,6 +171,30 @@ pub async fn validate_query_with_memory_limits( mem_limit_bytes: Option, query: &str, baseline_query: &str, +) { + let session_config = SessionConfig::new().with_target_partitions(4); // Make sure the configuration is the same if test is running on different machines + validate_query_with_memory_limits_and_config( + expected_mem_bytes, + mem_limit_bytes, + query, + baseline_query, + session_config, + None, + None, + ) + .await; +} + +/// Validate memory usage with a custom session configuration and optional +/// operator and spill assertions. +pub async fn validate_query_with_memory_limits_and_config( + expected_mem_bytes: i64, + mem_limit_bytes: Option, + query: &str, + baseline_query: &str, + session_config: SessionConfig, + expected_operator_name: Option<&str>, + expected_operator_spill: Option, ) { if std::env::var("DATAFUSION_TEST_MEM_LIMIT_VALIDATION").is_err() { println!("Skipping test because DATAFUSION_TEST_MEM_LIMIT_VALIDATION is not set"); @@ -151,18 +214,50 @@ pub async fn validate_query_with_memory_limits( None => runtime_builder.build_arc().unwrap(), }; - let session_config = SessionConfig::new().with_target_partitions(4); // Make sure the configuration is the same if test is running on different machines - let ctx = SessionContext::new_with_config_rt(session_config, runtime); let df = ctx.sql(query).await.unwrap(); + let physical_plan = df.create_physical_plan().await.unwrap(); + + if let Some(expected) = expected_operator_name { + let plan_display = displayable(physical_plan.as_ref()).indent(true).to_string(); + assert!( + plan_display.contains(expected), + "expected physical plan to contain `{expected}`, but got:\n{plan_display}", + ); + } + // Run a query with 10% data to estimate the constant overhead - let df_small = ctx.sql(baseline_query).await.unwrap(); + let baseline_plan = 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() + }) + .await; - let (_, baseline_max_rss) = - measure_max_rss(|| async { df_small.collect().await.unwrap() }).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() + }) + .await; - let (_, max_rss) = measure_max_rss(|| async { df.collect().await.unwrap() }).await; + if let (Some(operator), Some(expect_spill)) = + (expected_operator_name, expected_operator_spill) + { + let spill_count = operator_spill_count(physical_plan.as_ref(), operator); + assert_eq!( + spill_count > 0, + expect_spill, + "unexpected spill_count={spill_count} for {operator}", + ); + } println!( "Memory before: {}, Memory after: {}", diff --git a/datafusion/sqllogictest/test_files/sort_merge_join_spill.slt b/datafusion/sqllogictest/test_files/sort_merge_join_spill.slt new file mode 100644 index 0000000000000..1a3dcafa60f82 --- /dev/null +++ b/datafusion/sqllogictest/test_files/sort_merge_join_spill.slt @@ -0,0 +1,252 @@ +# 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. + +# End-to-end SortMergeJoinExec spilling tests. +# +# Each query runs as an unlimited-memory hash join for expected results, then as +# a memory-limited sort-merge join that must spill. + +hash-threshold 100 + +# Use multiple partitions so the planner can select SortMergeJoinExec. +statement ok +SET datafusion.execution.target_partitions = 2 + +statement ok +SET datafusion.execution.batch_size = 200 + +# Probe rows include one matching key; x=500 yields true, false, and NULL filters. +statement ok +CREATE VIEW probe AS +SELECT value AS k, 500 AS x FROM generate_series(1, 3); + +# Probe rows with no matching buffered key. +statement ok +CREATE VIEW probe_nomatch AS SELECT value AS k FROM generate_series(7, 9); + +# One 2,000-row key group with a 512-byte payload, split into 10 batches. +# Ordered generation avoids input sorts so only the join buffers the payload; +# x includes NULLs and values on both sides of 500. +statement ok +CREATE VIEW wide AS +SELECT 2 AS k, + value AS v, + CASE WHEN value % 10 = 0 THEN cast(NULL AS BIGINT) ELSE value % 1000 END AS x, + lpad(cast(value AS varchar), 512, 'x') AS p +FROM generate_series(1, 2000); + +# Keep output narrow while retaining the payload in the buffered input. + +query TT +EXPLAIN SELECT p.k, w.v, length(w.p) FROM probe p JOIN wide w ON p.k = w.k +---- +HashJoinExec + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p JOIN wide w ON p.k = w.k +---- +6000 values hashing to ae029ab21ba6942d04253c3eb1fafbee + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p LEFT JOIN wide w ON p.k = w.k +---- +6006 values hashing to 352109cc65a61f6224bb027dbee60df5 + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM wide w RIGHT JOIN probe p ON p.k = w.k +---- +6006 values hashing to 352109cc65a61f6224bb027dbee60df5 + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p FULL JOIN wide w ON p.k = w.k +---- +6006 values hashing to 352109cc65a61f6224bb027dbee60df5 + +# Use the unlimited-memory hash join as the reference for filtered joins. +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p +JOIN wide w ON p.k = w.k AND p.x < w.x +---- +2700 values hashing to 824832563a1e34fe419885d0b7cccc9d + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p +LEFT JOIN wide w ON p.k = w.k AND p.x < w.x +---- +2706 values hashing to 38c8a5628a41b463e41d592c2401ca8d + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM wide w +RIGHT JOIN probe p ON p.k = w.k AND p.x < w.x +---- +2706 values hashing to 38c8a5628a41b463e41d592c2401ca8d + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p +FULL JOIN wide w ON p.k = w.k AND p.x < w.x +---- +6006 values hashing to b6b875b2658ee19e8bca7cd6e993dee1 + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe_nomatch p FULL JOIN wide w ON p.k = w.k +---- +6009 values hashing to 126cd87356bc636448b65c5fb5f4bd2b + +# A 64 KB pool spills all 10 buffered batches; each result must match its +# unlimited-memory hash-join reference. + +statement ok +SET datafusion.optimizer.prefer_hash_join = false + +statement ok +SET datafusion.runtime.memory_limit = '64K' + +query TT +EXPLAIN ANALYZE +SELECT p.k, w.v, length(w.p) FROM probe p JOIN wide w ON p.k = w.k +---- +Plan with Metrics +SortMergeJoinExec: join_type=Inner, on=[(k@0, k@0)], metrics=[output_rows=2.00 K,spill_count=10, spilled_bytes=spilled_rows=2.00 K, peak_mem_used= + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p JOIN wide w ON p.k = w.k +---- +6000 values hashing to ae029ab21ba6942d04253c3eb1fafbee + +query TT +EXPLAIN ANALYZE +SELECT p.k, w.v, length(w.p) FROM probe p LEFT JOIN wide w ON p.k = w.k +---- +Plan with Metrics +SortMergeJoinExec: join_type=Left, on=[(k@0, k@0)], metrics=[output_rows=2.00 K,spill_count=10, spilled_bytes=spilled_rows=2.00 K, peak_mem_used= + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p LEFT JOIN wide w ON p.k = w.k +---- +6006 values hashing to 352109cc65a61f6224bb027dbee60df5 + +query TT +EXPLAIN ANALYZE +SELECT p.k, w.v, length(w.p) FROM wide w RIGHT JOIN probe p ON p.k = w.k +---- +Plan with Metrics +SortMergeJoinExec: join_type=Right, on=[(k@0, k@0)], metrics=[output_rows=2.00 K,spill_count=10, spilled_bytes=spilled_rows=2.00 K, peak_mem_used= + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM wide w RIGHT JOIN probe p ON p.k = w.k +---- +6006 values hashing to 352109cc65a61f6224bb027dbee60df5 + +query TT +EXPLAIN ANALYZE +SELECT p.k, w.v, length(w.p) FROM probe p FULL JOIN wide w ON p.k = w.k +---- +Plan with Metrics +SortMergeJoinExec: join_type=Full, on=[(k@0, k@0)], metrics=[output_rows=2.00 K,spill_count=10, spilled_bytes=spilled_rows=2.00 K, peak_mem_used= + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p FULL JOIN wide w ON p.k = w.k +---- +6006 values hashing to 352109cc65a61f6224bb027dbee60df5 + +# Filtered spills cover true, false, and NULL masks; outer joins defer unmatched +# rows until the whole key group is restored. + +query TT +EXPLAIN ANALYZE +SELECT p.k, w.v, length(w.p) FROM probe p +JOIN wide w ON p.k = w.k AND p.x < w.x +---- +Plan with Metrics +SortMergeJoinExec: join_type=Inner, on=[(k@0, k@0)], filter=x@0 < x@1, metrics=[output_rows=900,spill_count=10, spilled_bytes=spilled_rows=2.00 K, peak_mem_used= + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p +JOIN wide w ON p.k = w.k AND p.x < w.x +---- +2700 values hashing to 824832563a1e34fe419885d0b7cccc9d + +query TT +EXPLAIN ANALYZE +SELECT p.k, w.v, length(w.p) FROM probe p +LEFT JOIN wide w ON p.k = w.k AND p.x < w.x +---- +Plan with Metrics +SortMergeJoinExec: join_type=Left, on=[(k@0, k@0)], filter=x@0 < x@1, metrics=[output_rows=902,spill_count=10, spilled_bytes=spilled_rows=2.00 K, peak_mem_used= + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p +LEFT JOIN wide w ON p.k = w.k AND p.x < w.x +---- +2706 values hashing to 38c8a5628a41b463e41d592c2401ca8d + +query TT +EXPLAIN ANALYZE +SELECT p.k, w.v, length(w.p) FROM wide w +RIGHT JOIN probe p ON p.k = w.k AND p.x < w.x +---- +Plan with Metrics +SortMergeJoinExec: join_type=Right, on=[(k@0, k@0)], filter=x@1 < x@0, metrics=[output_rows=902,spill_count=10, spilled_bytes=spilled_rows=2.00 K, peak_mem_used= + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM wide w +RIGHT JOIN probe p ON p.k = w.k AND p.x < w.x +---- +2706 values hashing to 38c8a5628a41b463e41d592c2401ca8d + +query TT +EXPLAIN ANALYZE +SELECT p.k, w.v, length(w.p) FROM probe p +FULL JOIN wide w ON p.k = w.k AND p.x < w.x +---- +Plan with Metrics +SortMergeJoinExec: join_type=Full, on=[(k@0, k@0)], filter=x@0 < x@1, metrics=[output_rows=2.00 K,spill_count=10, spilled_bytes=spilled_rows=2.00 K, peak_mem_used= + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p +FULL JOIN wide w ON p.k = w.k AND p.x < w.x +---- +6006 values hashing to b6b875b2658ee19e8bca7cd6e993dee1 + +# Full join restores all buffered batches to emit unmatched rows. + +query TT +EXPLAIN ANALYZE +SELECT p.k, w.v, length(w.p) FROM probe_nomatch p FULL JOIN wide w ON p.k = w.k +---- +Plan with Metrics +SortMergeJoinExec: join_type=Full, on=[(k@0, k@0)], metrics=[output_rows=2.00 K,spill_count=10, spilled_bytes=spilled_rows=2.00 K, peak_mem_used= + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe_nomatch p FULL JOIN wide w ON p.k = w.k +---- +6009 values hashing to 126cd87356bc636448b65c5fb5f4bd2b + +statement ok +RESET datafusion.runtime.memory_limit + +statement ok +RESET datafusion.optimizer.prefer_hash_join + +statement ok +RESET datafusion.execution.batch_size + +statement ok +SET datafusion.execution.target_partitions = 4 + +statement ok +RESET datafusion.catalog.create_default_catalog_and_schema