update table schema job

Author: bubriks

python/hsfs/core/feature_group_api.py

@@ -21,7 +21,12 @@
 from hopsworks_common import client
 from hsfs import feature_group as fg_mod
 from hsfs import feature_group_commit, util
-from hsfs.core import explicit_provenance, ingestion_job, ingestion_job_conf
+from hsfs.core import (
+    explicit_provenance,
+    ingestion_job,
+    ingestion_job_conf,
+    job,
+)
 
 
 class FeatureGroupApi:
@@ -416,6 +421,36 @@ def ingestion(
             ),
         )
 
+    def update_table_schema(
+        self,
+        feature_group_instance: fg_mod.FeatureGroup,
+    ) -> job.Job:
+        """
+        Setup a Hopsworks job to update table schema
+        Args:
+        feature_group_instance: FeatureGroup, required
+            metadata object of feature group.
+        job_conf: the configuration for the job application
+        """
+
+        _client = client.get_instance()
+        path_params = [
+            "project",
+            _client._project_id,
+            "featurestores",
+            feature_group_instance.feature_store_id,
+            "featuregroups",
+            feature_group_instance.id,
+            "updatetableschema",
+        ]
+
+        headers = {"content-type": "application/json"}
+        return job.Job.from_response_json(
+            _client._send_request(
+                "POST", path_params, headers=headers
+            ),
+        )
+
     def get_parent_feature_groups(
         self,
         feature_group_instance: Union[

python/hsfs/core/feature_group_engine.py

@@ -249,6 +249,8 @@ def commit_delete(feature_group, delete_df, write_options):
     @staticmethod
     def delta_vacuum(feature_group, retention_hours):
         if feature_group.time_travel_format == "DELTA":
+            # TODO: This should change, DeltaEngine and HudiEngine always assumes spark client!
+            # Cannot properly manage what should happen when using python.
             delta_engine_instance = delta_engine.DeltaEngine(
                 feature_group.feature_store_id,
                 feature_group.feature_store_name,

python/hsfs/engine/python.py

@@ -1213,8 +1213,10 @@ def save_stream_dataframe(
         )
 
     def update_table_schema(self, feature_group: Union[FeatureGroup, ExternalFeatureGroup]) -> None:
-        """Wrapper around update_table_schema in order to provide no-op."""
-        pass
+        _job = self._feature_group_api.update_table_schema(feature_group)
+        _job._wait_for_job(
+            await_termination=True
+        )
 
     def _get_app_options(
         self, user_write_options: Optional[Dict[str, Any]] = None

python/hsfs/engine/spark.py

@@ -251,7 +251,7 @@ def reconcile_schema(
     ):
         if sorted(self._spark_session.table(fg_alias.alias).columns) != sorted(
             [feature.name for feature in fg_alias.feature_group._features] +
-            self.HUDI_SPEC_FEATURE_NAMES if fg_alias.feature_group.time_travel_format == "HUDI" else []
+            hudi_engine.HudiEngine.HUDI_SPEC_FEATURE_NAMES if fg_alias.feature_group.time_travel_format == "HUDI" else []
         ):
             full_fg = feature_group_api.FeatureGroupApi().get(
                 feature_store_id=fg_alias.feature_group._feature_store_id,

python/tests/engine/test_python.py

@@ -2565,15 +2565,22 @@ def test_save_stream_dataframe(self):
             == "Stream ingestion is not available on Python environments, because it requires Spark as engine."
         )
 
-    def test_update_table_schema(self):
+    def test_update_table_schema(self, mocker):
         # Arrange
+        mock_fg_api = mocker.patch("hsfs.core.feature_group_api.FeatureGroupApi")
+
         python_engine = python.Engine()
 
+        mock_fg_api.return_value.update_table_schema.return_value.job = job.Job(
+            1, "test_job", None, None, None, None
+        )
+
         # Act
         result = python_engine.update_table_schema(feature_group=None)
 
         # Assert
         assert result is None
+        assert mock_fg_api.return_value.update_table_schema.call_count == 1
 
     def test_get_app_options(self, mocker):
         # Arrange

utils/python/hsfs_utils.py

@@ -323,6 +323,18 @@ def offline_fg_materialization(spark: SparkSession, job_conf: Dict[Any, Any], in
     offset_df = spark.createDataFrame([offset_dict])
     offset_df.coalesce(1).write.mode("overwrite").json(offset_location)
 
+def update_table_schema_fg(spark: SparkSession, job_conf: Dict[Any, Any]) -> None:
+    """
+    Run table schema update job 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.stream = False
+    engine.get_instance().update_table_schema(entity)
+
 def _build_starting_offsets(initial_check_point_string: str):
     if not initial_check_point_string:
         return ""
@@ -358,6 +370,7 @@ def _build_starting_offsets(initial_check_point_string: str):
             "run_feature_monitoring",
             "delta_vacuum_fg",
             "offline_fg_materialization",
+            "update_table_schema_fg",
         ],
         help="Operation type",
     )
@@ -406,6 +419,8 @@ def parse_isoformat_date(da: str) -> datetime:
             delta_vacuum_fg(spark, job_conf)
         elif args.op == "offline_fg_materialization":
             offline_fg_materialization(spark, job_conf, args.initialCheckPointString)
+        elif args.op == "update_table_schema_fg":
+            update_table_schema_fg(spark, job_conf)
 
         success = True
     except Exception: