Function setup_context

datafusion/core/tests/memory_limit/mod.rs:558–592

Create a new `SessionContext` with specified disk limit, memory pool limit, and spill compression codec

(
    disk_limit: u64,
    memory_pool_limit: usize,
    spill_compression: SpillCompression,
)

// Create a new `SessionContext` with specified disk limit, memory pool limit, and spill compression codec
async fn setup_context(
    disk_limit: u64,
    memory_pool_limit: usize,
    spill_compression: SpillCompression,
) -> Result<SessionContext> {
    let disk_manager = DiskManagerBuilder::default()
        .with_mode(DiskManagerMode::OsTmpDirectory)
        .with_max_temp_directory_size(disk_limit)
        .build()?;

    let runtime = RuntimeEnvBuilder::new()
        .with_memory_pool(Arc::new(FairSpillPool::new(memory_pool_limit)))
        .build_arc()
        .unwrap();

    let runtime = Arc::new(RuntimeEnv {
        memory_pool: runtime.memory_pool.clone(),
        disk_manager: Arc::new(disk_manager),
        cache_manager: runtime.cache_manager.clone(),
        object_store_registry: runtime.object_store_registry.clone(),
        #[cfg(feature = "parquet_encryption")]
        parquet_encryption_factory_registry: runtime
            .parquet_encryption_factory_registry
            .clone(),
    });

    let config = SessionConfig::new()
        .with_sort_spill_reservation_bytes(64 * 1024) // 256KB
        .with_sort_in_place_threshold_bytes(0)
        .with_spill_compression(spill_compression)
        .with_batch_size(64) // To reduce test memory usage
        .with_target_partitions(1);

    Ok(SessionContext::new_with_config_rt(config, runtime))
}

/// If the spilled bytes exceed the disk limit, the query should fail
/// (specified by `max_temp_directory_size` in `DiskManager`)