[FSTORE-1537][APPEND] Managed feature group delta vacuum (#354)

Author: bubriks

python/hsfs/core/delta_engine.py

@@ -158,14 +158,10 @@ def _write_delta_dataset(self, dataset, write_options):
             self._spark_session, location
         )
 
-    def vacuum(self, retention_hours):
+    def vacuum(self, retention_hours: int):
         location = self._feature_group.prepare_spark_location()
-
-        delta_table = DeltaTable.forPath(self._spark_session, location)
-
-        # Vacuum the table
-        # https://docs.delta.io/1.0.1/api/python/index.html#delta.tables.DeltaTable.vacuum
-        delta_table.vacuum(retention_hours)
+        retention = f"RETAIN {retention_hours} HOURS" if retention_hours is not None else ""
+        self._spark_session.sql(f"VACUUM '{location}' {retention}")
 
     def _generate_merge_query(self, source_alias, updates_alias):
         merge_query_list = []

python/hsfs/core/feature_group_engine.py

@@ -247,7 +247,7 @@ def commit_delete(feature_group, delete_df, write_options):
             return hudi_engine_instance.delete_record(delete_df, write_options)
 
     @staticmethod
-    def clean(feature_group, write_options):
+    def delta_vacuum(feature_group, retention_hours):
         if feature_group.time_travel_format == "DELTA":
             delta_engine_instance = delta_engine.DeltaEngine(
                 feature_group.feature_store_id,
@@ -256,7 +256,7 @@ def clean(feature_group, write_options):
                 engine.get_instance()._spark_session,
                 engine.get_instance()._spark_context,
             )
-            return delta_engine_instance.vacuum(write_options.get("retention_hours", None))
+            return delta_engine_instance.vacuum(retention_hours)
         else:
             return None
 

python/hsfs/feature_group.py

@@ -3276,11 +3276,12 @@ def commit_delete_record(
         """
         self._feature_group_engine.commit_delete(self, delete_df, write_options or {})
 
-    def clean(
+    def delta_vacuum(
         self,
-        write_options: Optional[Dict[Any, Any]] = None,
+        retention_hours: int = None,
     ) -> None:
-        """ Clean up old files. This method can only be used on feature groups stored as DELTA.
+        """ Vacuum files that are no longer referenced by a Delta table and are older than the retention threshold.
+        This method can only be used on feature groups stored as DELTA.
 
         !!! example
             ```python
@@ -3290,15 +3291,15 @@ def clean(
             # get the Feature Group instance
             fg = fs.get_or_create_feature_group(...)
 
-            commit_details = fg.clean(write_options = {"retention_hours": 100})
+            commit_details = fg.delta_vacuum(retention_hours = 168)
 
         # Arguments
-            write_options: User provided write options. Defaults to `{}`.
+            retention_hours: User provided retention period. The default retention threshold for the files is 7 days.
 
         # Raises
             `hsfs.client.exceptions.RestAPIError`.
         """
-        self._feature_group_engine.clean(self, write_options or {})
+        self._feature_group_engine.delta_vacuum(self, retention_hours)
 
     def as_of(
         self,

python/tests/core/test_feature_group_engine.py

@@ -547,7 +547,7 @@ def test_clean_delta(self, mocker):
         )
 
         # Act
-        fg_engine.clean(feature_group=fg, write_options={})
+        fg_engine.delta_vacuum(feature_group=fg, retention_hours=200)
 
         # Assert
         assert mock_hudi_engine.return_value.vacuum.call_count == 1
@@ -571,7 +571,7 @@ def test_clean_hudi(self, mocker):
         )
 
         # Act
-        fg_engine.clean(feature_group=fg, write_options={})
+        fg_engine.delta_vacuum(feature_group=fg, retention_hours=200)
 
     def test_sql(self, mocker):
         # Arrange

utils/python/hsfs_utils.py

@@ -247,16 +247,16 @@ def run_feature_monitoring(job_conf: Dict[str, str]) -> None:
         raise e
 
 
-def clean_fg(spark: SparkSession, job_conf: Dict[Any, Any]) -> None:
+def delta_vacuum_fg(spark: SparkSession, job_conf: Dict[Any, Any]) -> None:
     """
-    Run clean on a feature group.
+    Run delta vacuum on a feature group.
     """
     feature_store = job_conf.pop("feature_store")
     fs = get_feature_store_handle(feature_store)
 
     entity = fs.get_feature_group(name=job_conf["name"], version=job_conf["version"])
 
-    entity.clean()
+    entity.delta_vacuum()
 
 
 if __name__ == "__main__":
@@ -277,7 +277,7 @@ def clean_fg(spark: SparkSession, job_conf: Dict[Any, Any]) -> None:
             "ge_validate",
             "import_fg",
             "run_feature_monitoring",
-            "clean_fg",
+            "delta_vacuum_fg",
         ],
         help="Operation type",
     )
@@ -316,8 +316,8 @@ def parse_isoformat_date(da: str) -> datetime:
             import_fg(job_conf)
         elif args.op == "run_feature_monitoring":
             run_feature_monitoring(job_conf)
-        elif args.op == "clean_fg":
-            clean_fg(spark, job_conf)
+        elif args.op == "delta_vacuum_fg":
+            delta_vacuum_fg(spark, job_conf)
 
         success = True
     except Exception: