From 738d10786ebd2ff41cd97aafc43f0f419c2069f5 Mon Sep 17 00:00:00 2001 From: Lu Qiu Date: Sun, 13 Sep 2026 12:06:53 -0700 Subject: [PATCH] fix(datafusion): size the memory pool by the effective partition count mem_pool_size() multiplied the 150MB per-partition default by target_partition.unwrap_or(1), but when target_partition is None the session config is not overridden and DataFusion falls back to its default target_partitions, which is the available parallelism (number of CPU cores). The FairSpillPool was therefore sized for one partition while sort-heavy plans (e.g. BTree scalar index training) spawned one ExternalSorterMerge consumer per core, each reserving ~40MB unspillable against a 150MB pool. On multi-core machines this caused load-dependent, flaky ResourcesExhausted failures under concurrent create_index workloads (observed in the frag_reuse clone tests, which had been worked around with serial_test). Size the pool by the effective partition count instead: when target_partition is None, multiply the per-partition default by datafusion_common's get_available_parallelism(), the same source DataFusion uses for its target_partitions default. The env-parse-error fallback branch uses the same value, so the two fallbacks agree. Explicit mem_pool_size (option or LANCE_MEM_POOL_SIZE env var) still sets the total pool size, unscaled, and target_partition: Some(n) still yields 150MB * n. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01DgAshYD7wVzPVdjXRuWPPs --- rust/lance-datafusion/src/exec.rs | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/rust/lance-datafusion/src/exec.rs b/rust/lance-datafusion/src/exec.rs index dfdb5865744..f9b8732b4be 100644 --- a/rust/lance-datafusion/src/exec.rs +++ b/rust/lance-datafusion/src/exec.rs @@ -39,7 +39,7 @@ use datafusion::{ }, }; use datafusion::{execution::memory_pool::TrackConsumersPool, physical_plan::metrics::MetricType}; -use datafusion_common::{DataFusionError, Statistics}; +use datafusion_common::{DataFusionError, Statistics, utils::get_available_parallelism}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use futures::{StreamExt, stream}; @@ -310,8 +310,20 @@ const DEFAULT_LANCE_MEM_POOL_SIZE_PER_PARTITION: u64 = 150 * 1024 * 1024; const DEFAULT_LANCE_MAX_TEMP_DIRECTORY_SIZE: u64 = 100 * 1024 * 1024 * 1024; // 100GB impl LanceExecutionOptions { + /// The number of partitions the DataFusion session will actually run with. + /// + /// When `target_partition` is not set we do not override the session config, + /// so the session falls back to DataFusion's default `target_partitions`, + /// which is the available parallelism (number of CPU cores). The memory pool + /// must be sized for that effective partition count, not for a single + /// partition, or sort-heavy plans exhaust the pool. + fn effective_target_partition(&self) -> u64 { + self.target_partition + .unwrap_or_else(get_available_parallelism) as u64 + } + pub fn mem_pool_size(&self) -> u64 { - let num_partitions = self.target_partition.unwrap_or(1) as u64; + let num_partitions = self.effective_target_partition(); self.mem_pool_size.unwrap_or_else(|| { std::env::var("LANCE_MEM_POOL_SIZE") .map(|s| match s.parse::() { @@ -1334,9 +1346,14 @@ mod tests { fn test_mem_pool_size_scales_with_partitions() { let default_per_partition = DEFAULT_LANCE_MEM_POOL_SIZE_PER_PARTITION; - // No partitions specified → defaults to 1 partition + // No partitions specified → the session runs with DataFusion's default + // target_partitions (available parallelism), so the pool must be sized + // for that effective partition count. let opts = LanceExecutionOptions::default(); - assert_eq!(opts.mem_pool_size(), default_per_partition); + assert_eq!( + opts.mem_pool_size(), + default_per_partition * get_available_parallelism() as u64 + ); // 4 partitions → 4x the per-partition size let opts = LanceExecutionOptions {