Skip to content

Commit b223feb

Browse files
franciscojavierarceofranciscojavierarceo
committed
type changes, hope i dont regret them
Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
1 parent 59940cf commit b223feb

File tree

8 files changed

+34
-25
lines changed

8 files changed

+34
-25
lines changed

‎sdk/python/feast/feature_service.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,9 @@ def __init__(
8484
if isinstance(feature_grouping, BaseFeatureView):
8585
self.feature_view_projections.append(feature_grouping.projection)
8686

87-
def infer_features(self, fvs_to_update: Dict[str, FeatureView]):
87+
def infer_features(
88+
self, fvs_to_update: Dict[str, Union[FeatureView, BaseFeatureView]]
89+
):
8890
"""
8991
Infers the features for the projections of this feature service, and updates this feature
9092
service in place.

‎sdk/python/feast/feature_store.py‎

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -893,10 +893,11 @@ def apply(
893893
data_sources_set_to_update.add(fv.stream_source)
894894
if isinstance(fv, OnDemandFeatureView):
895895
for source_fvp in fv.source_feature_view_projections:
896-
if fv.source_feature_view_projections[source_fvp].batch_source:
897-
data_sources_set_to_update.add(
898-
fv.source_feature_view_projections[source_fvp].batch_source
899-
)
896+
odfv_batch_source: Optional[DataSource] = (
897+
fv.source_feature_view_projections[source_fvp].batch_source
898+
)
899+
if odfv_batch_source is not None:
900+
data_sources_set_to_update.add(odfv_batch_source)
900901
else:
901902
pass
902903

@@ -1018,9 +1019,9 @@ def apply(
10181019
tables_to_delete: List[FeatureView] = (
10191020
views_to_delete + sfvs_to_delete if not partial else [] # type: ignore
10201021
)
1021-
tables_to_keep: List[FeatureView] = (
1022-
views_to_update + sfvs_to_update + odfvs_with_writes_to_update
1023-
) # type: ignore
1022+
tables_to_keep: List[
1023+
Union[FeatureView, StreamFeatureView, OnDemandFeatureView]
1024+
] = views_to_update + sfvs_to_update + odfvs_with_writes_to_update # type: ignore
10241025

10251026
self._get_provider().update_infra(
10261027
project=self.project,
@@ -1503,12 +1504,7 @@ def write_to_online_store(
15031504
raise DataFrameSerializationError
15041505

15051506
provider = self._get_provider()
1506-
if isinstance(feature_view, OnDemandFeatureView):
1507-
# TODO: add projection mapping
1508-
projection_mapping = {}
1509-
provider.ingest_df(feature_view, df, projection_mapping)
1510-
else:
1511-
provider.ingest_df(feature_view, df)
1507+
provider.ingest_df(feature_view, df)
15121508

15131509
def write_to_offline_store(
15141510
self,

‎sdk/python/feast/inference.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -209,7 +209,7 @@ def _infer_features_and_entities(
209209
fv, join_keys, run_inference_for_features, config
210210
)
211211

212-
entity_columns = []
212+
entity_columns: List[Field] = []
213213
columns_to_exclude = {
214214
fv.batch_source.timestamp_field,
215215
fv.batch_source.created_timestamp_column,

‎sdk/python/feast/infra/materialization/batch_materialization_engine.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from feast.infra.offline_stores.offline_store import OfflineStore
1313
from feast.infra.online_stores.online_store import OnlineStore
1414
from feast.infra.registry.base_registry import BaseRegistry
15+
from feast.on_demand_feature_view import OnDemandFeatureView
1516
from feast.repo_config import RepoConfig
1617
from feast.stream_feature_view import StreamFeatureView
1718

@@ -89,7 +90,7 @@ def update(
8990
Union[BatchFeatureView, StreamFeatureView, FeatureView]
9091
],
9192
views_to_keep: Sequence[
92-
Union[BatchFeatureView, StreamFeatureView, FeatureView]
93+
Union[BatchFeatureView, StreamFeatureView, FeatureView, OnDemandFeatureView]
9394
],
9495
entities_to_delete: Sequence[Entity],
9596
entities_to_keep: Sequence[Entity],

‎sdk/python/feast/infra/passthrough_provider.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from tqdm import tqdm
77

88
from feast import OnDemandFeatureView, importer
9+
from feast.base_feature_view import BaseFeatureView
910
from feast.batch_feature_view import BatchFeatureView
1011
from feast.data_source import DataSource
1112
from feast.entity import Entity
@@ -122,7 +123,7 @@ def update_infra(
122123
self,
123124
project: str,
124125
tables_to_delete: Sequence[FeatureView],
125-
tables_to_keep: Sequence[FeatureView],
126+
tables_to_keep: Sequence[Union[FeatureView, OnDemandFeatureView]],
126127
entities_to_delete: Sequence[Entity],
127128
entities_to_keep: Sequence[Entity],
128129
partial: bool,
@@ -274,12 +275,14 @@ def retrieve_online_documents(
274275

275276
def ingest_df(
276277
self,
277-
feature_view: FeatureView,
278+
feature_view: Union[BaseFeatureView, FeatureView, OnDemandFeatureView],
278279
df: pd.DataFrame,
279280
field_mapping: Optional[Dict] = None,
280281
):
281282
table = pa.Table.from_pandas(df)
282283
if isinstance(feature_view, OnDemandFeatureView):
284+
if not field_mapping:
285+
field_mapping = {}
283286
table = _run_pyarrow_field_mapping(table, field_mapping)
284287
join_keys = {
285288
entity.name: entity.dtype.to_value_type()

‎sdk/python/feast/infra/provider.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,13 +8,15 @@
88
from tqdm import tqdm
99

1010
from feast import FeatureService, errors
11+
from feast.base_feature_view import BaseFeatureView
1112
from feast.data_source import DataSource
1213
from feast.entity import Entity
1314
from feast.feature_view import FeatureView
1415
from feast.importer import import_class
1516
from feast.infra.infra_object import Infra
1617
from feast.infra.offline_stores.offline_store import RetrievalJob
1718
from feast.infra.registry.base_registry import BaseRegistry
19+
from feast.on_demand_feature_view import OnDemandFeatureView
1820
from feast.online_response import OnlineResponse
1921
from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto
2022
from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
@@ -47,7 +49,7 @@ def update_infra(
4749
self,
4850
project: str,
4951
tables_to_delete: Sequence[FeatureView],
50-
tables_to_keep: Sequence[FeatureView],
52+
tables_to_keep: Sequence[Union[FeatureView, OnDemandFeatureView]],
5153
entities_to_delete: Sequence[Entity],
5254
entities_to_keep: Sequence[Entity],
5355
partial: bool,
@@ -125,15 +127,17 @@ def online_write_batch(
125127

126128
def ingest_df(
127129
self,
128-
feature_view: FeatureView,
130+
feature_view: Union[BaseFeatureView, FeatureView, OnDemandFeatureView],
129131
df: pd.DataFrame,
132+
field_mapping: Optional[Dict] = None,
130133
):
131134
"""
132135
Persists a dataframe to the online store.
133136
134137
Args:
135138
feature_view: The feature view to which the dataframe corresponds.
136139
df: The dataframe to be persisted.
140+
field_mapping: A dictionary mapping dataframe column names to feature names.
137141
"""
138142
pass
139143

‎sdk/python/feast/types.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -207,6 +207,9 @@ def from_feast_to_pyarrow_type(feast_type: FeastType) -> pyarrow.DataType:
207207
Raises:
208208
ValueError: The conversion could not be performed.
209209
"""
210+
assert isinstance(
211+
feast_type, FeastType
212+
), f"Expected FeastType, got {type(feast_type)}"
210213
if feast_type in FEAST_TYPES_TO_PYARROW_TYPES:
211214
return FEAST_TYPES_TO_PYARROW_TYPES[feast_type]
212215

‎sdk/python/feast/utils.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -228,14 +228,14 @@ def _coerce_datetime(ts):
228228

229229
def _convert_arrow_to_proto(
230230
table: Union[pyarrow.Table, pyarrow.RecordBatch],
231-
feature_view: "FeatureView",
231+
feature_view: Union["FeatureView", "OnDemandFeatureView"],
232232
join_keys: Dict[str, ValueType],
233233
) -> List[Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]]:
234234
# This is a workaround for isinstance(feature_view, OnDemandFeatureView), which triggers a circular import
235235
if getattr(feature_view, "source_request_sources", None):
236-
return _convert_arrow_odfv_to_proto(table, feature_view, join_keys)
236+
return _convert_arrow_odfv_to_proto(table, feature_view, join_keys) # type: ignore[arg-type]
237237
else:
238-
return _convert_arrow_fv_to_proto(table, feature_view, join_keys)
238+
return _convert_arrow_fv_to_proto(table, feature_view, join_keys) # type: ignore[arg-type]
239239

240240

241241
def _convert_arrow_fv_to_proto(
@@ -301,7 +301,7 @@ def _convert_arrow_fv_to_proto(
301301

302302
def _convert_arrow_odfv_to_proto(
303303
table: Union[pyarrow.Table, pyarrow.RecordBatch],
304-
feature_view: "FeatureView",
304+
feature_view: "OnDemandFeatureView",
305305
join_keys: Dict[str, ValueType],
306306
) -> List[Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]]:
307307
# Avoid ChunkedArrays which guarantees `zero_copy_only` available.
@@ -1013,7 +1013,7 @@ def _prepare_entities_to_read_from_online_store(
10131013

10141014
num_rows = _validate_entity_values(entity_proto_values)
10151015

1016-
odfv_entities = []
1016+
odfv_entities: List[Entity] = []
10171017
request_source_keys = []
10181018
for on_demand_feature_view in requested_on_demand_feature_views:
10191019
odfv_entities.append(*getattr(on_demand_feature_view, "entities", None))

0 commit comments

Comments
 (0)