fix: BatchFeatureView transformation should persist in Registry Ser/Deserialization

ntkathole · ntkathole · commit 3364bad2be06 · 2025-10-04T12:06:03.000+05:30
Signed-off-by: ntkathole &lt;nikhilkathole2683@gmail.com&gt;
diff --git a/protos/feast/core/FeatureView.proto b/protos/feast/core/FeatureView.proto
@@ -26,6 +26,7 @@ import "google/protobuf/duration.proto";
 import "google/protobuf/timestamp.proto";
 import "feast/core/DataSource.proto";
 import "feast/core/Feature.proto";
+import "feast/core/Transformation.proto";
 
 message FeatureView {
     // User-specified specifications of this feature view.
@@ -35,7 +36,7 @@ message FeatureView {
     FeatureViewMeta meta = 2;
 }
 
-// Next available id: 13
+// Next available id: 16
 // 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.
@@ -81,6 +82,9 @@ message FeatureViewSpec {
     bool offline = 13;
 
     repeated FeatureViewSpec source_views = 14;
+
+    // Feature transformation for batch feature views
+    FeatureTransformationV2 feature_transformation = 15;
 }
 
 message FeatureViewMeta {
diff --git a/sdk/python/feast/feature_view.py b/sdk/python/feast/feature_view.py
@@ -36,6 +36,9 @@
 from feast.protos.feast.core.FeatureView_pb2 import (
     MaterializationInterval as MaterializationIntervalProto,
 )
+from feast.protos.feast.core.Transformation_pb2 import (
+    FeatureTransformationV2 as FeatureTransformationProto,
+)
 from feast.types import from_value_type
 from feast.value_type import ValueType
 
@@ -415,6 +418,27 @@ def to_proto_spec(
             source_view_protos = [
                 view._to_proto_internal(seen).spec for view in self.source_views
             ]
+
+        feature_transformation_proto = None
+        if hasattr(self, "feature_transformation") and self.feature_transformation:
+            from feast.protos.feast.core.Transformation_pb2 import (
+                SubstraitTransformationV2 as SubstraitTransformationProto,
+            )
+            from feast.protos.feast.core.Transformation_pb2 import (
+                UserDefinedFunctionV2 as UserDefinedFunctionProto,
+            )
+
+            transformation_proto = self.feature_transformation.to_proto()
+
+            if isinstance(transformation_proto, UserDefinedFunctionProto):
+                feature_transformation_proto = FeatureTransformationProto(
+                    user_defined_function=transformation_proto,
+                )
+            elif isinstance(transformation_proto, SubstraitTransformationProto):
+                feature_transformation_proto = FeatureTransformationProto(
+                    substrait_transformation=transformation_proto,
+                )
+
         return FeatureViewSpecProto(
             name=self.name,
             entities=self.entities,
@@ -429,6 +453,7 @@ def to_proto_spec(
             batch_source=batch_source_proto,
             stream_source=stream_source_proto,
             source_views=source_view_protos,
+            feature_transformation=feature_transformation_proto,
         )
 
     def to_proto_meta(self):
@@ -498,21 +523,61 @@ def _from_proto_internal(
             for view_spec in feature_view_proto.spec.source_views
         ]
 
-        feature_view = cls(
-            name=feature_view_proto.spec.name,
-            description=feature_view_proto.spec.description,
-            tags=dict(feature_view_proto.spec.tags),
-            owner=feature_view_proto.spec.owner,
-            online=feature_view_proto.spec.online,
-            offline=feature_view_proto.spec.offline,
-            ttl=(
-                timedelta(days=0)
-                if feature_view_proto.spec.ttl.ToNanoseconds() == 0
-                else feature_view_proto.spec.ttl.ToTimedelta()
-            ),
-            source=source_views if source_views else batch_source,
-            sink_source=batch_source if source_views else None,
-        )
+        has_transformation = feature_view_proto.spec.HasField("feature_transformation")
+
+        if has_transformation and cls == FeatureView:
+            from feast.batch_feature_view import BatchFeatureView
+            from feast.transformation.python_transformation import PythonTransformation
+            from feast.transformation.substrait_transformation import (
+                SubstraitTransformation,
+            )
+
+            feature_transformation_proto = (
+                feature_view_proto.spec.feature_transformation
+            )
+            transformation = None
+
+            if feature_transformation_proto.HasField("user_defined_function"):
+                transformation = PythonTransformation.from_proto(
+                    feature_transformation_proto.user_defined_function
+                )
+            elif feature_transformation_proto.HasField("substrait_transformation"):
+                transformation = SubstraitTransformation.from_proto(
+                    feature_transformation_proto.substrait_transformation
+                )
+
+            feature_view: FeatureView = BatchFeatureView(  # type: ignore[assignment]
+                name=feature_view_proto.spec.name,
+                description=feature_view_proto.spec.description,
+                tags=dict(feature_view_proto.spec.tags),
+                owner=feature_view_proto.spec.owner,
+                online=feature_view_proto.spec.online,
+                offline=feature_view_proto.spec.offline,
+                ttl=(
+                    timedelta(days=0)
+                    if feature_view_proto.spec.ttl.ToNanoseconds() == 0
+                    else feature_view_proto.spec.ttl.ToTimedelta()
+                ),
+                source=source_views if source_views else batch_source,  # type: ignore[arg-type]
+                sink_source=batch_source if source_views else None,
+                feature_transformation=transformation,
+            )
+        else:
+            feature_view = cls(  # type: ignore[assignment]
+                name=feature_view_proto.spec.name,
+                description=feature_view_proto.spec.description,
+                tags=dict(feature_view_proto.spec.tags),
+                owner=feature_view_proto.spec.owner,
+                online=feature_view_proto.spec.online,
+                offline=feature_view_proto.spec.offline,
+                ttl=(
+                    timedelta(days=0)
+                    if feature_view_proto.spec.ttl.ToNanoseconds() == 0
+                    else feature_view_proto.spec.ttl.ToTimedelta()
+                ),
+                source=source_views if source_views else batch_source,
+                sink_source=batch_source if source_views else None,
+            )
         if stream_source:
             feature_view.stream_source = stream_source
 
diff --git a/sdk/python/feast/infra/compute_engines/local/feature_builder.py b/sdk/python/feast/infra/compute_engines/local/feature_builder.py
@@ -70,8 +70,13 @@ def build_dedup_node(self, view, input_node):
 
     def build_transformation_node(self, view, input_nodes):
         transform_config = view.feature_transformation
+        transformation_fn = (
+            transform_config.udf
+            if hasattr(transform_config, "udf")
+            else transform_config
+        )
         node = LocalTransformationNode(
-            "transform", transform_config, self.backend, inputs=input_nodes
+            "transform", transformation_fn, self.backend, inputs=input_nodes
         )
         self.nodes.append(node)
         return node
diff --git a/sdk/python/feast/protos/feast/core/FeatureService_pb2.pyi b/sdk/python/feast/protos/feast/core/FeatureService_pb2.pyi
@@ -71,7 +71,7 @@ class FeatureServiceSpec(google.protobuf.message.Message):
     """Name of Feast project that this Feature Service belongs to."""
     @property
     def features(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[feast.core.FeatureViewProjection_pb2.FeatureViewProjection]:
-        """Represents a projection that's to be applied on top of the FeatureView.
+        """Represents a projection that's to be applied on top of the FeatureView.
         Contains data such as the features to use from a FeatureView.
         """
     @property
diff --git a/sdk/python/feast/protos/feast/core/FeatureView_pb2.py b/sdk/python/feast/protos/feast/core/FeatureView_pb2.py
diff --git a/sdk/python/feast/protos/feast/core/FeatureView_pb2.pyi b/sdk/python/feast/protos/feast/core/FeatureView_pb2.pyi
@@ -20,6 +20,7 @@ import builtins
 import collections.abc
 import feast.core.DataSource_pb2
 import feast.core.Feature_pb2
+import feast.core.Transformation_pb2
 import google.protobuf.descriptor
 import google.protobuf.duration_pb2
 import google.protobuf.internal.containers
@@ -57,7 +58,7 @@ class FeatureView(google.protobuf.message.Message):
 global___FeatureView = FeatureView
 
 class FeatureViewSpec(google.protobuf.message.Message):
-    """Next available id: 13
+    """Next available id: 16
     TODO(adchia): refactor common fields from this and ODFV into separate metadata proto
     """
 
@@ -92,6 +93,7 @@ class FeatureViewSpec(google.protobuf.message.Message):
     ONLINE_FIELD_NUMBER: builtins.int
     OFFLINE_FIELD_NUMBER: builtins.int
     SOURCE_VIEWS_FIELD_NUMBER: builtins.int
+    FEATURE_TRANSFORMATION_FIELD_NUMBER: builtins.int
     name: builtins.str
     """Name of the feature view. Must be unique. Not updated."""
     project: builtins.str
@@ -133,6 +135,9 @@ class FeatureViewSpec(google.protobuf.message.Message):
     """Whether these features should be written to the offline store"""
     @property
     def source_views(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___FeatureViewSpec]: ...
+    @property
+    def feature_transformation(self) -> feast.core.Transformation_pb2.FeatureTransformationV2:
+        """Feature transformation (UDF or Substrait) for batch feature views"""
     def __init__(
         self,
         *,
@@ -150,9 +155,10 @@ class FeatureViewSpec(google.protobuf.message.Message):
         online: builtins.bool = ...,
         offline: builtins.bool = ...,
         source_views: collections.abc.Iterable[global___FeatureViewSpec] | None = ...,
+        feature_transformation: feast.core.Transformation_pb2.FeatureTransformationV2 | None = ...,
     ) -> None: ...
-    def HasField(self, field_name: typing_extensions.Literal["batch_source", b"batch_source", "stream_source", b"stream_source", "ttl", b"ttl"]) -> builtins.bool: ...
-    def ClearField(self, field_name: typing_extensions.Literal["batch_source", b"batch_source", "description", b"description", "entities", b"entities", "entity_columns", b"entity_columns", "features", b"features", "name", b"name", "offline", b"offline", "online", b"online", "owner", b"owner", "project", b"project", "source_views", b"source_views", "stream_source", b"stream_source", "tags", b"tags", "ttl", b"ttl"]) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["batch_source", b"batch_source", "feature_transformation", b"feature_transformation", "stream_source", b"stream_source", "ttl", b"ttl"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["batch_source", b"batch_source", "description", b"description", "entities", b"entities", "entity_columns", b"entity_columns", "feature_transformation", b"feature_transformation", "features", b"features", "name", b"name", "offline", b"offline", "online", b"online", "owner", b"owner", "project", b"project", "source_views", b"source_views", "stream_source", b"stream_source", "tags", b"tags", "ttl", b"ttl"]) -> None: ...
 
 global___FeatureViewSpec = FeatureViewSpec
 
diff --git a/sdk/python/tests/unit/test_feature_views.py b/sdk/python/tests/unit/test_feature_views.py
