[FSTORE-1672] Allow multiple on-demand features to be returned from an on-demand transformation function and allow passing of local variables to a transformation function (logicalclocks#452) (logicalclocks#468)

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,8 @@ def insert(
         storage,
         write_options,
         validation_options: dict = None,
+        transformation_context: Dict[str, Any] = None,
+        transform: bool = True,
     ):
         dataframe_features = engine.get_instance().parse_schema_feature_group(
             feature_dataframe,
@@ -152,16 +160,20 @@ def insert(
         if (
             not isinstance(feature_group, fg.ExternalFeatureGroup)
             and feature_group.transformation_functions
+            and transform
         ):
             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 = (
-            self._update_feature_group_schema_on_demand_transformations(
-                feature_group=feature_group, features=dataframe_features
+            dataframe_features = (
+                self._update_feature_group_schema_on_demand_transformations(
+                    feature_group=feature_group, features=dataframe_features
+                )
             )
-        )
+
         util.validate_embedding_feature_type(
             feature_group.embedding_index, dataframe_features
         )
@@ -361,6 +373,8 @@ def insert_stream(
         timeout,
         checkpoint_dir,
         write_options,
+        transformation_context: Dict[str, Any] = None,
+        transform: bool = True,
     ):
         if not feature_group.online_enabled and not feature_group.stream:
             raise exceptions.FeatureStoreException(
@@ -377,9 +391,11 @@ def insert_stream(
             )
         )
 
-        if feature_group.transformation_functions:
+        if feature_group.transformation_functions and transform:
             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

@@ -392,6 +392,7 @@ def create_training_dataset(
         primary_keys=False,
         event_time=False,
         training_helper_columns=False,
+        transformation_context: Dict[str, Any] = None,
     ):
         self._set_event_time(feature_view_obj, training_dataset_obj)
         updated_instance = self._create_training_data_metadata(
@@ -405,6 +406,7 @@ def create_training_dataset(
             primary_keys=primary_keys,
             event_time=event_time,
             training_helper_columns=training_helper_columns,
+            transformation_context=transformation_context,
         )
         return updated_instance, td_job
 
@@ -420,6 +422,7 @@ def get_training_data(
         event_time=False,
         training_helper_columns=False,
         dataframe_type="default",
+        transformation_context: Dict[str, Any] = None,
     ):
         # check if provided td version has already existed.
         if training_dataset_version:
@@ -497,6 +500,7 @@ def get_training_data(
                 read_options,
                 dataframe_type,
                 training_dataset_version,
+                transformation_context=transformation_context,
             )
             self.compute_training_dataset_statistics(
                 feature_view_obj, td_updated, split_df
@@ -581,6 +585,7 @@ def recreate_training_dataset(
         statistics_config,
         user_write_options,
         spine=None,
+        transformation_context: Dict[str, Any] = None,
     ):
         training_dataset_obj = self._get_training_dataset_metadata(
             feature_view_obj, training_dataset_version
@@ -597,6 +602,7 @@ def recreate_training_dataset(
             user_write_options,
             training_dataset_obj=training_dataset_obj,
             spine=spine,
+            transformation_context=transformation_context,
         )
         # Set training dataset schema after training dataset has been generated
         training_dataset_obj.schema = self.get_training_dataset_schema(
@@ -757,6 +763,7 @@ def compute_training_dataset(
         primary_keys=False,
         event_time=False,
         training_helper_columns=False,
+        transformation_context: Dict[str, Any] = None,
     ):
         if training_dataset_obj:
             pass
@@ -791,6 +798,7 @@ def compute_training_dataset(
             user_write_options,
             self._OVERWRITE,
             feature_view_obj=feature_view_obj,
+            transformation_context=transformation_context,
         )
 
         # Set training dataset schema after training dataset has been generated
@@ -913,6 +921,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 +945,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