add clean fg (#7)

Author: bubriks

python/hsfs/core/delta_engine.py

@@ -158,6 +158,15 @@ def _write_delta_dataset(self, dataset, write_options):
             self._spark_session, location
         )
 
+    def vacuum(self, retention_hours):
+        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)
+
     def _generate_merge_query(self, source_alias, updates_alias):
         merge_query_list = []
         primary_key = self._feature_group.primary_key

python/hsfs/core/feature_group_engine.py

@@ -246,6 +246,20 @@ 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):
+        if feature_group.time_travel_format == "DELTA":
+            delta_engine_instance = delta_engine.DeltaEngine(
+                feature_group.feature_store_id,
+                feature_group.feature_store_name,
+                feature_group,
+                engine.get_instance()._spark_session,
+                engine.get_instance()._spark_context,
+            )
+            return delta_engine_instance.vacuum(write_options.get("retention_hours", None))
+        else:
+            return None
+
     def sql(self, query, feature_store_name, dataframe_type, online, read_options):
         if online and self._online_conn is None:
             self._online_conn = self._storage_connector_api.get_online_connector(

python/hsfs/engine/python.py

@@ -481,6 +481,12 @@ def register_external_temporary_table(
         # No op to avoid query failure
         pass
 
+    def register_delta_temporary_table(
+        self, delta_fg_alias, feature_store_id, feature_store_name, read_options
+    ):
+        # No op to avoid query failure
+        pass
+
     def register_hudi_temporary_table(
         self,
         hudi_fg_alias: "hsfs.constructor.hudi_feature_group_alias.HudiFeatureGroupAlias",

python/hsfs/feature_group.py

@@ -3277,6 +3277,30 @@ def commit_delete_record(
         """
         self._feature_group_engine.commit_delete(self, delete_df, write_options or {})
 
+    def clean(
+        self,
+        write_options: Optional[Dict[Any, Any]] = None,
+    ) -> None:
+        """ Clean up old files. This method can only be used on feature groups stored as DELTA.
+
+        !!! example
+            ```python
+            # connect to the Feature Store
+            fs = ...
+
+            # get the Feature Group instance
+            fg = fs.get_or_create_feature_group(...)
+
+            commit_details = fg.clean(write_options = {"retention_hours": 100})
+
+        # Arguments
+            write_options: User provided write options. Defaults to `{}`.
+
+        # Raises
+            `hsfs.client.exceptions.RestAPIError`.
+        """
+        self._feature_group_engine.clean(self, write_options or {})
+
     def as_of(
         self,
         wallclock_time: Optional[Union[str, int, datetime, date]] = None,

utils/python/hsfs_utils.py

@@ -247,6 +247,18 @@ def run_feature_monitoring(job_conf: Dict[str, str]) -> None:
         raise e
 
 
+def clean_fg(spark: SparkSession, job_conf: Dict[Any, Any]) -> None:
+    """
+    Run clean 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()
+
+
 if __name__ == "__main__":
     # Setup spark first so it fails faster in case of args errors
     # Otherwise the resource manager will wait until the spark application master
@@ -265,6 +277,7 @@ def run_feature_monitoring(job_conf: Dict[str, str]) -> None:
             "ge_validate",
             "import_fg",
             "run_feature_monitoring",
+            "clean_fg",
         ],
         help="Operation type",
     )
@@ -303,6 +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)
 
         success = True
     except Exception: