allowing multiple on-demand features from a single feature group and also supporting passing of context variables in transformation functions

Author: manu-sj

python/hsfs/core/feature_group_engine.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import List, Union
+from typing import Any, Dict, List, Union
 
 from hsfs import engine, feature, util
 from hsfs import feature_group as fg
@@ -49,12 +49,18 @@ def _update_feature_group_schema_on_demand_transformations(
             transformed_features = []
             dropped_features = []
             for tf in feature_group.transformation_functions:
-                transformed_features.append(
-                    feature.Feature(
-                        tf.hopsworks_udf.output_column_names[0],
-                        tf.hopsworks_udf.return_types[0],
-                        on_demand=True,
-                    )
+                transformed_features.extend(
+                    [
+                        feature.Feature(
+                            output_column_name,
+                            return_type,
+                            on_demand=True,
+                        )
+                        for output_column_name, return_type in zip(
+                            tf.hopsworks_udf.output_column_names,
+                            tf.hopsworks_udf.return_types,
+                        )
+                    ]
                 )
                 if tf.hopsworks_udf.dropped_features:
                     dropped_features.extend(tf.hopsworks_udf.dropped_features)
@@ -141,6 +147,7 @@ def insert(
         storage,
         write_options,
         validation_options: dict = None,
+        transformation_context: Dict[str, Any] = None,
     ):
         dataframe_features = engine.get_instance().parse_schema_feature_group(
             feature_dataframe,
@@ -154,7 +161,9 @@ def insert(
             and feature_group.transformation_functions
         ):
             feature_dataframe = engine.get_instance()._apply_transformation_function(
-                feature_group.transformation_functions, feature_dataframe
+                feature_group.transformation_functions,
+                feature_dataframe,
+                transformation_context=transformation_context,
             )
 
         dataframe_features = (
@@ -358,6 +367,7 @@ def insert_stream(
         timeout,
         checkpoint_dir,
         write_options,
+        transformation_context: Dict[str, Any] = None,
     ):
         if not feature_group.online_enabled and not feature_group.stream:
             raise exceptions.FeatureStoreException(
@@ -376,7 +386,9 @@ def insert_stream(
 
         if feature_group.transformation_functions:
             dataframe = engine.get_instance()._apply_transformation_function(
-                feature_group.transformation_functions, dataframe
+                feature_group.transformation_functions,
+                dataframe,
+                transformation_context=transformation_context,
             )
 
         util.validate_embedding_feature_type(

python/hsfs/core/feature_view_engine.py

@@ -913,6 +913,7 @@ def get_batch_data(
         inference_helper_columns=False,
         dataframe_type="default",
         transformed=True,
+        transformation_context: Dict[str, Any] = None,
     ):
         self._check_feature_group_accessibility(feature_view_obj)
 
@@ -936,7 +937,9 @@ def get_batch_data(
         ).read(read_options=read_options, dataframe_type=dataframe_type)
         if transformation_functions and transformed:
             return engine.get_instance()._apply_transformation_function(
-                transformation_functions, dataset=feature_dataframe
+                transformation_functions,
+                dataset=feature_dataframe,
+                transformation_context=transformation_context,
             )
         else:
             return feature_dataframe