feat(optimize): configure DataFusion spill settings for optimize operations

fvaleye · fvaleye · commit b0d47cfef2bb · 2025-10-14T16:00:37.000+02:00
Signed-off-by: Florian Valeye &lt;florian.valeye@gmail.com&gt;
diff --git a/crates/core/src/operations/optimize.rs b/crates/core/src/operations/optimize.rs
@@ -29,6 +29,7 @@ use arrow_array::RecordBatch;
 use arrow_schema::SchemaRef as ArrowSchemaRef;
 use datafusion::catalog::Session;
 use datafusion::execution::context::SessionState;
+use datafusion::execution::disk_manager::DiskManagerBuilder;
 use datafusion::execution::memory_pool::FairSpillPool;
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion::execution::SessionStateBuilder;
@@ -215,8 +216,6 @@ pub struct OptimizeBuilder<'a> {
     preserve_insertion_order: bool,
     /// Maximum number of concurrent tasks (default is number of cpus)
     max_concurrent_tasks: usize,
-    /// Maximum number of bytes allowed in memory before spilling to disk
-    max_spill_size: usize,
     /// Optimize type
     optimize_type: OptimizeType,
     /// Datafusion session state relevant for executing the input plan
@@ -234,6 +233,38 @@ impl super::Operation<()> for OptimizeBuilder<'_> {
     }
 }
 
+/// Create a SessionState configured for optimize operations with custom spill settings.
+///
+/// This is the recommended way to configure memory and disk limits for optimize operations.
+/// The created SessionState should be passed to [`OptimizeBuilder`] via [`with_session_state`](OptimizeBuilder::with_session_state).
+///
+/// # Arguments
+/// * `max_spill_size` - Maximum bytes in memory before spilling to disk. If `None`, uses DataFusion's default memory pool.
+/// * `max_temp_directory_size` - Maximum disk space for temporary spill files. If `None`, uses DataFusion's default disk manager.
+pub fn create_session_state_for_optimize(
+    max_spill_size: Option<usize>,
+    max_temp_directory_size: Option<u64>,
+) -> SessionState {
+    let mut runtime_builder = RuntimeEnvBuilder::new();
+
+    if let Some(spill_size) = max_spill_size {
+        let memory_pool = FairSpillPool::new(spill_size);
+        runtime_builder = runtime_builder.with_memory_pool(Arc::new(memory_pool));
+    }
+
+    if let Some(directory_size) = max_temp_directory_size {
+        let disk_manager_builder =
+            DiskManagerBuilder::default().with_max_temp_directory_size(directory_size);
+        runtime_builder = runtime_builder.with_disk_manager_builder(disk_manager_builder);
+    }
+
+    let runtime = runtime_builder.build_arc().unwrap();
+    SessionStateBuilder::new()
+        .with_default_features()
+        .with_runtime_env(runtime)
+        .build()
+}
+
 impl<'a> OptimizeBuilder<'a> {
     /// Create a new [`OptimizeBuilder`]
     pub fn new(log_store: LogStoreRef, snapshot: EagerSnapshot) -> Self {
@@ -246,7 +277,6 @@ impl<'a> OptimizeBuilder<'a> {
             commit_properties: CommitProperties::default(),
             preserve_insertion_order: false,
             max_concurrent_tasks: num_cpus::get(),
-            max_spill_size: 20 * 1024 * 1024 * 1024, // 20 GB.
             optimize_type: OptimizeType::Compact,
             min_commit_interval: None,
             session: None,
@@ -296,16 +326,6 @@ impl<'a> OptimizeBuilder<'a> {
         self
     }
 
-    /// Max spill size
-    #[deprecated(
-        since = "0.29.0",
-        note = "Pass in a `SessionState` configured with a `RuntimeEnv` and a `FairSpillPool`"
-    )]
-    pub fn with_max_spill_size(mut self, max_spill_size: usize) -> Self {
-        self.max_spill_size = max_spill_size;
-        self
-    }
-
     /// Min commit interval
     pub fn with_min_commit_interval(mut self, min_commit_interval: Duration) -> Self {
         self.min_commit_interval = Some(min_commit_interval);
@@ -349,17 +369,7 @@ impl<'a> std::future::IntoFuture for OptimizeBuilder<'a> {
             let session = this
                 .session
                 .and_then(|session| session.as_any().downcast_ref::<SessionState>().cloned())
-                .unwrap_or_else(|| {
-                    let memory_pool = FairSpillPool::new(this.max_spill_size);
-                    let runtime = RuntimeEnvBuilder::new()
-                        .with_memory_pool(Arc::new(memory_pool))
-                        .build_arc()
-                        .unwrap();
-                    SessionStateBuilder::new()
-                        .with_default_features()
-                        .with_runtime_env(runtime)
-                        .build()
-                });
+                .unwrap_or_else(|| create_session_state_for_optimize(None, None));
             let plan = create_merge_plan(
                 &this.log_store,
                 this.optimize_type,
diff --git a/python/deltalake/_internal.pyi b/python/deltalake/_internal.pyi
@@ -117,6 +117,8 @@ class RawDeltaTable:
         partition_filters: PartitionFilterType | None,
         target_size: int | None,
         max_concurrent_tasks: int | None,
+        max_spill_size: int | None,
+        max_temp_directory_size: int | None,
         min_commit_interval: int | None,
         writer_properties: WriterProperties | None,
         commit_properties: CommitProperties | None,
@@ -129,6 +131,7 @@ class RawDeltaTable:
         target_size: int | None,
         max_concurrent_tasks: int | None,
         max_spill_size: int | None,
+        max_temp_directory_size: int | None,
         min_commit_interval: int | None,
         writer_properties: WriterProperties | None,
         commit_properties: CommitProperties | None,
diff --git a/python/deltalake/table.py b/python/deltalake/table.py
@@ -1934,6 +1934,8 @@ def compact(
         partition_filters: FilterConjunctionType | None = None,
         target_size: int | None = None,
         max_concurrent_tasks: int | None = None,
+        max_spill_size: int | None = None,
+        max_temp_directory_size: int | None = None,
         min_commit_interval: int | timedelta | None = None,
         writer_properties: WriterProperties | None = None,
         post_commithook_properties: PostCommitHookProperties | None = None,
@@ -1956,6 +1958,8 @@ def compact(
             max_concurrent_tasks: the maximum number of concurrent tasks to use for
                                     file compaction. Defaults to number of CPUs. More concurrent tasks can make compaction
                                     faster, but will also use more memory.
+            max_spill_size: the maximum number of bytes allowed in memory before spilling to disk. If not specified, uses DataFusion's default.
+            max_temp_directory_size: the maximum disk space for temporary spill files. If not specified, uses DataFusion's default.
             min_commit_interval: minimum interval in seconds or as timedeltas before a new commit is
                                     created. Interval is useful for long running executions. Set to 0 or timedelta(0), if you
                                     want a commit per partition.
@@ -1989,6 +1993,8 @@ def compact(
             self.table._stringify_partition_values(partition_filters),
             target_size,
             max_concurrent_tasks,
+            max_spill_size,
+            max_temp_directory_size,
             min_commit_interval,
             writer_properties,
             commit_properties,
@@ -2003,7 +2009,8 @@ def z_order(
         partition_filters: FilterConjunctionType | None = None,
         target_size: int | None = None,
         max_concurrent_tasks: int | None = None,
-        max_spill_size: int = 20 * 1024 * 1024 * 1024,
+        max_spill_size: int | None = None,
+        max_temp_directory_size: int | None = None,
         min_commit_interval: int | timedelta | None = None,
         writer_properties: WriterProperties | None = None,
         post_commithook_properties: PostCommitHookProperties | None = None,
@@ -2023,7 +2030,8 @@ def z_order(
             max_concurrent_tasks: the maximum number of concurrent tasks to use for
                                     file compaction. Defaults to number of CPUs. More concurrent tasks can make compaction
                                     faster, but will also use more memory.
-            max_spill_size: the maximum number of bytes allowed in memory before spilling to disk. Defaults to 20GB.
+            max_spill_size: the maximum number of bytes allowed in memory before spilling to disk. If not specified, uses DataFusion's default.
+            max_temp_directory_size: the maximum disk space for temporary spill files. If not specified, uses DataFusion's default.
             min_commit_interval: minimum interval in seconds or as timedeltas before a new commit is
                                     created. Interval is useful for long running executions. Set to 0 or timedelta(0), if you
                                     want a commit per partition.
@@ -2059,6 +2067,7 @@ def z_order(
             target_size,
             max_concurrent_tasks,
             max_spill_size,
+            max_temp_directory_size,
             min_commit_interval,
             writer_properties,
             commit_properties,
diff --git a/python/src/lib.rs b/python/src/lib.rs
@@ -40,7 +40,9 @@ use deltalake::operations::delete::DeleteBuilder;
 use deltalake::operations::drop_constraints::DropConstraintBuilder;
 use deltalake::operations::filesystem_check::FileSystemCheckBuilder;
 use deltalake::operations::load_cdf::CdfLoadBuilder;
-use deltalake::operations::optimize::{OptimizeBuilder, OptimizeType};
+use deltalake::operations::optimize::{
+    create_session_state_for_optimize, OptimizeBuilder, OptimizeType,
+};
 use deltalake::operations::restore::RestoreBuilder;
 use deltalake::operations::set_tbl_properties::SetTablePropertiesBuilder;
 use deltalake::operations::update::UpdateBuilder;
@@ -592,6 +594,8 @@ impl RawDeltaTable {
         partition_filters = None,
         target_size = None,
         max_concurrent_tasks = None,
+        max_spill_size = None,
+        max_temp_directory_size = None,
         min_commit_interval = None,
         writer_properties=None,
         commit_properties=None,
@@ -604,6 +608,8 @@ impl RawDeltaTable {
         partition_filters: Option<Vec<(PyBackedStr, PyBackedStr, PartitionFilterValue)>>,
         target_size: Option<u64>,
         max_concurrent_tasks: Option<usize>,
+        max_spill_size: Option<usize>,
+        max_temp_directory_size: Option<u64>,
         min_commit_interval: Option<u64>,
         writer_properties: Option<PyWriterProperties>,
         commit_properties: Option<PyCommitProperties>,
@@ -612,6 +618,14 @@ impl RawDeltaTable {
         let (table, metrics) = py.allow_threads(|| {
             let mut cmd = OptimizeBuilder::new(self.log_store()?, self.cloned_state()?)
                 .with_max_concurrent_tasks(max_concurrent_tasks.unwrap_or_else(num_cpus::get));
+
+            // Configure session with custom spill settings if provided
+            if max_spill_size.is_some() || max_temp_directory_size.is_some() {
+                let session =
+                    create_session_state_for_optimize(max_spill_size, max_temp_directory_size);
+                cmd = cmd.with_session_state(Arc::new(session));
+            }
+
             if let Some(size) = target_size {
                 cmd = cmd.with_target_size(size);
             }
@@ -654,7 +668,8 @@ impl RawDeltaTable {
         partition_filters = None,
         target_size = None,
         max_concurrent_tasks = None,
-        max_spill_size = 20 * 1024 * 1024 * 1024,
+        max_spill_size = None,
+        max_temp_directory_size = None,
         min_commit_interval = None,
         writer_properties=None,
         commit_properties=None,
@@ -666,7 +681,8 @@ impl RawDeltaTable {
         partition_filters: Option<Vec<(PyBackedStr, PyBackedStr, PartitionFilterValue)>>,
         target_size: Option<u64>,
         max_concurrent_tasks: Option<usize>,
-        max_spill_size: usize,
+        max_spill_size: Option<usize>,
+        max_temp_directory_size: Option<u64>,
         min_commit_interval: Option<u64>,
         writer_properties: Option<PyWriterProperties>,
         commit_properties: Option<PyCommitProperties>,
@@ -675,8 +691,15 @@ impl RawDeltaTable {
         let (table, metrics) = py.allow_threads(|| {
             let mut cmd = OptimizeBuilder::new(self.log_store()?, self.cloned_state()?)
                 .with_max_concurrent_tasks(max_concurrent_tasks.unwrap_or_else(num_cpus::get))
-                .with_max_spill_size(max_spill_size)
                 .with_type(OptimizeType::ZOrder(z_order_columns));
+
+            // Configure session with custom spill settings if provided
+            if max_spill_size.is_some() || max_temp_directory_size.is_some() {
+                let session =
+                    create_session_state_for_optimize(max_spill_size, max_temp_directory_size);
+                cmd = cmd.with_session_state(Arc::new(session));
+            }
+
             if let Some(size) = target_size {
                 cmd = cmd.with_target_size(size);
             }
diff --git a/python/tests/test_optimize.py b/python/tests/test_optimize.py
@@ -266,3 +266,49 @@ def test_optimize_schema_evolved_3185(tmp_path):
     assert dt.version() == 2
     last_action = dt.history(1)[0]
     assert last_action["operation"] == "OPTIMIZE"
+
+
+def test_compact_with_spill_parameters(
+    tmp_path: pathlib.Path,
+    sample_table: Table,
+):
+    write_deltalake(tmp_path, sample_table, mode="append")
+    write_deltalake(tmp_path, sample_table, mode="append")
+    write_deltalake(tmp_path, sample_table, mode="append")
+
+    dt = DeltaTable(tmp_path)
+    old_version = dt.version()
+    old_num_files = len(dt.file_uris())
+
+    dt.optimize.compact(
+        max_spill_size=100 * 1024 * 1024 * 1024,  # 100 GB
+        max_temp_directory_size=500 * 1024 * 1024 * 1024,  # 500 GB
+    )
+
+    last_action = dt.history(1)[0]
+    assert last_action["operation"] == "OPTIMIZE"
+    assert dt.version() == old_version + 1
+    assert len(dt.file_uris()) <= old_num_files
+
+
+def test_z_order_with_spill_parameters(
+    tmp_path: pathlib.Path,
+    sample_table: Table,
+):
+    write_deltalake(tmp_path, sample_table, mode="append")
+    write_deltalake(tmp_path, sample_table, mode="append")
+    write_deltalake(tmp_path, sample_table, mode="append")
+
+    dt = DeltaTable(tmp_path)
+    old_version = dt.version()
+
+    dt.optimize.z_order(
+        columns=["sold", "price"],
+        max_spill_size=100 * 1024 * 1024 * 1024,  # 100 GB
+        max_temp_directory_size=500 * 1024 * 1024 * 1024,  # 500 GB
+    )
+
+    last_action = dt.history(1)[0]
+    assert last_action["operation"] == "OPTIMIZE"
+    assert dt.version() == old_version + 1
+    assert len(dt.file_uris()) == 1
