feat(optimize): add max_temp_directory_size parameter for z-order and compact operations for DataFusion

Signed-off-by: Florian Valeye <[email protected]>

Author: fvaleye

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;
@@ -217,6 +218,8 @@ pub struct OptimizeBuilder<'a> {
     max_concurrent_tasks: usize,
     /// Maximum number of bytes allowed in memory before spilling to disk
     max_spill_size: usize,
+    /// Maximum size of temporary directory for spilling (default 100GB)
+    max_temp_directory_size: u64,
     /// Optimize type
     optimize_type: OptimizeType,
     /// Datafusion session state relevant for executing the input plan
@@ -247,6 +250,7 @@ impl<'a> OptimizeBuilder<'a> {
             preserve_insertion_order: false,
             max_concurrent_tasks: num_cpus::get(),
             max_spill_size: 20 * 1024 * 1024 * 1024, // 20 GB.
+            max_temp_directory_size: 100 * 1024 * 1024 * 1024, // 100 GB.
             optimize_type: OptimizeType::Compact,
             min_commit_interval: None,
             session: None,
@@ -306,6 +310,16 @@ impl<'a> OptimizeBuilder<'a> {
         self
     }
 
+    /// Max temporary directory size for spilling
+    #[deprecated(
+        since = "0.29.0",
+        note = "Pass in a `SessionState` configured with a `RuntimeEnv` and a `DiskManager`"
+    )]
+    pub fn with_max_temp_directory_size(mut self, max_temp_directory_size: u64) -> Self {
+        self.max_temp_directory_size = max_temp_directory_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);
@@ -351,8 +365,11 @@ impl<'a> std::future::IntoFuture for OptimizeBuilder<'a> {
                 .and_then(|session| session.as_any().downcast_ref::<SessionState>().cloned())
                 .unwrap_or_else(|| {
                     let memory_pool = FairSpillPool::new(this.max_spill_size);
+                    let disk_manager_builder = DiskManagerBuilder::default()
+                        .with_max_temp_directory_size(this.max_temp_directory_size);
                     let runtime = RuntimeEnvBuilder::new()
                         .with_memory_pool(Arc::new(memory_pool))
+                        .with_disk_manager_builder(disk_manager_builder)
                         .build_arc()
                         .unwrap();
                     SessionStateBuilder::new()

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,

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. Defaults to 20GB if not specified.
+            max_temp_directory_size: the maximum size of temporary directory for spilling. Defaults to 100GB if not specified.
             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. Defaults to 20GB if not specified.
+            max_temp_directory_size: the maximum size of temporary directory for spilling. Defaults to 100GB if not specified.
             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,

python/src/lib.rs

@@ -592,6 +592,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 +606,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>,
@@ -615,6 +619,12 @@ impl RawDeltaTable {
             if let Some(size) = target_size {
                 cmd = cmd.with_target_size(size);
             }
+            if let Some(spill_size) = max_spill_size {
+                cmd = cmd.with_max_spill_size(spill_size);
+            }
+            if let Some(directory_size) = max_temp_directory_size {
+                cmd = cmd.with_max_temp_directory_size(directory_size);
+            }
             if let Some(commit_interval) = min_commit_interval {
                 cmd = cmd.with_min_commit_interval(time::Duration::from_secs(commit_interval));
             }
@@ -654,7 +664,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 +677,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,11 +687,16 @@ 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));
             if let Some(size) = target_size {
                 cmd = cmd.with_target_size(size);
             }
+            if let Some(spill_size) = max_spill_size {
+                cmd = cmd.with_max_spill_size(spill_size);
+            }
+            if let Some(directory_size) = max_temp_directory_size {
+                cmd = cmd.with_max_temp_directory_size(directory_size);
+            }
             if let Some(commit_interval) = min_commit_interval {
                 cmd = cmd.with_min_commit_interval(time::Duration::from_secs(commit_interval));
             }

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