fixed restore

oracle · Dec 20, 2023 · 28d1688 · 28d1688
1 parent 09cdd19
commit 28d1688
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ads/feature_store/execution_strategy/spark/spark_execution.py b/ads/feature_store/execution_strategy/spark/spark_execution.py
@@ -366,7 +366,7 @@ def update_feature_definition_features(self, feature_group, target_table):
         try:
             # Get the output features
             output_features = get_features(
-                self.spark_engine.get_columns_from_table(target_table), feature_group.id
+                self.spark_engine.get_output_columns_from_table_or_dataframe(table_name=target_table), feature_group.id
             )
             if output_features:
                 feature_group._with_features(output_features)
@@ -392,7 +392,7 @@ def update_dataset_features(self, dataset, target_table):
         try:
             # Get the output features
             output_features = get_features(
-                output_columns=self.spark_engine.get_columns_from_table(target_table),
+                output_columns=self.spark_engine.get_output_columns_from_table_or_dataframe(table_name=target_table),
                 parent_id=dataset.id,
                 entity_type=EntityType.DATASET,
             )