From 6cb6d2a2b226d676cfefda681c370658b91e91c0 Mon Sep 17 00:00:00 2001 From: buraksenn Date: Tue, 28 Jul 2026 21:53:23 +0300 Subject: [PATCH 1/2] add smj spill test coverage --- datafusion/core/tests/fuzz_cases/join_fuzz.rs | 258 ++++++++++-------- ...spilling_fuzz_in_memory_constrained_env.rs | 228 +++++++++++++++- .../memory_limit_validation/mod.rs | 1 + .../smj_mem_validation.rs | 125 +++++++++ .../sort_mem_validation.rs | 48 +--- .../memory_limit_validation/utils.rs | 111 +++++++- 6 files changed, 613 insertions(+), 158 deletions(-) create mode 100644 datafusion/core/tests/memory_limit/memory_limit_validation/smj_mem_validation.rs diff --git a/datafusion/core/tests/fuzz_cases/join_fuzz.rs b/datafusion/core/tests/fuzz_cases/join_fuzz.rs index 81c7c9f83928e..eb4c26351ab85 100644 --- a/datafusion/core/tests/fuzz_cases/join_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/join_fuzz.rs @@ -20,7 +20,7 @@ use std::time::SystemTime; use crate::fuzz_cases::join_fuzz::JoinTestType::{HjSmj, NljHj}; -use arrow::array::{ArrayRef, BinaryArray, Int32Array}; +use arrow::array::{ArrayRef, BinaryArray, Int32Array, StringArray}; use arrow::compute::SortOptions; use arrow::datatypes::Schema; use arrow::record_batch::RecordBatch; @@ -30,22 +30,24 @@ use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::logical_expr::{JoinType, Operator}; use datafusion::physical_expr::expressions::BinaryExpr; -use datafusion::physical_plan::collect; use datafusion::physical_plan::expressions::Column; use datafusion::physical_plan::joins::utils::{ColumnIndex, JoinFilter}; use datafusion::physical_plan::joins::{ HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, }; +use datafusion::physical_plan::{ExecutionPlan, collect}; use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_common::{NullEquality, ScalarValue}; use datafusion_execution::TaskContext; use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; +use datafusion_execution::memory_pool::FairSpillPool; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_physical_expr::PhysicalExprRef; use datafusion_physical_expr::expressions::Literal; use itertools::Itertools; -use rand::Rng; +use rand::rngs::SmallRng; +use rand::{Rng, SeedableRng}; use test_utils::stagger_batch_with_seed; // Determines what Fuzz tests needs to run @@ -1126,133 +1128,131 @@ impl JoinFuzzTestCase { } } -/// Fuzz test: compare SMJ (with spilling) against HJ (no spill) for filtered -/// outer joins under memory pressure. This exercises the deferred filtering + -/// spill read-back path that unit tests can't easily cover with random data. +/// Compare a guaranteed-spilling SMJ against an unlimited-memory hash join +/// for filtered materializing joins. #[tokio::test] async fn test_filtered_join_spill_fuzz() { - let join_types = [JoinType::Left, JoinType::Right, JoinType::Full]; + let join_types = [ + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + ]; + let input1 = make_spill_join_batches(256, 32, 512, 1); + let input2 = make_spill_join_batches(256, 32, 512, 2); + let schema1 = input1[0].schema(); + let schema2 = input2[0].schema(); + let filter = col_lt_col_filter(Arc::clone(&schema1), Arc::clone(&schema2)); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a", &schema1).unwrap()) as _, + Arc::new(Column::new_with_schema("a", &schema2).unwrap()) as _, + ), + ( + Arc::new(Column::new_with_schema("b", &schema1).unwrap()) as _, + Arc::new(Column::new_with_schema("b", &schema2).unwrap()) as _, + ), + ]; let runtime_spill = RuntimeEnvBuilder::new() - .with_memory_limit(4096, 1.0) + .with_memory_pool(Arc::new(FairSpillPool::new(1024))) .with_disk_manager_builder( DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), ) .build_arc() .unwrap(); - for join_type in &join_types { - for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] { - let input1 = make_staggered_batches_i32(1000, left_extra); - let input2 = make_staggered_batches_i32(1000, right_extra); + for join_type in join_types { + for batch_size in [2, 50] { + let session_config = SessionConfig::new().with_batch_size(batch_size); - let schema1 = input1[0].schema(); - let schema2 = input2[0].schema(); - let filter = col_lt_col_filter(schema1.clone(), schema2.clone()); - - let on = vec![ - ( - Arc::new(Column::new_with_schema("a", &schema1).unwrap()) as _, - Arc::new(Column::new_with_schema("a", &schema2).unwrap()) as _, - ), - ( - Arc::new(Column::new_with_schema("b", &schema1).unwrap()) as _, - Arc::new(Column::new_with_schema("b", &schema2).unwrap()) as _, - ), - ]; - - for batch_size in [2, 49, 100] { - let session_config = SessionConfig::new().with_batch_size(batch_size); - - // HJ baseline (no memory limit) - let left_hj = MemorySourceConfig::try_new_exec( - std::slice::from_ref(&input1), - schema1.clone(), - None, - ) - .unwrap(); - let right_hj = MemorySourceConfig::try_new_exec( - std::slice::from_ref(&input2), - schema2.clone(), + let left_hj = MemorySourceConfig::try_new_exec( + std::slice::from_ref(&input1), + Arc::clone(&schema1), + None, + ) + .unwrap(); + let right_hj = MemorySourceConfig::try_new_exec( + std::slice::from_ref(&input2), + Arc::clone(&schema2), + None, + ) + .unwrap(); + let hj = Arc::new( + HashJoinExec::try_new( + left_hj, + right_hj, + on.clone(), + Some(filter.clone()), + &join_type, None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, ) - .unwrap(); - let hj = Arc::new( - HashJoinExec::try_new( - left_hj, - right_hj, - on.clone(), - Some(filter.clone()), - join_type, - None, - PartitionMode::Partitioned, - NullEquality::NullEqualsNothing, - false, - ) - .unwrap(), - ); - let ctx_hj = SessionContext::new_with_config(session_config.clone()); - let hj_collected = collect(hj, ctx_hj.task_ctx()).await.unwrap(); + .unwrap(), + ); + let ctx_hj = SessionContext::new_with_config(session_config.clone()); + let hj_collected = collect(hj, ctx_hj.task_ctx()).await.unwrap(); - // SMJ with spilling - let left_smj = MemorySourceConfig::try_new_exec( - std::slice::from_ref(&input1), - schema1.clone(), - None, - ) - .unwrap(); - let right_smj = MemorySourceConfig::try_new_exec( - std::slice::from_ref(&input2), - schema2.clone(), - None, + let left_smj = MemorySourceConfig::try_new_exec( + std::slice::from_ref(&input1), + Arc::clone(&schema1), + None, + ) + .unwrap(); + let right_smj = MemorySourceConfig::try_new_exec( + std::slice::from_ref(&input2), + Arc::clone(&schema2), + None, + ) + .unwrap(); + let smj = Arc::new( + SortMergeJoinExec::try_new( + left_smj, + right_smj, + on.clone(), + Some(filter.clone()), + join_type, + vec![SortOptions::default(); on.len()], + NullEquality::NullEqualsNothing, ) - .unwrap(); - let smj = Arc::new( - SortMergeJoinExec::try_new( - left_smj, - right_smj, - on.clone(), - Some(filter.clone()), - *join_type, - vec![SortOptions::default(); on.len()], - NullEquality::NullEqualsNothing, - ) - .unwrap(), - ); - let task_ctx_spill = Arc::new( - TaskContext::default() - .with_session_config(session_config) - .with_runtime(Arc::clone(&runtime_spill)), - ); - let smj_collected = collect(smj, task_ctx_spill).await.unwrap(); + .unwrap(), + ); + let task_ctx_spill = Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(Arc::clone(&runtime_spill)), + ); + let smj_collected = + collect(Arc::clone(&smj) as Arc, task_ctx_spill) + .await + .unwrap(); + + assert!( + smj.metrics().unwrap().spill_count().unwrap_or_default() > 0, + "expected SMJ to spill for {join_type:?} batch_size={batch_size}", + ); - let hj_rows: usize = hj_collected.iter().map(|b| b.num_rows()).sum(); - let smj_rows: usize = smj_collected.iter().map(|b| b.num_rows()).sum(); + let hj_rows: usize = hj_collected.iter().map(|b| b.num_rows()).sum(); + let smj_rows: usize = smj_collected.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + hj_rows, smj_rows, + "row count mismatch for {join_type:?} batch_size={batch_size}: \ + HJ={hj_rows} SMJ={smj_rows}", + ); + if hj_rows > 0 { + let hj_fmt = pretty_format_batches(&hj_collected).unwrap().to_string(); + let smj_fmt = pretty_format_batches(&smj_collected).unwrap().to_string(); + let mut hj_sorted: Vec<&str> = hj_fmt.trim().lines().collect(); + hj_sorted.sort_unstable(); + let mut smj_sorted: Vec<&str> = smj_fmt.trim().lines().collect(); + smj_sorted.sort_unstable(); assert_eq!( - hj_rows, smj_rows, - "Row count mismatch for {join_type:?} batch_size={batch_size} \ - left_extra={left_extra} right_extra={right_extra}: \ - HJ={hj_rows} SMJ={smj_rows}" + hj_sorted, smj_sorted, + "content mismatch for {join_type:?} batch_size={batch_size}", ); - - if hj_rows > 0 { - let hj_fmt = - pretty_format_batches(&hj_collected).unwrap().to_string(); - let smj_fmt = - pretty_format_batches(&smj_collected).unwrap().to_string(); - - let mut hj_sorted: Vec<&str> = hj_fmt.trim().lines().collect(); - hj_sorted.sort_unstable(); - let mut smj_sorted: Vec<&str> = smj_fmt.trim().lines().collect(); - smj_sorted.sort_unstable(); - - assert_eq!( - hj_sorted, smj_sorted, - "Content mismatch for {join_type:?} batch_size={batch_size} \ - left_extra={left_extra} right_extra={right_extra}" - ); - } } } } @@ -1347,3 +1347,39 @@ fn make_staggered_batches_binary( // preserve your existing randomized partitioning stagger_batch_with_seed(batch, 42) } + +/// Sorted, low-cardinality inputs whose wide payloads force SMJ key-group +/// spilling. `(a, b)` is sorted (`b` is constant) and `x` is nullable to +/// exercise filter NULL handling. +fn make_spill_join_batches( + len: usize, + num_keys: i32, + payload_len: usize, + seed: u64, +) -> Vec { + let mut rng = SmallRng::seed_from_u64(seed); + + let mut keys: Vec = (0..len).map(|_| rng.random_range(0..num_keys)).collect(); + keys.sort_unstable(); + let a = Int32Array::from_iter_values(keys); + let b = Int32Array::from_iter_values(std::iter::repeat_n(0, len)); + let x = Int32Array::from_iter((0..len).map(|_| { + if rng.random_range(0..10) == 0 { + None + } else { + Some(rng.random_range(0..1000)) + } + })); + + let payload = "a".repeat(payload_len); + let p = StringArray::from_iter_values(std::iter::repeat_n(payload.as_str(), len)); + let batch = RecordBatch::try_from_iter(vec![ + ("a", Arc::new(a) as ArrayRef), + ("b", Arc::new(b) as ArrayRef), + ("x", Arc::new(x) as ArrayRef), + ("p", Arc::new(p) as ArrayRef), + ]) + .unwrap(); + + stagger_batch_with_seed(batch, 7) +} diff --git a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs index 103c3e03c06df..eddd9727cb35c 100644 --- a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs +++ b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs @@ -22,16 +22,18 @@ use std::sync::Arc; use crate::fuzz_cases::aggregate_fuzz::assert_spill_count_metric; use crate::fuzz_cases::once_exec::OnceExec; -use arrow::array::UInt64Array; +use arrow::array::{Array, Int64Array, UInt64Array}; use arrow::{array::StringArray, compute::SortOptions, record_batch::RecordBatch}; -use arrow_schema::{DataType, Field, Schema}; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; use datafusion::common::Result; use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::logical_expr::JoinType; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::expressions::PhysicalSortExpr; use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::prelude::SessionConfig; use datafusion_common::units::{KB, MB}; +use datafusion_common::{JoinSide, NullEquality}; use datafusion_execution::memory_pool::{ FairSpillPool, MemoryConsumer, MemoryReservation, }; @@ -44,6 +46,7 @@ use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, }; +use datafusion_physical_plan::joins::SortMergeJoinExec; use datafusion_physical_plan::metrics::MetricValue; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use futures::StreamExt; @@ -792,3 +795,224 @@ fn get_output_batches_from_metrics(metrics: &MetricsSet) -> usize { }) .expect("Must have output_batches metric since it exists in the baseline") } + +/// Sorted-by-key SMJ input: keys `0..num_keys`, each repeated `rows_per_key` +/// times (SMJ does not sort its inputs), with a `payload_len`-byte payload per +/// row. A key group spans batches when `rows_per_key > batch_size`, which is +/// what forces the buffered side to spill. +fn make_sorted_join_input( + schema: &SchemaRef, + num_keys: i64, + rows_per_key: usize, + payload_len: usize, + batch_size: usize, +) -> Vec { + let payload = "a".repeat(payload_len); + let total = num_keys as usize * rows_per_key; + let mut keys = Vec::with_capacity(total); + for k in 0..num_keys { + for _ in 0..rows_per_key { + keys.push(k); + } + } + + keys.chunks(batch_size) + .map(|chunk| { + let key_arr = Int64Array::from_iter_values(chunk.iter().copied()); + let val_arr = StringArray::from_iter_values(std::iter::repeat_n( + payload.as_str(), + chunk.len(), + )); + RecordBatch::try_new( + Arc::clone(schema), + vec![Arc::new(key_arr), Arc::new(val_arr)], + ) + .expect("valid record batch") + }) + .collect() +} + +/// Run one `SortMergeJoinExec` whose buffered key groups (~3 MB each) overflow +/// a 2 MB pool, asserting it spills and still produces the full, correct +/// output. The streamed side has one row per key, bounding the output to +/// `num_keys * rows_per_key_large` rows. +async fn run_sort_merge_join_test_with_limited_memory(join_type: JoinType) -> Result<()> { + let pool_size = 2 * MB as usize; + let batch_size = 1000; + let num_keys: i64 = 12; + let rows_per_key_large = 6000; + let payload_len_large = 512; + let payload_len_sparse = 16; + + let left_schema = Arc::new(Schema::new(vec![ + Field::new("k1", DataType::Int64, false), + Field::new("v1", DataType::Utf8, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("k2", DataType::Int64, false), + Field::new("v2", DataType::Utf8, false), + ])); + + // SMJ spills the buffered side — the opposite of the probe/streamed side. + let buffered_side = SortMergeJoinExec::probe_side(&join_type).negate(); + let large_left = matches!(buffered_side, JoinSide::Left); + let large_payload = "a".repeat(payload_len_large); + let sparse_payload = "a".repeat(payload_len_sparse); + let (expected_left_payload, expected_right_payload) = if large_left { + (large_payload.as_str(), sparse_payload.as_str()) + } else { + (sparse_payload.as_str(), large_payload.as_str()) + }; + + let left_batches = make_sorted_join_input( + &left_schema, + num_keys, + if large_left { rows_per_key_large } else { 1 }, + if large_left { + payload_len_large + } else { + payload_len_sparse + }, + batch_size, + ); + let right_batches = make_sorted_join_input( + &right_schema, + num_keys, + if large_left { 1 } else { rows_per_key_large }, + if large_left { + payload_len_sparse + } else { + payload_len_large + }, + batch_size, + ); + + // No unmatched rows on either side, so all four join types produce the + // same output row count. + let expected_output_rows = num_keys as usize * rows_per_key_large; + + let task_ctx = { + let memory_pool = Arc::new(FairSpillPool::new(pool_size)); + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(batch_size)) + .with_runtime(Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(memory_pool) + .build()?, + )) + }; + let task_ctx = Arc::new(task_ctx); + + let left: Arc = + Arc::new(OnceExec::new(Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&left_schema), + futures::stream::iter(left_batches.into_iter().map(Ok)), + )))); + let right: Arc = + Arc::new(OnceExec::new(Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&right_schema), + futures::stream::iter(right_batches.into_iter().map(Ok)), + )))); + + let on = vec![( + Arc::new(Column::new_with_schema("k1", &left_schema)?) as _, + Arc::new(Column::new_with_schema("k2", &right_schema)?) as _, + )]; + + let join = Arc::new(SortMergeJoinExec::try_new( + left, + right, + on, + None, + join_type, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + )?); + + let mut stream = join.execute(0, Arc::clone(&task_ctx))?; + let mut output_rows = 0; + let mut output_rows_per_key = vec![0; num_keys as usize]; + while let Some(batch) = stream.next().await { + let batch = batch?; + let left_keys = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("left key must be Int64"); + let left_values = batch + .column(1) + .as_any() + .downcast_ref::() + .expect("left payload must be Utf8"); + let right_keys = batch + .column(2) + .as_any() + .downcast_ref::() + .expect("right key must be Int64"); + let right_values = batch + .column(3) + .as_any() + .downcast_ref::() + .expect("right payload must be Utf8"); + + assert_eq!(left_keys.null_count(), 0); + assert_eq!(left_values.null_count(), 0); + assert_eq!(right_keys.null_count(), 0); + assert_eq!(right_values.null_count(), 0); + for row in 0..batch.num_rows() { + assert_eq!(left_keys.value(row), right_keys.value(row)); + let key = left_keys.value(row); + assert!( + (0..num_keys).contains(&key), + "unexpected output key {key} for {join_type:?}", + ); + output_rows_per_key[key as usize] += 1; + assert_eq!(left_values.value(row), expected_left_payload); + assert_eq!(right_values.value(row), expected_right_payload); + } + output_rows += batch.num_rows(); + } + + assert_eq!( + output_rows, expected_output_rows, + "unexpected output row count for {join_type:?}", + ); + for (key, count) in output_rows_per_key.into_iter().enumerate() { + assert_eq!( + count, rows_per_key_large, + "unexpected output count for key {key} in {join_type:?}", + ); + } + + let metrics = join.metrics().expect("must have metrics"); + assert!( + metrics.spill_count().unwrap_or_default() > 0, + "expected the buffered side to spill for {join_type:?}", + ); + assert!( + metrics.spilled_bytes().unwrap_or_default() > 0, + "expected spilled_bytes > 0 for {join_type:?}", + ); + assert!( + metrics.spilled_rows().unwrap_or_default() > 0, + "expected spilled_rows > 0 for {join_type:?}", + ); + + Ok(()) +} + +/// Memory-limit stress test: ~36 MB of buffered data joined under a 2 MB pool +/// for each materializing join type. Semi/anti/mark joins take a separate +/// execution path (`BitwiseSortMergeJoinStream`) and are out of scope here. +#[tokio::test] +async fn test_sort_merge_join_with_limited_memory() -> Result<()> { + for join_type in [ + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + ] { + run_sort_merge_join_test_with_limited_memory(join_type).await?; + } + Ok(()) +} 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..4034dbe0828e0 --- /dev/null +++ b/datafusion/core/tests/memory_limit/memory_limit_validation/smj_mem_validation.rs @@ -0,0 +1,125 @@ +// 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. + +//! Validates that the physical (RSS) memory usage of sort-merge join queries +//! stays within the configured memory limit. Tests must run in separate +//! processes for accurate measurement, so runners spawn each test case as its +//! own process. + +use datafusion::prelude::SessionConfig; + +use crate::memory_limit::memory_limit_validation::utils; + +/// Guards against validating a hash join the planner picked instead. +const SMJ_OPERATOR_NAME: &str = "SortMergeJoinExec"; + +/// Force the planner to pick a sort-merge join and let it spill under a tight +/// pool: SMJ is only chosen with hash joins disabled and partitions > 1 (kept +/// at the minimum of 2 since each partition multiplies the per-sorter merge +/// reservation), and `sort_spill_reservation_bytes` is shrunk from its 10MB +/// default, which would otherwise exhaust a small pool before any data spills. +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) +} + +/// Equi-join with one large right-side key group that overflows the configured +/// pools; `sum(v)` keeps the output to one row. +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 runners: +// Runners are split into multiple tests to run in parallel +// =========================================================================== + +#[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"); +} + +// =========================================================================== +// Test cases: +// All following tests need to be run through their individual test wrapper. +// When run directly, environment variable `DATAFUSION_TEST_MEM_LIMIT_VALIDATION` +// is not set, test will return with a no-op. +// +// If some tests consistently fail, suppress by setting a larger expected memory +// usage (e.g. 40_000_000 * 4 -> 40_000_000 * 5) +// =========================================================================== + +/// 40MB limit against one ~80MB buffered key group: the join must spill. +#[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), // Baseline query with ~10% of data + smj_session_config(), + Some(SMJ_OPERATOR_NAME), + Some(true), + ) + .await; +} + +/// Tighter 16MB limit forces more aggressive spilling. +#[tokio::test] +async fn smj_with_mem_limit_2() { + utils::validate_query_with_memory_limits_and_config( + 16_000_000 * 5, + Some(16_000_000), + &smj_sum_query(5_000_000), + &smj_sum_query(500_000), // Baseline query with ~10% of data + smj_session_config(), + Some(SMJ_OPERATOR_NAME), + Some(true), + ) + .await; +} + +/// No memory limit: the join must not spill. +#[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), // Baseline query with ~10% of data + 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..2cec1f9bae0a8 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,41 @@ where (result, peak_rss) } +/// Re-run one test from the current integration-test binary in an isolated +/// process so its RSS can be measured independently. +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 +170,33 @@ pub async fn validate_query_with_memory_limits( mem_limit_bytes: Option, query: &str, baseline_query: &str, +) { + // Fixed partition count so results are stable across machines. + let session_config = SessionConfig::new().with_target_partitions(4); + validate_query_with_memory_limits_and_config( + expected_mem_bytes, + mem_limit_bytes, + query, + baseline_query, + session_config, + None, + None, + ) + .await; +} + +/// Like [`validate_query_with_memory_limits`] but with an explicit +/// [`SessionConfig`] and optional operator expectations: the plan must contain +/// `expected_operator_name`, and `expected_operator_spill` requires its spill +/// count to be positive (`true`) or zero (`false`). +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 +216,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(); - // Run a query with 10% data to estimate the constant overhead - let df_small = ctx.sql(baseline_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}", + ); + } - let (_, baseline_max_rss) = - measure_max_rss(|| async { df_small.collect().await.unwrap() }).await; + // Run a query with 10% data to estimate the constant overhead. + 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 (_, max_rss) = measure_max_rss(|| async { df.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; + + 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: {}", From f45a5cd04ab080b83be244cf4b4454792b352723 Mon Sep 17 00:00:00 2001 From: buraksenn Date: Wed, 29 Jul 2026 11:47:35 +0300 Subject: [PATCH 2/2] move to slt tests as much as possible --- datafusion/core/tests/fuzz_cases/join_fuzz.rs | 258 ++++++++---------- ...spilling_fuzz_in_memory_constrained_env.rs | 228 +--------------- .../smj_mem_validation.rs | 48 +--- .../memory_limit_validation/utils.rs | 16 +- .../test_files/sort_merge_join_spill.slt | 252 +++++++++++++++++ 5 files changed, 386 insertions(+), 416 deletions(-) create mode 100644 datafusion/sqllogictest/test_files/sort_merge_join_spill.slt diff --git a/datafusion/core/tests/fuzz_cases/join_fuzz.rs b/datafusion/core/tests/fuzz_cases/join_fuzz.rs index eb4c26351ab85..81c7c9f83928e 100644 --- a/datafusion/core/tests/fuzz_cases/join_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/join_fuzz.rs @@ -20,7 +20,7 @@ use std::time::SystemTime; use crate::fuzz_cases::join_fuzz::JoinTestType::{HjSmj, NljHj}; -use arrow::array::{ArrayRef, BinaryArray, Int32Array, StringArray}; +use arrow::array::{ArrayRef, BinaryArray, Int32Array}; use arrow::compute::SortOptions; use arrow::datatypes::Schema; use arrow::record_batch::RecordBatch; @@ -30,24 +30,22 @@ use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::logical_expr::{JoinType, Operator}; use datafusion::physical_expr::expressions::BinaryExpr; +use datafusion::physical_plan::collect; use datafusion::physical_plan::expressions::Column; use datafusion::physical_plan::joins::utils::{ColumnIndex, JoinFilter}; use datafusion::physical_plan::joins::{ HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, }; -use datafusion::physical_plan::{ExecutionPlan, collect}; use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_common::{NullEquality, ScalarValue}; use datafusion_execution::TaskContext; use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; -use datafusion_execution::memory_pool::FairSpillPool; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_physical_expr::PhysicalExprRef; use datafusion_physical_expr::expressions::Literal; use itertools::Itertools; -use rand::rngs::SmallRng; -use rand::{Rng, SeedableRng}; +use rand::Rng; use test_utils::stagger_batch_with_seed; // Determines what Fuzz tests needs to run @@ -1128,131 +1126,133 @@ impl JoinFuzzTestCase { } } -/// Compare a guaranteed-spilling SMJ against an unlimited-memory hash join -/// for filtered materializing joins. +/// Fuzz test: compare SMJ (with spilling) against HJ (no spill) for filtered +/// outer joins under memory pressure. This exercises the deferred filtering + +/// spill read-back path that unit tests can't easily cover with random data. #[tokio::test] async fn test_filtered_join_spill_fuzz() { - let join_types = [ - JoinType::Inner, - JoinType::Left, - JoinType::Right, - JoinType::Full, - ]; - let input1 = make_spill_join_batches(256, 32, 512, 1); - let input2 = make_spill_join_batches(256, 32, 512, 2); - let schema1 = input1[0].schema(); - let schema2 = input2[0].schema(); - let filter = col_lt_col_filter(Arc::clone(&schema1), Arc::clone(&schema2)); - let on = vec![ - ( - Arc::new(Column::new_with_schema("a", &schema1).unwrap()) as _, - Arc::new(Column::new_with_schema("a", &schema2).unwrap()) as _, - ), - ( - Arc::new(Column::new_with_schema("b", &schema1).unwrap()) as _, - Arc::new(Column::new_with_schema("b", &schema2).unwrap()) as _, - ), - ]; + let join_types = [JoinType::Left, JoinType::Right, JoinType::Full]; let runtime_spill = RuntimeEnvBuilder::new() - .with_memory_pool(Arc::new(FairSpillPool::new(1024))) + .with_memory_limit(4096, 1.0) .with_disk_manager_builder( DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), ) .build_arc() .unwrap(); - for join_type in join_types { - for batch_size in [2, 50] { - let session_config = SessionConfig::new().with_batch_size(batch_size); + for join_type in &join_types { + for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] { + let input1 = make_staggered_batches_i32(1000, left_extra); + let input2 = make_staggered_batches_i32(1000, right_extra); - let left_hj = MemorySourceConfig::try_new_exec( - std::slice::from_ref(&input1), - Arc::clone(&schema1), - None, - ) - .unwrap(); - let right_hj = MemorySourceConfig::try_new_exec( - std::slice::from_ref(&input2), - Arc::clone(&schema2), - None, - ) - .unwrap(); - let hj = Arc::new( - HashJoinExec::try_new( - left_hj, - right_hj, - on.clone(), - Some(filter.clone()), - &join_type, + let schema1 = input1[0].schema(); + let schema2 = input2[0].schema(); + let filter = col_lt_col_filter(schema1.clone(), schema2.clone()); + + let on = vec![ + ( + Arc::new(Column::new_with_schema("a", &schema1).unwrap()) as _, + Arc::new(Column::new_with_schema("a", &schema2).unwrap()) as _, + ), + ( + Arc::new(Column::new_with_schema("b", &schema1).unwrap()) as _, + Arc::new(Column::new_with_schema("b", &schema2).unwrap()) as _, + ), + ]; + + for batch_size in [2, 49, 100] { + let session_config = SessionConfig::new().with_batch_size(batch_size); + + // HJ baseline (no memory limit) + let left_hj = MemorySourceConfig::try_new_exec( + std::slice::from_ref(&input1), + schema1.clone(), None, - PartitionMode::Partitioned, - NullEquality::NullEqualsNothing, - false, ) - .unwrap(), - ); - let ctx_hj = SessionContext::new_with_config(session_config.clone()); - let hj_collected = collect(hj, ctx_hj.task_ctx()).await.unwrap(); + .unwrap(); + let right_hj = MemorySourceConfig::try_new_exec( + std::slice::from_ref(&input2), + schema2.clone(), + None, + ) + .unwrap(); + let hj = Arc::new( + HashJoinExec::try_new( + left_hj, + right_hj, + on.clone(), + Some(filter.clone()), + join_type, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ); + let ctx_hj = SessionContext::new_with_config(session_config.clone()); + let hj_collected = collect(hj, ctx_hj.task_ctx()).await.unwrap(); - let left_smj = MemorySourceConfig::try_new_exec( - std::slice::from_ref(&input1), - Arc::clone(&schema1), - None, - ) - .unwrap(); - let right_smj = MemorySourceConfig::try_new_exec( - std::slice::from_ref(&input2), - Arc::clone(&schema2), - None, - ) - .unwrap(); - let smj = Arc::new( - SortMergeJoinExec::try_new( - left_smj, - right_smj, - on.clone(), - Some(filter.clone()), - join_type, - vec![SortOptions::default(); on.len()], - NullEquality::NullEqualsNothing, + // SMJ with spilling + let left_smj = MemorySourceConfig::try_new_exec( + std::slice::from_ref(&input1), + schema1.clone(), + None, ) - .unwrap(), - ); - let task_ctx_spill = Arc::new( - TaskContext::default() - .with_session_config(session_config) - .with_runtime(Arc::clone(&runtime_spill)), - ); - let smj_collected = - collect(Arc::clone(&smj) as Arc, task_ctx_spill) - .await - .unwrap(); - - assert!( - smj.metrics().unwrap().spill_count().unwrap_or_default() > 0, - "expected SMJ to spill for {join_type:?} batch_size={batch_size}", - ); + .unwrap(); + let right_smj = MemorySourceConfig::try_new_exec( + std::slice::from_ref(&input2), + schema2.clone(), + None, + ) + .unwrap(); + let smj = Arc::new( + SortMergeJoinExec::try_new( + left_smj, + right_smj, + on.clone(), + Some(filter.clone()), + *join_type, + vec![SortOptions::default(); on.len()], + NullEquality::NullEqualsNothing, + ) + .unwrap(), + ); + let task_ctx_spill = Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(Arc::clone(&runtime_spill)), + ); + let smj_collected = collect(smj, task_ctx_spill).await.unwrap(); - let hj_rows: usize = hj_collected.iter().map(|b| b.num_rows()).sum(); - let smj_rows: usize = smj_collected.iter().map(|b| b.num_rows()).sum(); - assert_eq!( - hj_rows, smj_rows, - "row count mismatch for {join_type:?} batch_size={batch_size}: \ - HJ={hj_rows} SMJ={smj_rows}", - ); + let hj_rows: usize = hj_collected.iter().map(|b| b.num_rows()).sum(); + let smj_rows: usize = smj_collected.iter().map(|b| b.num_rows()).sum(); - if hj_rows > 0 { - let hj_fmt = pretty_format_batches(&hj_collected).unwrap().to_string(); - let smj_fmt = pretty_format_batches(&smj_collected).unwrap().to_string(); - let mut hj_sorted: Vec<&str> = hj_fmt.trim().lines().collect(); - hj_sorted.sort_unstable(); - let mut smj_sorted: Vec<&str> = smj_fmt.trim().lines().collect(); - smj_sorted.sort_unstable(); assert_eq!( - hj_sorted, smj_sorted, - "content mismatch for {join_type:?} batch_size={batch_size}", + hj_rows, smj_rows, + "Row count mismatch for {join_type:?} batch_size={batch_size} \ + left_extra={left_extra} right_extra={right_extra}: \ + HJ={hj_rows} SMJ={smj_rows}" ); + + if hj_rows > 0 { + let hj_fmt = + pretty_format_batches(&hj_collected).unwrap().to_string(); + let smj_fmt = + pretty_format_batches(&smj_collected).unwrap().to_string(); + + let mut hj_sorted: Vec<&str> = hj_fmt.trim().lines().collect(); + hj_sorted.sort_unstable(); + let mut smj_sorted: Vec<&str> = smj_fmt.trim().lines().collect(); + smj_sorted.sort_unstable(); + + assert_eq!( + hj_sorted, smj_sorted, + "Content mismatch for {join_type:?} batch_size={batch_size} \ + left_extra={left_extra} right_extra={right_extra}" + ); + } } } } @@ -1347,39 +1347,3 @@ fn make_staggered_batches_binary( // preserve your existing randomized partitioning stagger_batch_with_seed(batch, 42) } - -/// Sorted, low-cardinality inputs whose wide payloads force SMJ key-group -/// spilling. `(a, b)` is sorted (`b` is constant) and `x` is nullable to -/// exercise filter NULL handling. -fn make_spill_join_batches( - len: usize, - num_keys: i32, - payload_len: usize, - seed: u64, -) -> Vec { - let mut rng = SmallRng::seed_from_u64(seed); - - let mut keys: Vec = (0..len).map(|_| rng.random_range(0..num_keys)).collect(); - keys.sort_unstable(); - let a = Int32Array::from_iter_values(keys); - let b = Int32Array::from_iter_values(std::iter::repeat_n(0, len)); - let x = Int32Array::from_iter((0..len).map(|_| { - if rng.random_range(0..10) == 0 { - None - } else { - Some(rng.random_range(0..1000)) - } - })); - - let payload = "a".repeat(payload_len); - let p = StringArray::from_iter_values(std::iter::repeat_n(payload.as_str(), len)); - let batch = RecordBatch::try_from_iter(vec![ - ("a", Arc::new(a) as ArrayRef), - ("b", Arc::new(b) as ArrayRef), - ("x", Arc::new(x) as ArrayRef), - ("p", Arc::new(p) as ArrayRef), - ]) - .unwrap(); - - stagger_batch_with_seed(batch, 7) -} diff --git a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs index eddd9727cb35c..103c3e03c06df 100644 --- a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs +++ b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs @@ -22,18 +22,16 @@ use std::sync::Arc; use crate::fuzz_cases::aggregate_fuzz::assert_spill_count_metric; use crate::fuzz_cases::once_exec::OnceExec; -use arrow::array::{Array, Int64Array, UInt64Array}; +use arrow::array::UInt64Array; use arrow::{array::StringArray, compute::SortOptions, record_batch::RecordBatch}; -use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use arrow_schema::{DataType, Field, Schema}; use datafusion::common::Result; use datafusion::execution::runtime_env::RuntimeEnvBuilder; -use datafusion::logical_expr::JoinType; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::expressions::PhysicalSortExpr; use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::prelude::SessionConfig; use datafusion_common::units::{KB, MB}; -use datafusion_common::{JoinSide, NullEquality}; use datafusion_execution::memory_pool::{ FairSpillPool, MemoryConsumer, MemoryReservation, }; @@ -46,7 +44,6 @@ use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, }; -use datafusion_physical_plan::joins::SortMergeJoinExec; use datafusion_physical_plan::metrics::MetricValue; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use futures::StreamExt; @@ -795,224 +792,3 @@ fn get_output_batches_from_metrics(metrics: &MetricsSet) -> usize { }) .expect("Must have output_batches metric since it exists in the baseline") } - -/// Sorted-by-key SMJ input: keys `0..num_keys`, each repeated `rows_per_key` -/// times (SMJ does not sort its inputs), with a `payload_len`-byte payload per -/// row. A key group spans batches when `rows_per_key > batch_size`, which is -/// what forces the buffered side to spill. -fn make_sorted_join_input( - schema: &SchemaRef, - num_keys: i64, - rows_per_key: usize, - payload_len: usize, - batch_size: usize, -) -> Vec { - let payload = "a".repeat(payload_len); - let total = num_keys as usize * rows_per_key; - let mut keys = Vec::with_capacity(total); - for k in 0..num_keys { - for _ in 0..rows_per_key { - keys.push(k); - } - } - - keys.chunks(batch_size) - .map(|chunk| { - let key_arr = Int64Array::from_iter_values(chunk.iter().copied()); - let val_arr = StringArray::from_iter_values(std::iter::repeat_n( - payload.as_str(), - chunk.len(), - )); - RecordBatch::try_new( - Arc::clone(schema), - vec![Arc::new(key_arr), Arc::new(val_arr)], - ) - .expect("valid record batch") - }) - .collect() -} - -/// Run one `SortMergeJoinExec` whose buffered key groups (~3 MB each) overflow -/// a 2 MB pool, asserting it spills and still produces the full, correct -/// output. The streamed side has one row per key, bounding the output to -/// `num_keys * rows_per_key_large` rows. -async fn run_sort_merge_join_test_with_limited_memory(join_type: JoinType) -> Result<()> { - let pool_size = 2 * MB as usize; - let batch_size = 1000; - let num_keys: i64 = 12; - let rows_per_key_large = 6000; - let payload_len_large = 512; - let payload_len_sparse = 16; - - let left_schema = Arc::new(Schema::new(vec![ - Field::new("k1", DataType::Int64, false), - Field::new("v1", DataType::Utf8, false), - ])); - let right_schema = Arc::new(Schema::new(vec![ - Field::new("k2", DataType::Int64, false), - Field::new("v2", DataType::Utf8, false), - ])); - - // SMJ spills the buffered side — the opposite of the probe/streamed side. - let buffered_side = SortMergeJoinExec::probe_side(&join_type).negate(); - let large_left = matches!(buffered_side, JoinSide::Left); - let large_payload = "a".repeat(payload_len_large); - let sparse_payload = "a".repeat(payload_len_sparse); - let (expected_left_payload, expected_right_payload) = if large_left { - (large_payload.as_str(), sparse_payload.as_str()) - } else { - (sparse_payload.as_str(), large_payload.as_str()) - }; - - let left_batches = make_sorted_join_input( - &left_schema, - num_keys, - if large_left { rows_per_key_large } else { 1 }, - if large_left { - payload_len_large - } else { - payload_len_sparse - }, - batch_size, - ); - let right_batches = make_sorted_join_input( - &right_schema, - num_keys, - if large_left { 1 } else { rows_per_key_large }, - if large_left { - payload_len_sparse - } else { - payload_len_large - }, - batch_size, - ); - - // No unmatched rows on either side, so all four join types produce the - // same output row count. - let expected_output_rows = num_keys as usize * rows_per_key_large; - - let task_ctx = { - let memory_pool = Arc::new(FairSpillPool::new(pool_size)); - TaskContext::default() - .with_session_config(SessionConfig::new().with_batch_size(batch_size)) - .with_runtime(Arc::new( - RuntimeEnvBuilder::new() - .with_memory_pool(memory_pool) - .build()?, - )) - }; - let task_ctx = Arc::new(task_ctx); - - let left: Arc = - Arc::new(OnceExec::new(Box::pin(RecordBatchStreamAdapter::new( - Arc::clone(&left_schema), - futures::stream::iter(left_batches.into_iter().map(Ok)), - )))); - let right: Arc = - Arc::new(OnceExec::new(Box::pin(RecordBatchStreamAdapter::new( - Arc::clone(&right_schema), - futures::stream::iter(right_batches.into_iter().map(Ok)), - )))); - - let on = vec![( - Arc::new(Column::new_with_schema("k1", &left_schema)?) as _, - Arc::new(Column::new_with_schema("k2", &right_schema)?) as _, - )]; - - let join = Arc::new(SortMergeJoinExec::try_new( - left, - right, - on, - None, - join_type, - vec![SortOptions::default()], - NullEquality::NullEqualsNothing, - )?); - - let mut stream = join.execute(0, Arc::clone(&task_ctx))?; - let mut output_rows = 0; - let mut output_rows_per_key = vec![0; num_keys as usize]; - while let Some(batch) = stream.next().await { - let batch = batch?; - let left_keys = batch - .column(0) - .as_any() - .downcast_ref::() - .expect("left key must be Int64"); - let left_values = batch - .column(1) - .as_any() - .downcast_ref::() - .expect("left payload must be Utf8"); - let right_keys = batch - .column(2) - .as_any() - .downcast_ref::() - .expect("right key must be Int64"); - let right_values = batch - .column(3) - .as_any() - .downcast_ref::() - .expect("right payload must be Utf8"); - - assert_eq!(left_keys.null_count(), 0); - assert_eq!(left_values.null_count(), 0); - assert_eq!(right_keys.null_count(), 0); - assert_eq!(right_values.null_count(), 0); - for row in 0..batch.num_rows() { - assert_eq!(left_keys.value(row), right_keys.value(row)); - let key = left_keys.value(row); - assert!( - (0..num_keys).contains(&key), - "unexpected output key {key} for {join_type:?}", - ); - output_rows_per_key[key as usize] += 1; - assert_eq!(left_values.value(row), expected_left_payload); - assert_eq!(right_values.value(row), expected_right_payload); - } - output_rows += batch.num_rows(); - } - - assert_eq!( - output_rows, expected_output_rows, - "unexpected output row count for {join_type:?}", - ); - for (key, count) in output_rows_per_key.into_iter().enumerate() { - assert_eq!( - count, rows_per_key_large, - "unexpected output count for key {key} in {join_type:?}", - ); - } - - let metrics = join.metrics().expect("must have metrics"); - assert!( - metrics.spill_count().unwrap_or_default() > 0, - "expected the buffered side to spill for {join_type:?}", - ); - assert!( - metrics.spilled_bytes().unwrap_or_default() > 0, - "expected spilled_bytes > 0 for {join_type:?}", - ); - assert!( - metrics.spilled_rows().unwrap_or_default() > 0, - "expected spilled_rows > 0 for {join_type:?}", - ); - - Ok(()) -} - -/// Memory-limit stress test: ~36 MB of buffered data joined under a 2 MB pool -/// for each materializing join type. Semi/anti/mark joins take a separate -/// execution path (`BitwiseSortMergeJoinStream`) and are out of scope here. -#[tokio::test] -async fn test_sort_merge_join_with_limited_memory() -> Result<()> { - for join_type in [ - JoinType::Inner, - JoinType::Left, - JoinType::Right, - JoinType::Full, - ] { - run_sort_merge_join_test_with_limited_memory(join_type).await?; - } - Ok(()) -} 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 index 4034dbe0828e0..3af642fffe101 100644 --- 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 @@ -15,23 +15,19 @@ // specific language governing permissions and limitations // under the License. -//! Validates that the physical (RSS) memory usage of sort-merge join queries -//! stays within the configured memory limit. Tests must run in separate -//! processes for accurate measurement, so runners spawn each test case as its -//! own process. +//! 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; -/// Guards against validating a hash join the planner picked instead. +/// Ensures the planner selected a sort-merge join. const SMJ_OPERATOR_NAME: &str = "SortMergeJoinExec"; -/// Force the planner to pick a sort-merge join and let it spill under a tight -/// pool: SMJ is only chosen with hash joins disabled and partitions > 1 (kept -/// at the minimum of 2 since each partition multiplies the per-sorter merge -/// reservation), and `sort_spill_reservation_bytes` is shrunk from its 10MB -/// default, which would otherwise exhaust a small pool before any data spills. +/// 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) @@ -39,8 +35,7 @@ fn smj_session_config() -> SessionConfig { .set_bool("datafusion.optimizer.prefer_hash_join", false) } -/// Equi-join with one large right-side key group that overflows the configured -/// pools; `sum(v)` keeps the output to one row. +/// 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) \ @@ -49,11 +44,6 @@ fn smj_sum_query(series_len: usize) -> String { ) } -// =========================================================================== -// Test runners: -// Runners are split into multiple tests to run in parallel -// =========================================================================== - #[test] fn smj_with_mem_limit_1_runner() { utils::spawn_test_process("smj_mem_validation", "smj_with_mem_limit_1"); @@ -69,24 +59,14 @@ fn smj_no_mem_limit_runner() { utils::spawn_test_process("smj_mem_validation", "smj_no_mem_limit"); } -// =========================================================================== -// Test cases: -// All following tests need to be run through their individual test wrapper. -// When run directly, environment variable `DATAFUSION_TEST_MEM_LIMIT_VALIDATION` -// is not set, test will return with a no-op. -// -// If some tests consistently fail, suppress by setting a larger expected memory -// usage (e.g. 40_000_000 * 4 -> 40_000_000 * 5) -// =========================================================================== - -/// 40MB limit against one ~80MB buffered key group: the join must spill. +/// 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), // Baseline query with ~10% of data + &smj_sum_query(500_000), smj_session_config(), Some(SMJ_OPERATOR_NAME), Some(true), @@ -94,14 +74,15 @@ async fn smj_with_mem_limit_1() { .await; } -/// Tighter 16MB limit forces more aggressive spilling. +/// 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 * 5, + 16_000_000 * 12, Some(16_000_000), &smj_sum_query(5_000_000), - &smj_sum_query(500_000), // Baseline query with ~10% of data + &smj_sum_query(500_000), smj_session_config(), Some(SMJ_OPERATOR_NAME), Some(true), @@ -109,14 +90,13 @@ async fn smj_with_mem_limit_2() { .await; } -/// No memory limit: the join must not spill. #[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), // Baseline query with ~10% of data + &smj_sum_query(500_000), smj_session_config(), Some(SMJ_OPERATOR_NAME), Some(false), 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 2cec1f9bae0a8..788b8f4942ee4 100644 --- a/datafusion/core/tests/memory_limit/memory_limit_validation/utils.rs +++ b/datafusion/core/tests/memory_limit/memory_limit_validation/utils.rs @@ -101,8 +101,9 @@ where (result, peak_rss) } -/// Re-run one test from the current integration-test binary in an isolated -/// process so its RSS can be measured independently. +/// 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"); @@ -171,8 +172,7 @@ pub async fn validate_query_with_memory_limits( query: &str, baseline_query: &str, ) { - // Fixed partition count so results are stable across machines. - let session_config = SessionConfig::new().with_target_partitions(4); + 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, @@ -185,10 +185,8 @@ pub async fn validate_query_with_memory_limits( .await; } -/// Like [`validate_query_with_memory_limits`] but with an explicit -/// [`SessionConfig`] and optional operator expectations: the plan must contain -/// `expected_operator_name`, and `expected_operator_spill` requires its spill -/// count to be positive (`true`) or zero (`false`). +/// 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, @@ -229,7 +227,7 @@ pub async fn validate_query_with_memory_limits_and_config( ); } - // Run a query with 10% data to estimate the constant overhead. + // Run a query with 10% data to estimate the constant overhead let baseline_plan = ctx .sql(baseline_query) .await 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