type changes, hope i dont regret them

Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
feast-dev · franciscojavierarceo · Oct 10, 2024 · Sep 24, 2024 · Aug 17, 2024 · Sep 24, 2024
commit b223feb88dd2ef0f9f08d6095bfd5524df5276fc
@@ -84,7 +84,9 @@ def __init__(
             if isinstance(feature_grouping, BaseFeatureView):
                 self.feature_view_projections.append(feature_grouping.projection)
 
-    def infer_features(self, fvs_to_update: Dict[str, FeatureView]):
+    def infer_features(
+        self, fvs_to_update: Dict[str, Union[FeatureView, BaseFeatureView]]
+    ):
         """
         Infers the features for the projections of this feature service, and updates this feature
         service in place.

@@ -893,10 +893,11 @@ def apply(
                     data_sources_set_to_update.add(fv.stream_source)
             if isinstance(fv, OnDemandFeatureView):
                 for source_fvp in fv.source_feature_view_projections:
-                    if fv.source_feature_view_projections[source_fvp].batch_source:
-                        data_sources_set_to_update.add(
-                            fv.source_feature_view_projections[source_fvp].batch_source
-                        )
+                    odfv_batch_source: Optional[DataSource] = (
+                        fv.source_feature_view_projections[source_fvp].batch_source
+                    )
+                    if odfv_batch_source is not None:
+                        data_sources_set_to_update.add(odfv_batch_source)
             else:
                 pass
 
@@ -1018,9 +1019,9 @@ def apply(
         tables_to_delete: List[FeatureView] = (
             views_to_delete + sfvs_to_delete if not partial else []  # type: ignore
         )
-        tables_to_keep: List[FeatureView] = (
-            views_to_update + sfvs_to_update + odfvs_with_writes_to_update
-        )  # type: ignore
+        tables_to_keep: List[
+            Union[FeatureView, StreamFeatureView, OnDemandFeatureView]
+        ] = views_to_update + sfvs_to_update + odfvs_with_writes_to_update  # type: ignore
 
         self._get_provider().update_infra(
             project=self.project,
@@ -1503,12 +1504,7 @@ def write_to_online_store(
                     raise DataFrameSerializationError
 
         provider = self._get_provider()
-        if isinstance(feature_view, OnDemandFeatureView):
-            # TODO: add projection mapping
-            projection_mapping = {}
-            provider.ingest_df(feature_view, df, projection_mapping)
-        else:
-            provider.ingest_df(feature_view, df)
+        provider.ingest_df(feature_view, df)
 
     def write_to_offline_store(
         self,

@@ -209,7 +209,7 @@ def _infer_features_and_entities(
             fv, join_keys, run_inference_for_features, config
         )
 
-    entity_columns = []
+    entity_columns: List[Field] = []
     columns_to_exclude = {
         fv.batch_source.timestamp_field,
         fv.batch_source.created_timestamp_column,

@@ -12,6 +12,7 @@
 from feast.infra.offline_stores.offline_store import OfflineStore
 from feast.infra.online_stores.online_store import OnlineStore
 from feast.infra.registry.base_registry import BaseRegistry
+from feast.on_demand_feature_view import OnDemandFeatureView
 from feast.repo_config import RepoConfig
 from feast.stream_feature_view import StreamFeatureView
 
@@ -89,7 +90,7 @@ def update(
             Union[BatchFeatureView, StreamFeatureView, FeatureView]
         ],
         views_to_keep: Sequence[
-            Union[BatchFeatureView, StreamFeatureView, FeatureView]
+            Union[BatchFeatureView, StreamFeatureView, FeatureView, OnDemandFeatureView]
         ],
         entities_to_delete: Sequence[Entity],
         entities_to_keep: Sequence[Entity],

@@ -6,6 +6,7 @@
 from tqdm import tqdm
 
 from feast import OnDemandFeatureView, importer
+from feast.base_feature_view import BaseFeatureView
 from feast.batch_feature_view import BatchFeatureView
 from feast.data_source import DataSource
 from feast.entity import Entity
@@ -122,7 +123,7 @@ def update_infra(
         self,
         project: str,
         tables_to_delete: Sequence[FeatureView],
-        tables_to_keep: Sequence[FeatureView],
+        tables_to_keep: Sequence[Union[FeatureView, OnDemandFeatureView]],
         entities_to_delete: Sequence[Entity],
         entities_to_keep: Sequence[Entity],
         partial: bool,
@@ -274,12 +275,14 @@ def retrieve_online_documents(
 
     def ingest_df(
         self,
-        feature_view: FeatureView,
+        feature_view: Union[BaseFeatureView, FeatureView, OnDemandFeatureView],
         df: pd.DataFrame,
         field_mapping: Optional[Dict] = None,
     ):
         table = pa.Table.from_pandas(df)
         if isinstance(feature_view, OnDemandFeatureView):
+            if not field_mapping:
+                field_mapping = {}
             table = _run_pyarrow_field_mapping(table, field_mapping)
             join_keys = {
                 entity.name: entity.dtype.to_value_type()

@@ -8,13 +8,15 @@
 from tqdm import tqdm
 
 from feast import FeatureService, errors
+from feast.base_feature_view import BaseFeatureView
 from feast.data_source import DataSource
 from feast.entity import Entity
 from feast.feature_view import FeatureView
 from feast.importer import import_class
 from feast.infra.infra_object import Infra
 from feast.infra.offline_stores.offline_store import RetrievalJob
 from feast.infra.registry.base_registry import BaseRegistry
+from feast.on_demand_feature_view import OnDemandFeatureView
 from feast.online_response import OnlineResponse
 from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto
 from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
@@ -47,7 +49,7 @@ def update_infra(
         self,
         project: str,
         tables_to_delete: Sequence[FeatureView],
-        tables_to_keep: Sequence[FeatureView],
+        tables_to_keep: Sequence[Union[FeatureView, OnDemandFeatureView]],
         entities_to_delete: Sequence[Entity],
         entities_to_keep: Sequence[Entity],
         partial: bool,
@@ -125,15 +127,17 @@ def online_write_batch(
 
     def ingest_df(
         self,
-        feature_view: FeatureView,
+        feature_view: Union[BaseFeatureView, FeatureView, OnDemandFeatureView],
         df: pd.DataFrame,
+        field_mapping: Optional[Dict] = None,
     ):
         """
         Persists a dataframe to the online store.
 
         Args:
             feature_view: The feature view to which the dataframe corresponds.
             df: The dataframe to be persisted.
+            field_mapping: A dictionary mapping dataframe column names to feature names.
         """
         pass
 

@@ -207,6 +207,9 @@ def from_feast_to_pyarrow_type(feast_type: FeastType) -> pyarrow.DataType:
     Raises:
         ValueError: The conversion could not be performed.
     """
+    assert isinstance(
+        feast_type, FeastType
+    ), f"Expected FeastType, got {type(feast_type)}"
     if feast_type in FEAST_TYPES_TO_PYARROW_TYPES:
         return FEAST_TYPES_TO_PYARROW_TYPES[feast_type]
 

@@ -228,14 +228,14 @@ def _coerce_datetime(ts):
 
 def _convert_arrow_to_proto(
     table: Union[pyarrow.Table, pyarrow.RecordBatch],
-    feature_view: "FeatureView",
+    feature_view: Union["FeatureView", "OnDemandFeatureView"],
     join_keys: Dict[str, ValueType],
 ) -> List[Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]]:
     # This is a workaround for isinstance(feature_view, OnDemandFeatureView), which triggers a circular import
     if getattr(feature_view, "source_request_sources", None):
-        return _convert_arrow_odfv_to_proto(table, feature_view, join_keys)
+        return _convert_arrow_odfv_to_proto(table, feature_view, join_keys)  # type: ignore[arg-type]
     else:
-        return _convert_arrow_fv_to_proto(table, feature_view, join_keys)
+        return _convert_arrow_fv_to_proto(table, feature_view, join_keys)  # type: ignore[arg-type]
 
 
 def _convert_arrow_fv_to_proto(
@@ -301,7 +301,7 @@ def _convert_arrow_fv_to_proto(
 
 def _convert_arrow_odfv_to_proto(
     table: Union[pyarrow.Table, pyarrow.RecordBatch],
-    feature_view: "FeatureView",
+    feature_view: "OnDemandFeatureView",
     join_keys: Dict[str, ValueType],
 ) -> List[Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]]:
     # Avoid ChunkedArrays which guarantees `zero_copy_only` available.
@@ -1013,7 +1013,7 @@ def _prepare_entities_to_read_from_online_store(
 
     num_rows = _validate_entity_values(entity_proto_values)
 
-    odfv_entities = []
+    odfv_entities: List[Entity] = []
     request_source_keys = []
     for on_demand_feature_view in requested_on_demand_feature_views:
         odfv_entities.append(*getattr(on_demand_feature_view, "entities", None))