fix: Correct the returning class proto type of StreamFeatureView to StreamFeatureViewProto instead of FeatureViewProto. (#3843)

shuchu · web-flow · commit 86d62215f233 · 2024-01-24T11:06:54.000-05:00
diff --git a/sdk/python/feast/feature_view.py b/sdk/python/feast/feature_view.py
@@ -17,6 +17,7 @@
 from typing import Dict, List, Optional, Tuple, Type
 
 from google.protobuf.duration_pb2 import Duration
+from google.protobuf.message import Message
 from typeguard import typechecked
 
 from feast import utils
@@ -274,7 +275,7 @@ def ensure_valid(self):
             raise ValueError("Feature view has no entities.")
 
     @property
-    def proto_class(self) -> Type[FeatureViewProto]:
+    def proto_class(self) -> Type[Message]:
         return FeatureViewProto
 
     def with_join_key_map(self, join_key_map: Dict[str, str]):
diff --git a/sdk/python/feast/stream_feature_view.py b/sdk/python/feast/stream_feature_view.py
@@ -3,9 +3,10 @@
 import warnings
 from datetime import datetime, timedelta
 from types import FunctionType
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Type, Union
 
 import dill
+from google.protobuf.message import Message
 from typeguard import typechecked
 
 from feast import flags_helper, utils
@@ -298,6 +299,10 @@ def __copy__(self):
         fv.projection = copy.copy(self.projection)
         return fv
 
+    @property
+    def proto_class(self) -> Type[Message]:
+        return StreamFeatureViewProto
+
 
 def stream_feature_view(
     *,
diff --git a/sdk/python/tests/unit/test_feature_views.py b/sdk/python/tests/unit/test_feature_views.py
@@ -10,6 +10,9 @@
 from feast.feature_view import FeatureView
 from feast.field import Field
 from feast.infra.offline_stores.file_source import FileSource
+from feast.protos.feast.core.StreamFeatureView_pb2 import (
+    StreamFeatureView as StreamFeatureViewProto,
+)
 from feast.protos.feast.types.Value_pb2 import ValueType
 from feast.stream_feature_view import StreamFeatureView, stream_feature_view
 from feast.types import Float32
@@ -277,3 +280,22 @@ def test_hash():
 def test_field_types():
     with pytest.raises(TypeError):
         Field(name="name", dtype=ValueType.INT32)
+
+
+def test_stream_feature_view_proto_type():
+    stream_source = KafkaSource(
+        name="kafka",
+        timestamp_field="event_timestamp",
+        kafka_bootstrap_servers="",
+        message_format=AvroFormat(""),
+        topic="topic",
+        batch_source=FileSource(path="some path"),
+    )
+    sfv = StreamFeatureView(
+        name="test stream featureview proto class",
+        entities=[],
+        ttl=timedelta(days=30),
+        source=stream_source,
+        aggregations=[],
+    )
+    assert sfv.proto_class is StreamFeatureViewProto
