minor(test): strengthen sort-merge join spilling coverage

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;

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;
+}

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");
 }
 
 // ===========================================================================

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::<usize>()
+}
+
 /// 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<i64>,
     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<i64>,
+    query: &str,
+    baseline_query: &str,
+    session_config: SessionConfig,
+    expected_operator_name: Option<&str>,
+    expected_operator_spill: Option<bool>,
 ) {
     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: {}",