spark transformation changes

oracle · Sep 30, 2023 · dc4c9de · dc4c9de
1 parent cc3fb53
commit dc4c9de
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/ads/feature_store/common/enums.py b/ads/feature_store/common/enums.py
@@ -225,6 +225,7 @@ class TransformationMode(Enum):
 
     SQL = "sql"
     PANDAS = "pandas"
+    SPARK = "spark"
 
 
 class FilterOperators(Enum):

diff --git a/ads/feature_store/common/utils/transformation_utils.py b/ads/feature_store/common/utils/transformation_utils.py
@@ -69,7 +69,7 @@ def apply_transformation(
                     temporary_table_view, **transformation_kwargs_dict
                 )
             )
-        elif transformation.transformation_mode == TransformationMode.PANDAS.value:
+        elif transformation.transformation_mode in [TransformationMode.PANDAS.value, TransformationMode.SPARK.value]:
             transformed_data = transformation_function_caller(
                 dataframe, **transformation_kwargs_dict
             )