SMJ + aggregate finish with one partition but stuck forever when number of partitions is big

[SemyonSinchenko]

Is your feature request related to a problem or challenge?

I'm experimenting with out-of-core data processing with DataFusion. I'm trying to run algorithms over data that is much bigger than available memory using DataFusion's SMJ and manually checkpointing everything to parquet files. It works if I specify the number of partitions equal to 1 but it slow and does not utilize all the available CPUs. It stuck with something that looks like a dead-lock with big number of partitions. It sometimes works (sometimes not) when number of partitions is in between.

I have some hypothesis of the root cause but I got it mostly after talking with AI so I would like to avoid to speculate about the reason.

("datafusion.optimizer.prefer_hash_join", false)

is important for me because in my scenarios non of the join side can fit to the memeory and from what I read in documentation, SMJ is the only join that can spill to disk both sides.

Describe the solution you'd like

At least it would be nice to get fast-fail error like "not enough memory". Overall it looks strange to me that the same heavy out-of-core SMJ passes for a single partition but fails for multiple partitions.

Describe alternatives you've considered

It would be nice to get some advice on how to tune DataFusion for out-of-core scenarios.

Additional context

This is minimal reproducable example:

use std::path::{Path, PathBuf};
use std::process;
use std::sync::Arc;
use std::time::Duration;

use datafusion::arrow::util::pretty;
use datafusion::dataframe::DataFrameWriteOptions;
use datafusion::error::{DataFusionError, Result};
use datafusion::execution::memory_pool::FairSpillPool;
use datafusion::execution::runtime_env::RuntimeEnvBuilder;
use datafusion::execution::session_state::SessionStateBuilder;
use datafusion::functions_aggregate::sum::sum;
use datafusion::prelude::*;

// ---- dataset shape ---------------------------------------------------------
// Tuned so that a debug build still builds/runs the whole thing in well under
// a minute, while guaranteeing heavy spilling once `pool_size` is small.
const NUM_VERTICES: usize = 300_000;
const NUM_HUBS: usize = 5; // hub keys: 0..NUM_HUBS
const HUB_FAN: usize = 200_000; // edges per hub  -> 1_000_000 hub edges
const NUM_TAIL: usize = 1_000_000; // uniformly distributed edges
// => 2_000_000 edges total (~32 MiB raw: two int64 columns)

const DEFAULT_TIMEOUT_SECS: u64 = 180;

/// Tiny deterministic PRNG (LCG) so the generated dataset is reproducible and
/// we don't depend on a particular `rand` API.
struct Lcg(u64);

impl Lcg {
    fn new(seed: u64) -> Self {
        Self(seed)
    }
    fn next_u64(&mut self) -> u64 {
        // Numerical Recipes LCG constants.
        self.0 = self
            .0
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        self.0
    }
    fn range(&mut self, n: usize) -> usize {
        if n == 0 {
            0
        } else {
            (self.next_u64() % n as u64) as usize
        }
    }
}

/// Parse a human memory size like `10M`, `1G`, `512M`, `3G` into bytes
/// (`M`/`MiB` = 2^20, `G`/`GiB` = 2^30). Mirrors `run-algorithm`'s `parse_mem`.
fn parse_mem(s: &str) -> Result<usize> {
    let s = s.trim();
    let split = s
        .find(|c: char| !c.is_ascii_digit())
        .ok_or_else(|| DataFusionError::Execution(format!("missing size unit in '{s}'")))?;
    let (digits, unit) = s.split_at(split);
    let num: usize = digits
        .parse()
        .map_err(|e| DataFusionError::Execution(format!("invalid number in '{s}': {e}")))?;
    let mult = match unit.to_ascii_lowercase().as_str() {
        "g" | "gb" | "gib" => 1024 * 1024 * 1024,
        "m" | "mb" | "mib" => 1024 * 1024,
        other => {
            return Err(DataFusionError::Execution(format!(
                "wrong memory unit '{other}' (use M or G)"
            )))
        }
    };
    Ok(num * mult)
}

fn file_uri(dir: &Path) -> String {
    // Absolute path like "/tmp/x" -> "file:///tmp/x/".
    format!("file://{}/", dir.to_string_lossy())
}

/// Build the `state(id, value)` table in memory via the `dataframe!` macro.
fn gen_state() -> Result<DataFrame> {
    let ids: Vec<i64> = (0..NUM_VERTICES as i64).collect();
    let values: Vec<f64> = (0..NUM_VERTICES)
        .map(|i| (i as f64) / (NUM_VERTICES as f64))
        .collect();
    Ok(dataframe!(
        "id" => ids,
        "value" => values,
    )?)
}

/// Build the `edges(src, dst)` table in memory via the `dataframe!` macro,
/// with heavy skew on `src` (a few hub vertices carry most of the edges).
fn gen_edges() -> Result<DataFrame> {
    let total = HUB_FAN * NUM_HUBS + NUM_TAIL;
    let mut src = Vec::with_capacity(total);
    let mut dst = Vec::with_capacity(total);
    let mut rng = Lcg::new(0xC0FFEE);

    // Hub edges: src in 0..NUM_HUBS, each repeated HUB_FAN times.
    for hub in 0..NUM_HUBS {
        for _ in 0..HUB_FAN {
            src.push(hub as i64);
            dst.push(rng.range(NUM_VERTICES) as i64);
        }
    }
    // Tail edges: src uniform across all vertices.
    for _ in 0..NUM_TAIL {
        src.push(rng.range(NUM_VERTICES) as i64);
        dst.push(rng.range(NUM_VERTICES) as i64);
    }

    Ok(dataframe!(
        "src" => src,
        "dst" => dst,
    )?)
}

/// Write a DataFrame to a parquet directory as `num_partitions` files.
/// Uses the DataFrame's own (default, unbounded) session so the write never
/// competes with the tight `FairSpillPool` used for the read/join/agg.
async fn write_parquet_dir(
    df: DataFrame,
    num_partitions: usize,
    dir: &Path,
    label: &str,
) -> Result<()> {
    std::fs::create_dir_all(dir)?;
    let uri = file_uri(dir);
    println!("[setup] writing {label} -> {uri} ({num_partitions} partitions)");
    // RoundRobinBatch repartition gives us `num_partitions` output files.
    let df = if num_partitions > 1 {
        df.repartition(Partitioning::RoundRobinBatch(num_partitions))?
    } else {
        df
    };
    df.write_parquet(&uri, DataFrameWriteOptions::new(), None)
        .await?;
    Ok(())
}

/// Build the *run* context: tiny `FairSpillPool`, forced `SortMergeJoin`,
/// `target_partitions = num_partitions`.
fn build_run_ctx(num_partitions: usize, pool_bytes: usize) -> Result<SessionContext> {
    let config = SessionConfig::new()
        .with_target_partitions(num_partitions)
        .set_bool("datafusion.optimizer.prefer_hash_join", false);

    let env = RuntimeEnvBuilder::new()
        .with_memory_pool(Arc::new(FairSpillPool::new(pool_bytes)))
        .build_arc()?;

    let state = SessionStateBuilder::new()
        .with_config(config)
        .with_runtime_env(env)
        .with_default_features()
        .build();

    Ok(SessionContext::from(state))
}

/// Build the query: state ⋈ edges ON state.id = edges.src,
/// project (dst, value), then GROUP BY dst SUM(value).
async fn build_query(ctx: &SessionContext, state_uri: &str, edges_uri: &str) -> Result<DataFrame> {
    let state = ctx.read_parquet(state_uri, ParquetReadOptions::new()).await?;
    let edges = ctx.read_parquet(edges_uri, ParquetReadOptions::new()).await?;

    let triplets = state.join(edges, JoinType::Inner, &["id"], &["src"], None)?;
    let messages = triplets.select(vec![col("dst"), col("value")])?;
    let agg = messages.aggregate(vec![col("dst")], vec![sum(col("value"))])?;
    Ok(agg)
}

#[tokio::main]
async fn main() -> Result<()> {
    env_logger::init();

    let args: Vec<String> = std::env::args().collect();
    if args.len() < 3 {
        eprintln!(
            "usage: {} <num_partitions> <pool_size> [timeout_secs]\n\
             example: {} 4 10M\n\
             example: {} 6 1M 0   (0 = no timeout, attach gdb to the hung PID)",
            args[0], args[0], args[0]
        );
        std::process::exit(2);
    }

    let num_partitions: usize = args[1]
        .parse()
        .map_err(|e| DataFusionError::Execution(format!("invalid num_partitions: {e}")))?;
    let pool_bytes = parse_mem(&args[2])?;
    let timeout_secs: u64 = if args.len() >= 4 {
        args[3].parse().map_err(|e| {
            DataFusionError::Execution(format!("invalid timeout_secs: {e}"))
        })?
    } else {
        DEFAULT_TIMEOUT_SECS
    };

    println!("============================================================");
    println!("reproduce_deadlock");
    println!("  pid             : {}", process::id());
    println!("  num_partitions  : {num_partitions}");
    println!("  FairSpillPool   : {} bytes ({} MiB)", pool_bytes, pool_bytes / (1024 * 1024));
    println!("  prefer_hash_join: false  (forces SortMergeJoin)");
    println!("  vertices        : {NUM_VERTICES}");
    println!(
        "  edges           : {} (skewed: {NUM_HUBS} hubs x {HUB_FAN} + {NUM_TAIL} tail)",
        HUB_FAN * NUM_HUBS + NUM_TAIL
    );
    println!("  query timeout   : {} (0 = none)", if timeout_secs == 0 {
        "disabled".to_string()
    } else {
        format!("{timeout_secs}s")
    });
    println!("============================================================");

    // ---- 1-4. Generate + write both tables to parquet ----------------------
    // A fresh temp dir per run; left in place if we hang so you can inspect.
    let base: PathBuf = std::env::temp_dir()
        .join(format!("reproduce_deadlock_{}", process::id()));
    let state_dir = base.join("state");
    let edges_dir = base.join("edges");
    std::fs::create_dir_all(&base)?;
    let state_uri = file_uri(&state_dir);
    let edges_uri = file_uri(&edges_dir);

    println!("\n[setup] generating in-memory datasets with dataframe!() ...");
    let state_df = gen_state()?;
    let edges_df = gen_edges()?;
    println!("[setup] state rows: {}, edges rows: {}", state_df.clone().count().await?, edges_df.clone().count().await?);

    write_parquet_dir(state_df, num_partitions, &state_dir, "state").await?;
    write_parquet_dir(edges_df, num_partitions, &edges_dir, "edges").await?;
    println!("[setup] done writing parquet.\n");

    // ---- 5-6. Read back, build the query, print the EXPLAIN ----------------
    let ctx = build_run_ctx(num_partitions, pool_bytes)?;
    let agg = build_query(&ctx, &state_uri, &edges_uri).await?;

    let explained = agg.clone().explain(true, false)?;
    let explain_batches = explained.collect().await?;
    let explain_str = pretty::pretty_format_batches(&explain_batches)
        .map_err(|e| DataFusionError::Execution(format!("failed to format explain: {e}")))?;
    let explain_path = base.join("explain.txt");
    std::fs::write(&explain_path, explain_str.to_string())?;
    println!("[explain] written to {}", explain_path.display());

    // ---- 7. Run the query (collect) ----------------------------------------
    // If it deadlocks, this never returns. With timeout_secs > 0 we surface a
    // clear message; with 0 you are expected to attach gdb to the printed PID.
    println!("[run] collecting aggregate result ...");
    let run_fut = async {
        let batches = agg.collect().await?;
        println!("[run] completed: {} output batches, {} rows",
            batches.len(),
            batches.iter().map(|b| b.num_rows()).sum::<usize>());
        Ok::<_, DataFusionError>(())
    };

    if timeout_secs == 0 {
        run_fut.await?;
    } else {
        match tokio::time::timeout(Duration::from_secs(timeout_secs), run_fut).await {
            Ok(res) => res?,
            Err(_) => {
                eprintln!("\n!! DEADLOCKED: query did not finish within {timeout_secs}s.");
                eprintln!("!! Re-run with a trailing '0' to disable the timeout, then attach gdb:");
                eprintln!("!!     {} {} {} {} 0", args[0], args[1], args[2], args.get(3).map(|s| s.as_str()).unwrap_or("<timeout>"));
                eprintln!("!!     gdb -p {}", process::id());
            }
        }
    }

    // ---- cleanup (only reached if it did not hang) -------------------------
    let _ = std::fs::remove_dir_all(&base);
    Ok(())
}

I'm running it in three scenarios:

1. Many partitions.

./target/debug/reproduce_deadlock 10 150M --- it stuck in the middle:

sem@fedora:~/github/graphframes-rs$ ./target/debug/reproduce_deadlock 10 150M
============================================================
reproduce_deadlock
  pid             : 33467
  num_partitions  : 10
  FairSpillPool   : 157286400 bytes (150 MiB)
  prefer_hash_join: false  (forces SortMergeJoin)
  vertices        : 300000
  edges           : 2000000 (skewed: 5 hubs x 200000 + 1000000 tail)
  query timeout   : 180s (0 = none)
============================================================

[setup] generating in-memory datasets with dataframe!() ...
[setup] state rows: 300000, edges rows: 2000000
[setup] writing state -> file:///tmp/reproduce_deadlock_33467/state/ (10 partitions)
[setup] writing edges -> file:///tmp/reproduce_deadlock_33467/edges/ (10 partitions)
[setup] done writing parquet.

[explain] written to /tmp/reproduce_deadlock_33467/explain.txt
[run] collecting aggregate result ...
!! DEADLOCKED: query did not finish within 180s.
!! Re-run with a trailing '0' to disable the timeout, then attach gdb:
!!     ./target/debug/reproduce_deadlock 10 150M <timeout> 0
!!     gdb -p 33467

If I run rust-gdb --pid 33467 I see:

(gdb) info threads 
  Id   Target Id                                           Frame 
* 1    Thread 0x7ff205d30d40 (LWP 33467) "reproduce_deadl" syscall () at ../sysdeps/unix/sysv/linux/x86_64/syscall.S:38
  2    Thread 0x7ff2049256c0 (LWP 33479) "tokio-rt-worker" syscall () at ../sysdeps/unix/sysv/linux/x86_64/syscall.S:38
  3    Thread 0x7ff204b266c0 (LWP 33478) "tokio-rt-worker" syscall () at ../sysdeps/unix/sysv/linux/x86_64/syscall.S:38
  4    Thread 0x7ff204d276c0 (LWP 33477) "tokio-rt-worker" syscall () at ../sysdeps/unix/sysv/linux/x86_64/syscall.S:38
  5    Thread 0x7ff204f286c0 (LWP 33476) "tokio-rt-worker" syscall () at ../sysdeps/unix/sysv/linux/x86_64/syscall.S:38
  6    Thread 0x7ff2051296c0 (LWP 33475) "tokio-rt-worker" syscall () at ../sysdeps/unix/sysv/linux/x86_64/syscall.S:38
  7    Thread 0x7ff20532a6c0 (LWP 33474) "tokio-rt-worker" syscall () at ../sysdeps/unix/sysv/linux/x86_64/syscall.S:38
  8    Thread 0x7ff1fd1ff6c0 (LWP 33473) "tokio-rt-worker" syscall () at ../sysdeps/unix/sysv/linux/x86_64/syscall.S:38
  9    Thread 0x7ff20552b6c0 (LWP 33472) "tokio-rt-worker" syscall () at ../sysdeps/unix/sysv/linux/x86_64/syscall.S:38
  10   Thread 0x7ff20572c6c0 (LWP 33471) "tokio-rt-worker" syscall () at ../sysdeps/unix/sysv/linux/x86_64/syscall.S:38
  11   Thread 0x7ff20592d6c0 (LWP 33470) "tokio-rt-worker" __syscall_cancel_arch () at ../sysdeps/unix/sysv/linux/x86_64/syscall_cancel.S:56
  12   Thread 0x7ff205b2e6c0 (LWP 33469) "tokio-rt-worker" syscall () at ../sysdeps/unix/sysv/linux/x86_64/syscall.S:38
  13   Thread 0x7ff205d2f6c0 (LWP 33468) "tokio-rt-worker" syscall () at ../sysdeps/unix/sysv/linux/x86_64/syscall.S:38

(gdb) thread apply 1 bt

Thread 1 (Thread 0x7ff205d30d40 (LWP 33467) "reproduce_deadl"):
#0  syscall () at ../sysdeps/unix/sysv/linux/x86_64/syscall.S:38
#1  0x000055c801c1ebe1 in parking_lot_core::thread_parker::imp::ThreadParker::futex_wait (self=0x7ff205d30ca8, ts=...) at /var/home/sem/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/parking_lot_core-0.9.11/src/thread_parker/linux.rs:112
#2  0x000055c801c1ea2c in parking_lot_core::thread_parker::imp::{impl#0}::park (self=0x7ff205d30ca8) at /var/home/sem/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/parking_lot_core-0.9.11/src/thread_parker/linux.rs:66
#3  0x000055c801c1cc4a in parking_lot_core::parking_lot::park::{closure#0}<parking_lot::condvar::{impl#1}::wait_until_internal::{closure_env#0}, parking_lot::condvar::{impl#1}::wait_until_internal::{closure_env#1}, parking_lot::condvar::{impl#1}::wait_until_internal::{closure_env#2}> (thread_data=0x7ff205d30c88) at /var/home/sem/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/parking_lot_core-0.9.11/src/parking_lot.rs:635
#4  0x000055c801c1b4de in parking_lot_core::parking_lot::with_thread_data<parking_lot_core::parking_lot::ParkResult, parking_lot_core::parking_lot::park::{closure_env#0}<parking_lot::condvar::{impl#1}::wait_until_internal::{closure_env#0}, parking_lot::condvar::{impl#1}::wait_until_internal::{closure_env#1}, parking_lot::condvar::{impl#1}::wait_until_internal::{closure_env#2}>> (f=...) at /var/home/sem/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/parking_lot_core-0.9.11/src/parking_lot.rs:207
#5  parking_lot_core::parking_lot::park<parking_lot::condvar::{impl#1}::wait_until_internal::{closure_env#0}, parking_lot::condvar::{impl#1}::wait_until_internal::{closure_env#1}, parking_lot::condvar::{impl#1}::wait_until_internal::{closure_env#2}> (key=94318091852488, validate=..., before_sleep=..., timed_out=..., park_token=..., timeout=...) at /var/home/sem/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/parking_lot_core-0.9.11/src/parking_lot.rs:600
#6  0x000055c801c126da in parking_lot::condvar::Condvar::wait_until_internal (self=0x55c8245c5ac8, mutex=0x55c8245c5ad0, timeout=...) at src/condvar.rs:334
#7  0x000055c801ba9b02 in parking_lot::condvar::Condvar::wait<()> (self=0x55c8245c5ac8, mutex_guard=0x7fff325fbbd8) at /var/home/sem/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/parking_lot-0.12.4/src/condvar.rs:256
....

For me it looks like: all threads are parked, nothing happens. As well I do not see any signs of i/o from the process or any CPU activities.

2. Single partition

It passes in seconds.

sem@fedora:~/github/graphframes-rs$ ./target/debug/reproduce_deadlock 1 150M
============================================================
reproduce_deadlock
  pid             : 63694
  num_partitions  : 1
  FairSpillPool   : 157286400 bytes (150 MiB)
  prefer_hash_join: false  (forces SortMergeJoin)
  vertices        : 300000
  edges           : 2000000 (skewed: 5 hubs x 200000 + 1000000 tail)
  query timeout   : 180s (0 = none)
============================================================

[setup] generating in-memory datasets with dataframe!() ...
[setup] state rows: 300000, edges rows: 2000000
[setup] writing state -> file:///tmp/reproduce_deadlock_63694/state/ (1 partitions)
[setup] writing edges -> file:///tmp/reproduce_deadlock_63694/edges/ (1 partitions)
[setup] done writing parquet.

[explain] written to /tmp/reproduce_deadlock_63694/explain.txt
[run] collecting aggregate result ...
[run] completed: 36 output batches, 294586 rows
sem@fedora:~/github/graphframes-rs$ 

3. Small amount of partitions

It passes as well.

em@fedora:~/github/graphframes-rs$ ./target/debug/reproduce_deadlock 4 150M
============================================================
reproduce_deadlock
  pid             : 63879
  num_partitions  : 4
  FairSpillPool   : 157286400 bytes (150 MiB)
  prefer_hash_join: false  (forces SortMergeJoin)
  vertices        : 300000
  edges           : 2000000 (skewed: 5 hubs x 200000 + 1000000 tail)
  query timeout   : 180s (0 = none)
============================================================

[setup] generating in-memory datasets with dataframe!() ...
[setup] state rows: 300000, edges rows: 2000000
[setup] writing state -> file:///tmp/reproduce_deadlock_63879/state/ (4 partitions)
[setup] writing edges -> file:///tmp/reproduce_deadlock_63879/edges/ (4 partitions)
[setup] done writing parquet.

[explain] written to /tmp/reproduce_deadlock_63879/explain.txt
[run] collecting aggregate result ...
[run] completed: 36 output batches, 294586 rows

Additional information
I can run any gdb command or anything else and paste the result. Just tell me what to do please.

Thanks in advance!

[SemyonSinchenko]

That is the whole "explain" of the query I saved in the middle of the example.

explain.txt

[Omega359]

possible tangentially related: #21625

[pepijnve]

Quick check with tokio-console seems to confirm that things are stalling somewhere. I'll see if I can figure out where and why.

[Omega359]

I pointed codex at this pr and it thinks the issue is related to FairSpillPool. It ran the reproducer above among other things and came up with this:

Likely cause: FairSpillPool fair-share accounting interacts badly with high partition fan-out. The 10-partition plan registers many spillable consumers before the query has really made progress: grouped aggregates, two sorters per SMJ partition, and several RepartitionExec reservations. FairSpillPool::try_grow divides available spillable memory by num_spill, so once ~60 spillable consumers are registered, each active consumer only gets a small per-consumer allowance. That makes RepartitionExec, ExternalSorter, and later grouped aggregate spill paths repeatedly hit memory pressure instead of progressing.

Key code points:

• dev/datafusion/datafusion/execution/src/memory_pool/pool.rs:236: FairSpillPool::try_grow caps each spillable reservation at (pool_size - unspillable) / num_spill.
• dev/datafusion/datafusion/physical-plan/src/repartition/mod.rs:452: RepartitionExec registers one spillable reservation per output partition.
• dev/datafusion/datafusion/physical-plan/src/sorts/sort.rs:255: each ExternalSorter registers a spillable sorter reservation and an unspillable merge reservation.
• dev/datafusion/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs:515: grouped hash aggregate marks its reservation spillable when spilling is enabled.

Evidence from runs using a temporary repro crate outside the repo:

• FairSpillPool, 150M: 4, 6, 7 partitions completed; 8, 9, 10 partitions timed out.
• FairSpillPool, 10 partitions: 180M timed out; 220M, 300M, 600M completed.
• GreedyMemoryPool, 10 partitions, same 150M: completed quickly.
• Logging wrapper showed roughly 60 spillable consumers registered in the 10-partition case, with repeated try_grow failures first around RepartitionExec / ExternalSorter, then grouped aggregate.

Results:

  pool    partitions  memory  sort reservation  result
  fair    8           150M    default           timed out at 90s
  greedy  8           150M    default           completed
  fair    10          150M    default           timed out at 90s
  greedy  10          150M    default           completed
  greedy  12          150M    default           fast OOM error
  greedy  16          150M    default           fast OOM error
  greedy  20          150M    default           fast OOM error
  greedy  12          220M    default           completed
  greedy  20          300M    default           completed
  greedy  20          150M    1 MiB             completed
  fair    12          220M    default           completed
  fair    20          150M    1 MiB             timed out at 90s

  The greedy failures at 12+ partitions with 150M were not stalls. They failed quickly with:

  Not enough memory to continue external sort ...
  Failed to allocate additional 10.0 MB for ExternalSorterMerge[...]

  So the behavior is pretty distinct:

  - GreedyMemoryPool either completes or fast-fails with a clear OOM-style error.
  - FairSpillPool is the one producing the timeout/stall pattern under the same 8/10 partition, 150M setup.