fix: Add mode field to Transformation proto for proper serialization

Signed-off-by: ntkathole <[email protected]>

Author: ntkathole

protos/feast/core/FeatureView.proto

@@ -36,7 +36,7 @@ message FeatureView {
     FeatureViewMeta meta = 2;
 }
 
-// Next available id: 16
+// Next available id: 17
 // TODO(adchia): refactor common fields from this and ODFV into separate metadata proto
 message FeatureViewSpec {
     // Name of the feature view. Must be unique. Not updated.
@@ -51,18 +51,9 @@ message FeatureViewSpec {
     // List of specifications for each feature defined as part of this feature view.
     repeated FeatureSpecV2 features = 4;
 
-    // List of specifications for each entity defined as part of this feature view.
-    repeated FeatureSpecV2 entity_columns = 12;
-
-    // Description of the feature view.
-    string description = 10;
-
     // User defined metadata
     map<string,string> tags = 5;
 
-    // Owner of the feature view.
-    string owner = 11;
-
     // Features in this feature view can only be retrieved from online serving
     // younger than ttl. Ttl is measured as the duration of time between
     // the feature's event timestamp and when the feature is retrieved
@@ -71,20 +62,33 @@ message FeatureViewSpec {
 
     // Batch/Offline DataSource where this view can retrieve offline feature data.
     DataSource batch_source = 7;
-    // Streaming DataSource from where this view can consume "online" feature data.
-    DataSource stream_source = 9;
 
     // Whether these features should be served online or not
     // This is also used to determine whether the features should be written to the online store
     bool online = 8;
 
+    // Streaming DataSource from where this view can consume "online" feature data.
+    DataSource stream_source = 9;
+
+    // Description of the feature view.
+    string description = 10;
+
+    // Owner of the feature view.
+    string owner = 11;
+
+    // List of specifications for each entity defined as part of this feature view.
+    repeated FeatureSpecV2 entity_columns = 12;
+
     // Whether these features should be written to the offline store
     bool offline = 13;
 
     repeated FeatureViewSpec source_views = 14;
 
     // Feature transformation for batch feature views
     FeatureTransformationV2 feature_transformation = 15;
+
+    // The transformation mode (e.g., "python", "pandas", "spark", "sql", "ray")
+    string mode = 16;
 }
 
 message FeatureViewMeta {

protos/feast/core/Transformation.proto

@@ -15,6 +15,9 @@ message UserDefinedFunctionV2 {
 
     // The string representation of the udf
     string body_text = 3;
+
+    // The transformation mode (e.g., "python", "pandas", "spark", "sql")
+    string mode = 4;
 }
 
 // A feature transformation executed as a user-defined function

sdk/python/feast/batch_feature_view.py

@@ -136,6 +136,7 @@ def __init__(
             schema=schema,
             source=source,  # type: ignore[arg-type]
             sink_source=sink_source,
+            mode=mode,
         )
 
     def get_feature_transformation(self) -> Optional[Transformation]:

sdk/python/feast/feature_view.py

@@ -39,6 +39,7 @@
 from feast.protos.feast.core.Transformation_pb2 import (
     FeatureTransformationV2 as FeatureTransformationProto,
 )
+from feast.transformation.mode import TransformationMode
 from feast.types import from_value_type
 from feast.value_type import ValueType
 
@@ -102,6 +103,7 @@ class FeatureView(BaseFeatureView):
     tags: Dict[str, str]
     owner: str
     materialization_intervals: List[Tuple[datetime, datetime]]
+    mode: Optional[Union["TransformationMode", str]]
 
     def __init__(
         self,
@@ -117,6 +119,7 @@ def __init__(
         description: str = "",
         tags: Optional[Dict[str, str]] = None,
         owner: str = "",
+        mode: Optional[Union["TransformationMode", str]] = None,
     ):
         """
         Creates a FeatureView object.
@@ -140,6 +143,7 @@ def __init__(
             tags (optional): A dictionary of key-value pairs to store arbitrary metadata.
             owner (optional): The owner of the feature view, typically the email of the
                 primary maintainer.
+            mode (optional): The transformation mode to use (e.g., python, pandas, spark, sql).
 
         Raises:
             ValueError: A field mapping conflicts with an Entity or a Feature.
@@ -252,6 +256,7 @@ def __init__(
         )
         self.online = online
         self.offline = offline
+        self.mode = mode
         self.materialization_intervals = []
 
     def __hash__(self):
@@ -439,6 +444,14 @@ def to_proto_spec(
                     substrait_transformation=transformation_proto,
                 )
 
+        mode_str = ""
+        if self.mode:
+            mode_str = (
+                self.mode.value
+                if isinstance(self.mode, TransformationMode)
+                else self.mode
+            )
+
         return FeatureViewSpecProto(
             name=self.name,
             entities=self.entities,
@@ -454,6 +467,7 @@ def to_proto_spec(
             stream_source=stream_source_proto,
             source_views=source_view_protos,
             feature_transformation=feature_transformation_proto,
+            mode=mode_str,
         )
 
     def to_proto_meta(self):
@@ -527,6 +541,7 @@ def _from_proto_internal(
 
         if has_transformation and cls == FeatureView:
             from feast.batch_feature_view import BatchFeatureView
+            from feast.transformation.factory import get_transformation_class_from_type
             from feast.transformation.python_transformation import PythonTransformation
             from feast.transformation.substrait_transformation import (
                 SubstraitTransformation,
@@ -538,14 +553,30 @@ def _from_proto_internal(
             transformation = None
 
             if feature_transformation_proto.HasField("user_defined_function"):
-                transformation = PythonTransformation.from_proto(
-                    feature_transformation_proto.user_defined_function
-                )
+                udf_proto = feature_transformation_proto.user_defined_function
+                if udf_proto.mode:
+                    try:
+                        transformation_class = get_transformation_class_from_type(
+                            udf_proto.mode
+                        )
+                        transformation = transformation_class.from_proto(udf_proto)
+                    except (ValueError, KeyError):
+                        transformation = PythonTransformation.from_proto(udf_proto)
+                else:
+                    transformation = PythonTransformation.from_proto(udf_proto)
             elif feature_transformation_proto.HasField("substrait_transformation"):
                 transformation = SubstraitTransformation.from_proto(
                     feature_transformation_proto.substrait_transformation
                 )
 
+            mode: Union[TransformationMode, str]
+            if feature_view_proto.spec.mode:
+                mode = feature_view_proto.spec.mode
+            elif transformation and hasattr(transformation, "mode"):
+                mode = transformation.mode
+            else:
+                mode = TransformationMode.PYTHON
+
             feature_view: FeatureView = BatchFeatureView(  # type: ignore[assignment]
                 name=feature_view_proto.spec.name,
                 description=feature_view_proto.spec.description,
@@ -560,9 +591,14 @@ def _from_proto_internal(
                 ),
                 source=source_views if source_views else batch_source,  # type: ignore[arg-type]
                 sink_source=batch_source if source_views else None,
+                mode=mode,
                 feature_transformation=transformation,
             )
         else:
+            mode_from_spec = (
+                feature_view_proto.spec.mode if feature_view_proto.spec.mode else None
+            )
+
             feature_view = cls(  # type: ignore[assignment]
                 name=feature_view_proto.spec.name,
                 description=feature_view_proto.spec.description,
@@ -577,6 +613,7 @@ def _from_proto_internal(
                 ),
                 source=source_views if source_views else batch_source,
                 sink_source=batch_source if source_views else None,
+                mode=mode_from_spec,
             )
         if stream_source:
             feature_view.stream_source = stream_source