From 2f8394380f7af1f5ce44958aef4e3436e58299e7 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 3 Aug 2026 22:47:15 +0800 Subject: [PATCH 1/2] fix: raise sort spill reservation ceiling for default memory pool The sort spill reservation was sized as pool/3 but hard-capped at 40 MiB. Under the default 150 MiB FairSpillPool that leaves only 40 MiB free for ExternalSorterMerge, while BTREE training on large string columns (e.g. 10M uuid keys in merge_insert datagen) peaks near 48 MiB and fails with ResourcesExhausted. Raise the ceiling to 64 MiB so the default pool keeps its intended pool/3 (= 50 MiB) reservation. --- rust/lance-datafusion/src/exec.rs | 54 ++++++++++++++++++++++++++----- 1 file changed, 46 insertions(+), 8 deletions(-) diff --git a/rust/lance-datafusion/src/exec.rs b/rust/lance-datafusion/src/exec.rs index 994e16b783d..0d143d142df 100644 --- a/rust/lance-datafusion/src/exec.rs +++ b/rust/lance-datafusion/src/exec.rs @@ -309,6 +309,26 @@ impl std::fmt::Debug for LanceExecutionOptions { 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 +/// Ceiling for [`sort_spill_reservation_bytes`]: multi-GB pools should not +/// pre-reserve multi-GB for the unspillable merge phase. +const MAX_SORT_SPILL_RESERVATION_BYTES: u64 = 64 * 1024 * 1024; + +/// Bytes reserved for DataFusion's unspillable in-memory sort/merge while an +/// external sort spills. +/// +/// Sized at ~1/3 of the memory pool (see #7675). DataFusion's default 10 MiB +/// is too small: `ExternalSorter` double-counts buffered batches so roughly +/// half the pool is already spoken for when merge starts, and merge then needs +/// additional headroom for cursors and accumulation. +/// +/// The previous 40 MiB ceiling was below what merge needs under the default +/// 150 MiB pool (`pool / 3` = 50 MiB). Training a BTREE index on large string +/// columns (e.g. 10M uuid keys in merge_insert datagen) peaked near 48 MiB in +/// `ExternalSorterMerge` and failed with `ResourcesExhausted`. +fn sort_spill_reservation_bytes(mem_pool_size: u64) -> usize { + (mem_pool_size / 3).min(MAX_SORT_SPILL_RESERVATION_BYTES) as usize +} + impl LanceExecutionOptions { pub fn mem_pool_size(&self) -> u64 { let num_partitions = self.target_partition.unwrap_or(1) as u64; @@ -362,14 +382,9 @@ pub fn new_session_context(options: &LanceExecutionOptions) -> SessionContext { session_config = session_config.with_target_partitions(target_partition); } if options.use_spilling() { - // The default 10MB sort spill reservation seems to be too small for many common cases. - // - // There currently is no reasonable guidance provided by DataFusion for setting this value. - // We bump this to 40MB but try a smaller value if the mem pool is small. - let sort_spill_reservation_bytes = - (options.mem_pool_size() / 3).min(40 * 1024 * 1024) as usize; - session_config = - session_config.with_sort_spill_reservation_bytes(sort_spill_reservation_bytes); + session_config = session_config.with_sort_spill_reservation_bytes( + sort_spill_reservation_bytes(options.mem_pool_size()), + ); let disk_manager_builder = DiskManagerBuilder::default() .with_max_temp_directory_size(options.max_temp_directory_size()); runtime_env_builder = runtime_env_builder @@ -1343,4 +1358,27 @@ mod tests { }; assert_eq!(opts.mem_pool_size(), 50 * 1024 * 1024); } + + #[test] + fn test_sort_spill_reservation_bytes() { + // Small pools: ~1/3, no floor that could exceed the pool. + assert_eq!( + sort_spill_reservation_bytes(4 * 1024 * 1024), + 4 * 1024 * 1024 / 3 + ); + + // Default pool: pool/3 (= 50 MiB), not the old 40 MiB hard cap. + let default_pool = DEFAULT_LANCE_MEM_POOL_SIZE_PER_PARTITION; + assert_eq!( + sort_spill_reservation_bytes(default_pool), + (default_pool / 3) as usize + ); + assert!(sort_spill_reservation_bytes(default_pool) > 40 * 1024 * 1024); + + // Large pools: ceiling, not unbounded growth. + assert_eq!( + sort_spill_reservation_bytes(8 * 1024 * 1024 * 1024), + MAX_SORT_SPILL_RESERVATION_BYTES as usize + ); + } } From 89b92e2f2ece2d5e215d4be7eaaf87e2da65c83e Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 4 Aug 2026 14:02:30 +0800 Subject: [PATCH 2/2] fix(ci): enlarge mem pool only around merge_insert index build BTREE training on merge_insert_narrow's 10M string keys can exhaust the default 150 MiB FairSpillPool during ExternalSorterMerge and abort bench_regress datagen. Raising the library-wide sort spill reservation is not a reliable fix for this class of failure. Enlarge LANCE_MEM_POOL_SIZE to 1 GiB only for the create_scalar_index loop in merge_insert datagen, then restore the previous env so measured benchmark paths keep the default pool. --- .../ci_benchmarks/datagen/merge_insert.py | 31 +++++++++-- rust/lance-datafusion/src/exec.rs | 54 +++---------------- 2 files changed, 36 insertions(+), 49 deletions(-) diff --git a/python/python/ci_benchmarks/datagen/merge_insert.py b/python/python/ci_benchmarks/datagen/merge_insert.py index 33f165d8772..d69a7f4c85f 100644 --- a/python/python/ci_benchmarks/datagen/merge_insert.py +++ b/python/python/ci_benchmarks/datagen/merge_insert.py @@ -43,6 +43,9 @@ from __future__ import annotations +import os +from contextlib import contextmanager + import lance import numpy as np import pyarrow as pa @@ -53,6 +56,13 @@ # Tag marking the pristine version that benchmarks restore to. BASE_TAG = "merge_insert_base" +# BTREE training sorts the full column under DataFusion's FairSpillPool. The +# default 150 MiB pool is tight for 10M 32-byte string keys (uuid columns): +# ExternalSorterMerge has peaked near the pool edge and aborted datagen with +# ResourcesExhausted. Index build is one-shot setup, not a measured path, so +# temporarily enlarge the pool only around create_scalar_index. +_INDEX_BUILD_MEM_POOL_SIZE = str(1024 * 1024 * 1024) + NARROW_NUM_ROWS = 10_000_000 NARROW_ROWS_PER_FRAGMENT = 1_000_000 @@ -272,6 +282,20 @@ def _already_generated(uri: str, expected_rows: int) -> bool: return ds.checkout_version(base_version).count_rows() == expected_rows +@contextmanager +def _enlarged_mem_pool_for_index_build(): + """Raise LANCE_MEM_POOL_SIZE only for the duration of index training.""" + previous = os.environ.get("LANCE_MEM_POOL_SIZE") + os.environ["LANCE_MEM_POOL_SIZE"] = _INDEX_BUILD_MEM_POOL_SIZE + try: + yield + finally: + if previous is None: + os.environ.pop("LANCE_MEM_POOL_SIZE", None) + else: + os.environ["LANCE_MEM_POOL_SIZE"] = previous + + def _gen( name: str, data, @@ -294,9 +318,10 @@ def _gen( max_rows_per_file=rows_per_fragment, max_rows_per_group=min(rows_per_fragment, 100_000), ) - for column in indexed_columns: - LOGGER.info("Building BTREE index on %s.%s", name, column) - ds.create_scalar_index(column, "BTREE") + with _enlarged_mem_pool_for_index_build(): + for column in indexed_columns: + LOGGER.info("Building BTREE index on %s.%s", name, column) + ds.create_scalar_index(column, "BTREE") _tag_base(ds) return ds diff --git a/rust/lance-datafusion/src/exec.rs b/rust/lance-datafusion/src/exec.rs index 0d143d142df..994e16b783d 100644 --- a/rust/lance-datafusion/src/exec.rs +++ b/rust/lance-datafusion/src/exec.rs @@ -309,26 +309,6 @@ impl std::fmt::Debug for LanceExecutionOptions { 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 -/// Ceiling for [`sort_spill_reservation_bytes`]: multi-GB pools should not -/// pre-reserve multi-GB for the unspillable merge phase. -const MAX_SORT_SPILL_RESERVATION_BYTES: u64 = 64 * 1024 * 1024; - -/// Bytes reserved for DataFusion's unspillable in-memory sort/merge while an -/// external sort spills. -/// -/// Sized at ~1/3 of the memory pool (see #7675). DataFusion's default 10 MiB -/// is too small: `ExternalSorter` double-counts buffered batches so roughly -/// half the pool is already spoken for when merge starts, and merge then needs -/// additional headroom for cursors and accumulation. -/// -/// The previous 40 MiB ceiling was below what merge needs under the default -/// 150 MiB pool (`pool / 3` = 50 MiB). Training a BTREE index on large string -/// columns (e.g. 10M uuid keys in merge_insert datagen) peaked near 48 MiB in -/// `ExternalSorterMerge` and failed with `ResourcesExhausted`. -fn sort_spill_reservation_bytes(mem_pool_size: u64) -> usize { - (mem_pool_size / 3).min(MAX_SORT_SPILL_RESERVATION_BYTES) as usize -} - impl LanceExecutionOptions { pub fn mem_pool_size(&self) -> u64 { let num_partitions = self.target_partition.unwrap_or(1) as u64; @@ -382,9 +362,14 @@ pub fn new_session_context(options: &LanceExecutionOptions) -> SessionContext { session_config = session_config.with_target_partitions(target_partition); } if options.use_spilling() { - session_config = session_config.with_sort_spill_reservation_bytes( - sort_spill_reservation_bytes(options.mem_pool_size()), - ); + // The default 10MB sort spill reservation seems to be too small for many common cases. + // + // There currently is no reasonable guidance provided by DataFusion for setting this value. + // We bump this to 40MB but try a smaller value if the mem pool is small. + let sort_spill_reservation_bytes = + (options.mem_pool_size() / 3).min(40 * 1024 * 1024) as usize; + session_config = + session_config.with_sort_spill_reservation_bytes(sort_spill_reservation_bytes); let disk_manager_builder = DiskManagerBuilder::default() .with_max_temp_directory_size(options.max_temp_directory_size()); runtime_env_builder = runtime_env_builder @@ -1358,27 +1343,4 @@ mod tests { }; assert_eq!(opts.mem_pool_size(), 50 * 1024 * 1024); } - - #[test] - fn test_sort_spill_reservation_bytes() { - // Small pools: ~1/3, no floor that could exceed the pool. - assert_eq!( - sort_spill_reservation_bytes(4 * 1024 * 1024), - 4 * 1024 * 1024 / 3 - ); - - // Default pool: pool/3 (= 50 MiB), not the old 40 MiB hard cap. - let default_pool = DEFAULT_LANCE_MEM_POOL_SIZE_PER_PARTITION; - assert_eq!( - sort_spill_reservation_bytes(default_pool), - (default_pool / 3) as usize - ); - assert!(sort_spill_reservation_bytes(default_pool) > 40 * 1024 * 1024); - - // Large pools: ceiling, not unbounded growth. - assert_eq!( - sort_spill_reservation_bytes(8 * 1024 * 1024 * 1024), - MAX_SORT_SPILL_RESERVATION_BYTES as usize - ); - } }