Partitioned hash join keeps its hash maps allocated but unreserved after it finishes (dynamic filter pushdown) #26980

//! Memory a hash join leaves allocated after its stream is dropped, next to
//! what the memory pool has reserved. Run with:
//! cargo test -p datafusion --test hj_map_retention -- --nocapture

use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::Arc;
use std::sync::atomic::{AtomicIsize, Ordering::Relaxed};

use datafusion::arrow::array::Int64Array;
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool};
use datafusion::execution::runtime_env::RuntimeEnvBuilder;
use datafusion::parquet::arrow::ArrowWriter;
use datafusion::physical_plan::execute_stream;
use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext};
use futures::StreamExt;

/// Counts live heap bytes.
struct Counting;
static LIVE: AtomicIsize = AtomicIsize::new(0);

unsafe impl GlobalAlloc for Counting {
    unsafe fn alloc(&self, l: Layout) -> *mut u8 {
        LIVE.fetch_add(l.size() as isize, Relaxed);
        unsafe { System.alloc(l) }
    }
    unsafe fn dealloc(&self, p: *mut u8, l: Layout) {
        LIVE.fetch_sub(l.size() as isize, Relaxed);
        unsafe { System.dealloc(p, l) }
    }
}

#[global_allocator]
static GLOBAL: Counting = Counting;

fn mib(bytes: isize) -> f64 {
    bytes as f64 / (1024.0 * 1024.0)
}

/// One Int64 column `k` holding `0, step, 2 * step, ...`, in 4 Parquet files.
fn write_table(dir: &std::path::Path, name: &str, rows: i64, step: i64) {
    let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)]));
    std::fs::create_dir_all(dir.join(name)).unwrap();
    for f in 0..4 {
        let path = dir.join(name).join(format!("{f}.parquet"));
        let file = std::fs::File::create(path).unwrap();
        let mut w = ArrowWriter::try_new(file, Arc::clone(&schema), None).unwrap();
        let rows = (f * rows / 4..(f + 1) * rows / 4).collect::<Vec<_>>();
        for chunk in rows.chunks(65_536) {
            let col = Int64Array::from_iter_values(chunk.iter().map(|v| v * step));
            let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(col)]);
            w.write(&batch.unwrap()).unwrap();
        }
        w.close().unwrap();
    }
}

async fn run(dir: &std::path::Path, label: &str, repartition: bool, filter: bool) {
    let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(8 << 30));
    let runtime = RuntimeEnvBuilder::new()
        .with_memory_pool(Arc::clone(&pool))
        .build_arc()
        .unwrap();
    let mut config = SessionConfig::new().with_target_partitions(4);
    let options = config.options_mut();
    options.optimizer.enable_join_dynamic_filter_pushdown = filter;
    // false gives a CollectLeft join.
    options.optimizer.repartition_joins = repartition;
    let ctx = SessionContext::new_with_config_rt(config, runtime);
    for t in ["build", "probe"] {
        let path = dir.join(t);
        let options = ParquetReadOptions::default();
        ctx.register_parquet(t, path.to_str().unwrap(), options)
            .await
            .unwrap();
    }

    let base = LIVE.load(Relaxed);
    let sql = "SELECT COUNT(*) FROM probe p JOIN build b ON p.k = b.k";
    let plan = ctx.sql(sql).await.unwrap().create_physical_plan().await.unwrap();
    let mut stream = execute_stream(Arc::clone(&plan), ctx.task_ctx()).unwrap();
    while let Some(batch) = stream.next().await {
        batch.unwrap();
    }
    drop(stream);
    let allocated = LIVE.load(Relaxed) - base;
    let reserved = pool.reserved() as isize;
    drop(plan);
    let after_plan_drop = LIVE.load(Relaxed) - base;
    println!(
        "{label}: after the stream is dropped: allocated {:.1} MiB, reserved {:.1} MiB; \
         after the plan is dropped: allocated {:.1} MiB",
        mib(allocated),
        mib(reserved),
        mib(after_plan_drop),
    );
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn hj_map_retention() {
    let dir = tempfile::tempdir().unwrap();
    // 4M build keys, every 1000th value, so the join builds a hash table.
    write_table(dir.path(), "build", 4_000_000, 1000);
    write_table(dir.path(), "probe", 8_000_000, 1);

    run(dir.path(), "1. Partitioned, dynamic filter off", true, false).await;
    run(dir.path(), "2. Partitioned, dynamic filter on ", true, true).await;
    run(dir.path(), "3. CollectLeft, dynamic filter on ", false, true).await;
}