feat: support enable/disable auto-cleanup (#4037)

Author: yanghua

python/python/lance/dataset.py

@@ -2697,6 +2697,17 @@ def update_config(self, upsert_values: Dict[str, str]) -> None:
         """
         self._ds.update_config(upsert_values)
 
+    def delete_config_keys(self, keys: list[str]) -> None:
+        """Delete specified configuration keys from the dataset.
+
+        Parameters
+        ----------
+        keys : list of str
+            A list of configuration keys to remove from the dataset.
+            Non-existent keys will be silently ignored.
+        """
+        self._ds.delete_config_keys(keys)
+
     @property
     def optimize(self) -> "DatasetOptimizer":
         return DatasetOptimizer(self)
@@ -3836,6 +3847,29 @@ def optimize_indices(self, **kwargs):
         """
         self._dataset._ds.optimize_indices(**kwargs)
 
+    def enable_auto_cleanup(self, auto_cleanup_config: AutoCleanupConfig, **kwargs):
+        """Enable autocleaning for an existing dataset.
+
+        Parameters
+        ----------
+        auto_cleanup_config: AutoCleanupConfig
+            Config options for automatic cleanup of the dataset.
+            If set, dataset's old versions will be automatically
+            cleaned up according to this parameter.
+        """
+        self._dataset._ds.update_config(
+            {
+                "lance.auto_cleanup.interval": str(auto_cleanup_config["interval"]),
+                "lance.auto_cleanup.older_than": f"{auto_cleanup_config['older_than_seconds']}s",  # noqa E501
+            }
+        )
+
+    def disable_auto_cleanup(self, **kwargs):
+        """Disable autocleaning via delete related keys."""
+        self._dataset._ds.delete_config_keys(
+            ["lance.auto_cleanup.interval", "lance.auto_cleanup.older_than"]
+        )
+
 
 class Tags:
     """

python/python/tests/test_dataset.py

@@ -1166,6 +1166,41 @@ def test_auto_cleanup_invalid(tmp_path):
     assert len(dataset.versions()) == 4
 
 
+def test_enable_disable_auto_cleanup(tmp_path):
+    table = pa.Table.from_pydict({"a": range(100), "b": range(100)})
+    base_dir = tmp_path / "test"
+    ds = lance.write_dataset(table, base_dir, mode="create")
+    auto_cleanup_options = AutoCleanupConfig(
+        interval=1,
+        older_than_seconds=1,
+    )
+    # enable auto cleanup
+    ds.optimize.enable_auto_cleanup(auto_cleanup_options)
+    lance.write_dataset(table, base_dir, mode="append")
+    lance.write_dataset(table, base_dir, mode="append")
+    lance.write_dataset(table, base_dir, mode="append")
+
+    time.sleep(5)
+
+    # trigger cleanup
+    lance.write_dataset(table, base_dir, mode="append")
+    assert len(ds.versions()) == 2
+
+    # this is a transactional commit, so will increase a version
+    ds.optimize.disable_auto_cleanup()
+
+    lance.write_dataset(table, base_dir, mode="append")
+    lance.write_dataset(table, base_dir, mode="append")
+    lance.write_dataset(table, base_dir, mode="append")
+
+    time.sleep(5)
+
+    # wait to see if cleanup would be trigger
+    lance.write_dataset(table, base_dir, mode="append")
+
+    assert len(ds.versions()) == 7
+
+
 def test_create_from_commit(tmp_path: Path):
     table = pa.Table.from_pydict({"a": range(100), "b": range(100)})
     base_dir = tmp_path / "test"

python/src/dataset.rs

@@ -1768,6 +1768,16 @@ impl Dataset {
         self.ds = Arc::new(new_self);
         Ok(())
     }
+
+    #[pyo3(signature = (keys))]
+    fn delete_config_keys(&mut self, keys: Vec<String>) -> PyResult<()> {
+        let key_refs: Vec<&str> = keys.iter().map(|k| k.as_str()).collect();
+        let mut new_self = self.ds.as_ref().clone();
+        RT.block_on(None, new_self.delete_config_keys(&key_refs))?
+            .map_err(|err| PyIOError::new_err(err.to_string()))?;
+        self.ds = Arc::new(new_self);
+        Ok(())
+    }
 }
 
 #[derive(FromPyObject)]