update_table_schema also for delta table

Author: bubriks

python/hsfs/core/feature_group_engine.py

@@ -296,10 +296,7 @@ def append_features(self, feature_group, new_features):
         )
 
         # write empty dataframe to update parquet schema
-        if feature_group.time_travel_format == "DELTA":
-            engine.get_instance().add_cols_to_delta_table(feature_group)
-        else:
-            engine.get_instance().save_empty_dataframe(feature_group)
+        engine.get_instance().update_table_schema(feature_group)
 
     def update_description(self, feature_group, description):
         """Updates the description of a feature group."""

python/hsfs/core/hudi_engine.py

@@ -234,25 +234,6 @@ def _setup_hudi_read_opts(self, hudi_fg_alias, read_options):
 
         return hudi_options
 
-    def reconcile_hudi_schema(
-        self, save_empty_dataframe_callback, hudi_fg_alias, read_options
-    ):
-        if sorted(self._spark_session.table(hudi_fg_alias.alias).columns) != sorted(
-            [feature.name for feature in hudi_fg_alias.feature_group._features] + self.HUDI_SPEC_FEATURE_NAMES
-        ):
-            full_fg = self._feature_group_api.get(
-                feature_store_id=hudi_fg_alias.feature_group._feature_store_id,
-                name=hudi_fg_alias.feature_group.name,
-                version=hudi_fg_alias.feature_group.version,
-            )
-
-            save_empty_dataframe_callback(full_fg)
-
-            self.register_temporary_table(
-                hudi_fg_alias,
-                read_options,
-            )
-
     @staticmethod
     def _get_last_commit_metadata(spark_context, base_path):
         hopsfs_conf = spark_context._jvm.org.apache.hadoop.fs.FileSystem.get(

python/hsfs/engine/python.py

@@ -1212,12 +1212,8 @@ def save_stream_dataframe(
             "Stream ingestion is not available on Python environments, because it requires Spark as engine."
         )
 
-    def save_empty_dataframe(self, feature_group: Union[FeatureGroup, ExternalFeatureGroup]) -> None:
-        """Wrapper around save_dataframe in order to provide no-op."""
-        pass
-
-    def add_cols_to_delta_table(self, feature_group: FeatureGroup) -> None:
-        """Wrapper around add_cols_to_delta_table in order to provide no-op."""
+    def update_table_schema(self, feature_group: Union[FeatureGroup, ExternalFeatureGroup]) -> None:
+        """Wrapper around update_table_schema in order to provide no-op."""
         pass
 
     def _get_app_options(

python/hsfs/engine/spark.py

@@ -221,8 +221,8 @@ def register_hudi_temporary_table(
             read_options,
         )
 
-        hudi_engine_instance.reconcile_hudi_schema(
-            self.save_empty_dataframe, hudi_fg_alias, read_options
+        self.reconcile_schema(
+            hudi_fg_alias, read_options, hudi_engine_instance
         )
 
     def register_delta_temporary_table(
@@ -241,6 +241,30 @@ def register_delta_temporary_table(
             read_options,
         )
 
+        self.reconcile_schema(
+            delta_fg_alias, read_options, delta_engine_instance
+        )
+
+    def reconcile_schema(
+        self, fg_alias, read_options, engine_instance
+    ):
+        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 []
+        ):
+            full_fg = self._feature_group_api.get(
+                feature_store_id=fg_alias.feature_group._feature_store_id,
+                name=fg_alias.feature_group.name,
+                version=fg_alias.feature_group.version,
+            )
+
+            self.update_table_schema(full_fg)
+
+            engine_instance.register_temporary_table(
+                fg_alias,
+                read_options,
+            )
+
     def _return_dataframe_type(self, dataframe, dataframe_type):
         if dataframe_type.lower() in ["default", "spark"]:
             return dataframe
@@ -1324,7 +1348,13 @@ def is_spark_dataframe(self, dataframe):
             return True
         return False
 
-    def save_empty_dataframe(self, feature_group):
+    def update_table_schema(self, feature_group):
+        if feature_group.time_travel_format == "DELTA":
+            self._add_cols_to_delta_table(feature_group)
+        else:
+            self._save_empty_dataframe(feature_group)
+
+    def _save_empty_dataframe(self, feature_group):
         location = feature_group.prepare_spark_location()
 
         dataframe = self._spark_session.read.format("hudi").load(location)
@@ -1343,7 +1373,7 @@ def save_empty_dataframe(self, feature_group):
             {},
         )
 
-    def add_cols_to_delta_table(self, feature_group):
+    def _add_cols_to_delta_table(self, feature_group):
         location = feature_group.prepare_spark_location()
 
         dataframe = self._spark_session.read.format("delta").load(location)

python/tests/client/test_base_client.py

@@ -20,7 +20,6 @@
 import requests
 from hsfs.client.base import Client
 from hsfs.client.exceptions import RestAPIError
-
 from tests.util import changes_environ
 
 

python/tests/core/test_feature_group_engine.py

@@ -709,7 +709,7 @@ def test_append_features(self, mocker):
 
         # Assert
         assert (
-            mock_engine_get_instance.return_value.save_empty_dataframe.call_count == 1
+            mock_engine_get_instance.return_value.update_table_schema.call_count == 1
         )
         assert len(mock_fg_engine_update_features_metadata.call_args[0][1]) == 4
 

python/tests/engine/test_python.py

@@ -2565,12 +2565,12 @@ def test_save_stream_dataframe(self):
             == "Stream ingestion is not available on Python environments, because it requires Spark as engine."
         )
 
-    def test_save_empty_dataframe(self):
+    def test_update_table_schema(self):
         # Arrange
         python_engine = python.Engine()
 
         # Act
-        result = python_engine.save_empty_dataframe(feature_group=None)
+        result = python_engine.update_table_schema(feature_group=None)
 
         # Assert
         assert result is None

python/tests/engine/test_spark.py

@@ -203,6 +203,7 @@ def test_register_hudi_temporary_table(self, mocker):
         # Arrange
         mock_hudi_engine = mocker.patch("hsfs.core.hudi_engine.HudiEngine")
         mocker.patch("hsfs.feature_group.FeatureGroup.from_response_json")
+        mock_reconcile_schema = mocker.patch("hsfs.engine.spark.Engine.reconcile_schema")
 
         spark_engine = spark.Engine()
 
@@ -220,6 +221,31 @@ def test_register_hudi_temporary_table(self, mocker):
 
         # Assert
         assert mock_hudi_engine.return_value.register_temporary_table.call_count == 1
+        assert mock_reconcile_schema.call_count == 1
+
+    def test_register_delta_temporary_table(self, mocker):
+        # Arrange
+        mock_delta_engine = mocker.patch("hsfs.core.delta_engine.DeltaEngine")
+        mocker.patch("hsfs.feature_group.FeatureGroup.from_response_json")
+        mock_reconcile_schema = mocker.patch("hsfs.engine.spark.Engine.reconcile_schema")
+
+        spark_engine = spark.Engine()
+
+        hudi_fg_alias = hudi_feature_group_alias.HudiFeatureGroupAlias(
+            feature_group=None, alias=None
+        )
+
+        # Act
+        spark_engine.register_delta_temporary_table(
+            delta_fg_alias=hudi_fg_alias,
+            feature_store_id=None,
+            feature_store_name=None,
+            read_options=None,
+        )
+
+        # Assert
+        assert mock_delta_engine.return_value.register_temporary_table.call_count == 1
+        assert mock_reconcile_schema.call_count == 1
 
     def test_return_dataframe_type_default(self, mocker):
         # Arrange
@@ -4540,7 +4566,7 @@ def test_is_spark_dataframe_spark_dataframe(self):
         # Assert
         assert result is True
 
-    def test_save_empty_dataframe(self, mocker):
+    def test_update_table_schema_hudi(self, mocker):
         # Arrange
         mock_spark_engine_save_dataframe = mocker.patch(
             "hsfs.engine.spark.Engine.save_dataframe"
@@ -4560,15 +4586,42 @@ def test_save_empty_dataframe(self, mocker):
             partition_key=[],
             id=10,
             featurestore_name="test_featurestore",
+            time_travel_format="HUDI",
         )
 
         # Act
-        spark_engine.save_empty_dataframe(feature_group=fg)
+        spark_engine.update_table_schema(feature_group=fg)
 
         # Assert
         assert mock_spark_engine_save_dataframe.call_count == 1
         assert mock_spark_read.format.call_count == 1
 
+    def test_update_table_schema_delta(self, mocker):
+        # Arrange
+        mock_spark_read = mocker.patch("pyspark.sql.SparkSession.read")
+        mock_format = mocker.Mock()
+        mock_spark_read.format.return_value = mock_format
+
+        # Arrange
+        spark_engine = spark.Engine()
+
+        fg = feature_group.FeatureGroup(
+            name="test",
+            version=1,
+            featurestore_id=99,
+            primary_key=[],
+            partition_key=[],
+            id=10,
+            featurestore_name="test_featurestore",
+            time_travel_format="DELTA",
+        )
+
+        # Act
+        spark_engine.update_table_schema(feature_group=fg)
+
+        # Assert
+        assert mock_spark_read.format.call_count == 1
+
     def test_apply_transformation_function_single_output_udf_default_mode(self, mocker):
         # Arrange
         mocker.patch("hopsworks_common.client.get_instance")