Skip to content

Commit 06e0ad8

Browse files
antznette1antznette1
authored
refactor: Centralize feature view object lookup (#5898)
* refactor: centralize feature view object lookup Signed-off-by: antznette1 <Ochiezeanthonette@gmail.com> * fix: address PR review nits and typing Signed-off-by: antznette1 <Ochiezeanthonette@gmail.com> --------- Signed-off-by: antznette1 <Ochiezeanthonette@gmail.com>
1 parent 9fafc26 commit 06e0ad8

File tree

3 files changed

+81
-40
lines changed

3 files changed

+81
-40
lines changed

‎sdk/python/feast/feature_logging.py‎

Lines changed: 13 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,9 @@
88
from feast.embedded_go.type_map import FEAST_TYPE_TO_ARROW_TYPE, PA_TIMESTAMP_TYPE
99
from feast.errors import (
1010
FeastObjectNotFoundException,
11-
FeatureViewNotFoundException,
12-
OnDemandFeatureViewNotFoundException,
1311
)
1412
from feast.feature_view import DUMMY_ENTITY_ID
13+
from feast.feature_view_utils import get_feature_view_from_registry
1514
from feast.protos.feast.core.FeatureService_pb2 import (
1615
LoggingConfig as LoggingConfigProto,
1716
)
@@ -58,25 +57,22 @@ def get_schema(self, registry: "BaseRegistry") -> pa.Schema:
5857
# Go code can be found here:
5958
# https://github.com/feast-dev/feast/blob/master/go/internal/feast/server/logging/memorybuffer.go#L51
6059
try:
61-
feature_view = registry.get_feature_view(projection.name, self._project)
62-
except FeatureViewNotFoundException:
63-
try:
64-
on_demand_feature_view = registry.get_on_demand_feature_view(
65-
projection.name, self._project
66-
)
67-
except OnDemandFeatureViewNotFoundException:
68-
raise FeastObjectNotFoundException(
69-
f"Can't recognize feature view with a name {projection.name}"
70-
)
60+
feast_object = get_feature_view_from_registry(
61+
registry,
62+
projection.name,
63+
self._project,
64+
)
65+
except FeastObjectNotFoundException:
66+
raise FeastObjectNotFoundException(
67+
f"Can't recognize feature view with a name {projection.name}"
68+
)
7169

72-
for (
73-
request_source
74-
) in on_demand_feature_view.source_request_sources.values():
70+
if hasattr(feast_object, "source_request_sources"):
71+
for request_source in feast_object.source_request_sources.values():
7572
for field in request_source.schema:
7673
fields[field.name] = FEAST_TYPE_TO_ARROW_TYPE[field.dtype]
77-
7874
else:
79-
for entity_column in feature_view.entity_columns:
75+
for entity_column in feast_object.entity_columns:
8076
if entity_column.name == DUMMY_ENTITY_ID:
8177
continue
8278

‎sdk/python/feast/feature_server.py‎

Lines changed: 7 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -51,9 +51,9 @@
5151
from feast.data_source import PushMode
5252
from feast.errors import (
5353
FeastError,
54-
FeatureViewNotFoundException,
5554
)
5655
from feast.feast_object import FeastObject
56+
from feast.feature_view_utils import get_feature_view_from_feature_store
5757
from feast.permissions.action import WRITE, AuthzedAction
5858
from feast.permissions.security_manager import assert_permissions
5959
from feast.permissions.server.rest import inject_user_details
@@ -482,27 +482,12 @@ async def _push_with_to(push_to: PushMode) -> None:
482482
async def _get_feast_object(
483483
feature_view_name: str, allow_registry_cache: bool
484484
) -> FeastObject:
485-
# FIXME: this logic repeated at least 3 times in the codebase - should be centralized
486-
# in logging, in server and in feature_store (Python SDK)
487-
try:
488-
return await run_in_threadpool(
489-
store.get_feature_view,
490-
feature_view_name,
491-
allow_registry_cache=allow_registry_cache,
492-
)
493-
except FeatureViewNotFoundException:
494-
try:
495-
return await run_in_threadpool(
496-
store.get_on_demand_feature_view,
497-
feature_view_name,
498-
allow_registry_cache=allow_registry_cache,
499-
)
500-
except FeatureViewNotFoundException:
501-
return await run_in_threadpool(
502-
store.get_stream_feature_view,
503-
feature_view_name,
504-
allow_registry_cache=allow_registry_cache,
505-
)
485+
return await run_in_threadpool(
486+
get_feature_view_from_feature_store,
487+
store,
488+
feature_view_name,
489+
allow_registry_cache,
490+
)
506491

507492
@app.post("/write-to-online-store", dependencies=[Depends(inject_user_details)])
508493
async def write_to_online_store(request: WriteToFeatureStoreRequest) -> None:

‎sdk/python/feast/feature_view_utils.py‎

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,16 +5,29 @@
55
import logging
66
import typing
77
from dataclasses import dataclass
8-
from typing import Callable, Optional
8+
from typing import Callable, Optional, Union
9+
10+
from feast.errors import (
11+
FeastObjectNotFoundException,
12+
FeatureViewNotFoundException,
13+
OnDemandFeatureViewNotFoundException,
14+
)
915

1016
if typing.TYPE_CHECKING:
1117
from feast.data_source import DataSource
18+
from feast.feature_store import FeatureStore
1219
from feast.feature_view import FeatureView
20+
from feast.infra.registry.base_registry import BaseRegistry
21+
from feast.on_demand_feature_view import OnDemandFeatureView
1322
from feast.repo_config import RepoConfig
23+
from feast.stream_feature_view import StreamFeatureView
1424

1525
logger = logging.getLogger(__name__)
1626

1727

28+
FeatureViewLike = Union["FeatureView", "OnDemandFeatureView", "StreamFeatureView"]
29+
30+
1831
@dataclass
1932
class FeatureViewSourceInfo:
2033
"""Information about a feature view's data source resolution."""
@@ -227,3 +240,50 @@ def resolve_feature_view_source_with_fallback(
227240
raise ValueError(
228241
f"Unable to resolve any data source for feature view {feature_view.name}"
229242
)
243+
244+
245+
def get_feature_view_from_feature_store(
246+
store: "FeatureStore",
247+
name: str,
248+
allow_registry_cache: bool = False,
249+
) -> FeatureViewLike:
250+
try:
251+
return store.get_feature_view(name, allow_registry_cache=allow_registry_cache)
252+
except FeatureViewNotFoundException:
253+
try:
254+
return store.get_on_demand_feature_view(
255+
name, allow_registry_cache=allow_registry_cache
256+
)
257+
except (FeatureViewNotFoundException, OnDemandFeatureViewNotFoundException):
258+
try:
259+
return store.get_stream_feature_view(
260+
name, allow_registry_cache=allow_registry_cache
261+
)
262+
except FeatureViewNotFoundException as e:
263+
raise FeastObjectNotFoundException(
264+
f"Can't recognize feast object with a name {name}"
265+
) from e
266+
267+
268+
def get_feature_view_from_registry(
269+
registry: "BaseRegistry",
270+
name: str,
271+
project: str,
272+
allow_cache: bool = False,
273+
) -> FeatureViewLike:
274+
try:
275+
return registry.get_feature_view(name, project, allow_cache=allow_cache)
276+
except FeatureViewNotFoundException:
277+
try:
278+
return registry.get_on_demand_feature_view(
279+
name, project, allow_cache=allow_cache
280+
)
281+
except (FeatureViewNotFoundException, OnDemandFeatureViewNotFoundException):
282+
try:
283+
return registry.get_stream_feature_view(
284+
name, project, allow_cache=allow_cache
285+
)
286+
except FeatureViewNotFoundException as e:
287+
raise FeastObjectNotFoundException(
288+
f"Can't recognize feast object with a name {name}"
289+
) from e

0 commit comments

Comments
 (0)