Skip to content

Commit ac381b2

Browse files
tokokotokoko
authored
feat: Add registry methods for dealing with all FV types (#4435)
* add new registry method for working with any fv type Signed-off-by: tokoko <[email protected]> * fix: different project for each test in test_universal_registry Signed-off-by: tokoko <[email protected]> * revert project names to project in test_universal_registry Signed-off-by: tokoko <[email protected]> * remove print statements from test_universal_registry Signed-off-by: tokoko <[email protected]> --------- Signed-off-by: tokoko <[email protected]>
1 parent da24656 commit ac381b2

14 files changed

Lines changed: 446 additions & 62 deletions

File tree

‎protos/feast/registry/RegistryServer.proto‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,9 +32,13 @@ service RegistryServer{
3232

3333
// FeatureView RPCs
3434
rpc ApplyFeatureView (ApplyFeatureViewRequest) returns (google.protobuf.Empty) {}
35+
rpc DeleteFeatureView (DeleteFeatureViewRequest) returns (google.protobuf.Empty) {}
36+
rpc GetAnyFeatureView (GetAnyFeatureViewRequest) returns (GetAnyFeatureViewResponse) {}
37+
rpc ListAllFeatureViews (ListAllFeatureViewsRequest) returns (ListAllFeatureViewsResponse) {}
38+
39+
// plain FeatureView RPCs
3540
rpc GetFeatureView (GetFeatureViewRequest) returns (feast.core.FeatureView) {}
3641
rpc ListFeatureViews (ListFeatureViewsRequest) returns (ListFeatureViewsResponse) {}
37-
rpc DeleteFeatureView (DeleteFeatureViewRequest) returns (google.protobuf.Empty) {}
3842

3943
// StreamFeatureView RPCs
4044
rpc GetStreamFeatureView (GetStreamFeatureViewRequest) returns (feast.core.StreamFeatureView) {}
@@ -208,6 +212,35 @@ message DeleteFeatureViewRequest {
208212
bool commit = 3;
209213
}
210214

215+
message AnyFeatureView {
216+
oneof any_feature_view {
217+
feast.core.FeatureView feature_view = 1;
218+
feast.core.OnDemandFeatureView on_demand_feature_view = 2;
219+
feast.core.StreamFeatureView stream_feature_view = 3;
220+
}
221+
}
222+
223+
message GetAnyFeatureViewRequest {
224+
string name = 1;
225+
string project = 2;
226+
bool allow_cache = 3;
227+
}
228+
229+
message GetAnyFeatureViewResponse {
230+
AnyFeatureView any_feature_view = 1;
231+
}
232+
233+
message ListAllFeatureViewsRequest {
234+
string project = 1;
235+
bool allow_cache = 2;
236+
map<string,string> tags = 3;
237+
}
238+
239+
message ListAllFeatureViewsResponse {
240+
repeated AnyFeatureView feature_views = 1;
241+
}
242+
243+
211244
// StreamFeatureView
212245

213246
message GetStreamFeatureViewRequest {

‎sdk/python/feast/cli_utils.py‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -175,7 +175,7 @@ def handle_fv_verbose_permissions_command(
175175
tags=tags_filter # type: ignore[assignment]
176176
)
177177
for fv in feature_views:
178-
if p.match_resource(fv):
178+
if p.match_resource(fv): # type: ignore[arg-type]
179179
feature_views_names.add(fv.name)
180180
if len(feature_views_names) > 0:
181181
Node(
@@ -207,8 +207,7 @@ def handle_not_verbose_permissions_command(
207207
def fetch_all_feast_objects(store: FeatureStore) -> list[FeastObject]:
208208
objects: list[FeastObject] = []
209209
objects.extend(store.list_entities())
210-
objects.extend(store.list_all_feature_views())
211-
objects.extend(store.list_batch_feature_views())
210+
objects.extend(store.list_all_feature_views()) # type: ignore[arg-type]
212211
objects.extend(store.list_feature_services())
213212
objects.extend(store.list_data_sources())
214213
objects.extend(store.list_validation_references())

‎sdk/python/feast/feature_store.py‎

Lines changed: 18 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414
import itertools
15-
import logging
1615
import os
1716
import warnings
1817
from datetime import datetime, timedelta
@@ -247,9 +246,26 @@ def list_feature_services(
247246
"""
248247
return self._registry.list_feature_services(self.project, tags=tags)
249248

249+
def _list_all_feature_views(
250+
self, allow_cache: bool = False, tags: Optional[dict[str, str]] = None
251+
) -> List[BaseFeatureView]:
252+
feature_views = []
253+
for fv in self.registry.list_all_feature_views(
254+
self.project, allow_cache=allow_cache, tags=tags
255+
):
256+
if (
257+
isinstance(fv, FeatureView)
258+
and fv.entities
259+
and fv.entities[0] == DUMMY_ENTITY_NAME
260+
):
261+
fv.entities = []
262+
fv.entity_columns = []
263+
feature_views.append(fv)
264+
return feature_views
265+
250266
def list_all_feature_views(
251267
self, allow_cache: bool = False, tags: Optional[dict[str, str]] = None
252-
) -> List[Union[FeatureView, StreamFeatureView, OnDemandFeatureView]]:
268+
) -> List[BaseFeatureView]:
253269
"""
254270
Retrieves the list of feature views from the registry.
255271
@@ -274,10 +290,6 @@ def list_feature_views(
274290
Returns:
275291
A list of feature views.
276292
"""
277-
logging.warning(
278-
"list_feature_views will make breaking changes. Please use list_batch_feature_views instead. "
279-
"list_feature_views will behave like list_all_feature_views in the future."
280-
)
281293
return utils._list_feature_views(
282294
self._registry, self.project, allow_cache, tags=tags
283295
)
@@ -297,44 +309,6 @@ def list_batch_feature_views(
297309
"""
298310
return self._list_batch_feature_views(allow_cache=allow_cache, tags=tags)
299311

300-
def _list_all_feature_views(
301-
self,
302-
allow_cache: bool = False,
303-
tags: Optional[dict[str, str]] = None,
304-
) -> List[Union[FeatureView, StreamFeatureView, OnDemandFeatureView]]:
305-
all_feature_views = (
306-
utils._list_feature_views(
307-
self._registry, self.project, allow_cache, tags=tags
308-
)
309-
+ self._list_stream_feature_views(allow_cache, tags=tags)
310-
+ self.list_on_demand_feature_views(allow_cache, tags=tags)
311-
)
312-
return all_feature_views
313-
314-
def _list_feature_views(
315-
self,
316-
allow_cache: bool = False,
317-
hide_dummy_entity: bool = True,
318-
tags: Optional[dict[str, str]] = None,
319-
) -> List[FeatureView]:
320-
logging.warning(
321-
"_list_feature_views will make breaking changes. Please use _list_batch_feature_views instead. "
322-
"_list_feature_views will behave like _list_all_feature_views in the future."
323-
)
324-
feature_views = []
325-
for fv in self._registry.list_feature_views(
326-
self.project, allow_cache=allow_cache, tags=tags
327-
):
328-
if (
329-
hide_dummy_entity
330-
and fv.entities
331-
and fv.entities[0] == DUMMY_ENTITY_NAME
332-
):
333-
fv.entities = []
334-
fv.entity_columns = []
335-
feature_views.append(fv)
336-
return feature_views
337-
338312
def _list_batch_feature_views(
339313
self,
340314
allow_cache: bool = False,

‎sdk/python/feast/infra/registry/base_registry.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -391,6 +391,44 @@ def list_feature_views(
391391
"""
392392
raise NotImplementedError
393393

394+
@abstractmethod
395+
def get_any_feature_view(
396+
self, name: str, project: str, allow_cache: bool = False
397+
) -> BaseFeatureView:
398+
"""
399+
Retrieves a feature view of any type.
400+
401+
Args:
402+
name: Name of feature view
403+
project: Feast project that this feature view belongs to
404+
allow_cache: Allow returning feature view from the cached registry
405+
406+
Returns:
407+
Returns either the specified feature view, or raises an exception if
408+
none is found
409+
"""
410+
raise NotImplementedError
411+
412+
@abstractmethod
413+
def list_all_feature_views(
414+
self,
415+
project: str,
416+
allow_cache: bool = False,
417+
tags: Optional[dict[str, str]] = None,
418+
) -> List[BaseFeatureView]:
419+
"""
420+
Retrieve a list of feature views of all types from the registry
421+
422+
Args:
423+
allow_cache: Allow returning feature views from the cached registry
424+
project: Filter feature views based on project name
425+
tags: Filter by tags
426+
427+
Returns:
428+
List of feature views
429+
"""
430+
raise NotImplementedError
431+
394432
@abstractmethod
395433
def apply_materialization(
396434
self,

‎sdk/python/feast/infra/registry/caching_registry.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from threading import Lock
88
from typing import List, Optional
99

10+
from feast.base_feature_view import BaseFeatureView
1011
from feast.data_source import DataSource
1112
from feast.entity import Entity
1213
from feast.feature_service import FeatureService
@@ -102,6 +103,39 @@ def list_entities(
102103
)
103104
return self._list_entities(project, tags)
104105

106+
@abstractmethod
107+
def _get_any_feature_view(self, name: str, project: str) -> BaseFeatureView:
108+
pass
109+
110+
def get_any_feature_view(
111+
self, name: str, project: str, allow_cache: bool = False
112+
) -> BaseFeatureView:
113+
if allow_cache:
114+
self._refresh_cached_registry_if_necessary()
115+
return proto_registry_utils.get_any_feature_view(
116+
self.cached_registry_proto, name, project
117+
)
118+
return self._get_any_feature_view(name, project)
119+
120+
@abstractmethod
121+
def _list_all_feature_views(
122+
self, project: str, tags: Optional[dict[str, str]]
123+
) -> List[BaseFeatureView]:
124+
pass
125+
126+
def list_all_feature_views(
127+
self,
128+
project: str,
129+
allow_cache: bool = False,
130+
tags: Optional[dict[str, str]] = None,
131+
) -> List[BaseFeatureView]:
132+
if allow_cache:
133+
self._refresh_cached_registry_if_necessary()
134+
return proto_registry_utils.list_all_feature_views(
135+
self.cached_registry_proto, project, tags
136+
)
137+
return self._list_all_feature_views(project, tags)
138+
105139
@abstractmethod
106140
def _get_feature_view(self, name: str, project: str) -> FeatureView:
107141
pass

‎sdk/python/feast/infra/registry/proto_registry_utils.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from typing import List, Optional
33

44
from feast import utils
5+
from feast.base_feature_view import BaseFeatureView
56
from feast.data_source import DataSource
67
from feast.entity import Entity
78
from feast.errors import (
@@ -93,6 +94,33 @@ def get_feature_service(
9394
raise FeatureServiceNotFoundException(name, project=project)
9495

9596

97+
def get_any_feature_view(
98+
registry_proto: RegistryProto, name: str, project: str
99+
) -> BaseFeatureView:
100+
for feature_view_proto in registry_proto.feature_views:
101+
if (
102+
feature_view_proto.spec.name == name
103+
and feature_view_proto.spec.project == project
104+
):
105+
return FeatureView.from_proto(feature_view_proto)
106+
107+
for feature_view_proto in registry_proto.stream_feature_views:
108+
if (
109+
feature_view_proto.spec.name == name
110+
and feature_view_proto.spec.project == project
111+
):
112+
return StreamFeatureView.from_proto(feature_view_proto)
113+
114+
for on_demand_feature_view in registry_proto.on_demand_feature_views:
115+
if (
116+
on_demand_feature_view.spec.project == project
117+
and on_demand_feature_view.spec.name == name
118+
):
119+
return OnDemandFeatureView.from_proto(on_demand_feature_view)
120+
121+
raise FeatureViewNotFoundException(name, project)
122+
123+
96124
def get_feature_view(
97125
registry_proto: RegistryProto, name: str, project: str
98126
) -> FeatureView:
@@ -179,6 +207,17 @@ def list_feature_services(
179207
return feature_services
180208

181209

210+
@registry_proto_cache_with_tags
211+
def list_all_feature_views(
212+
registry_proto: RegistryProto, project: str, tags: Optional[dict[str, str]]
213+
) -> List[BaseFeatureView]:
214+
return (
215+
list_feature_views(registry_proto, project, tags)
216+
+ list_stream_feature_views(registry_proto, project, tags)
217+
+ list_on_demand_feature_views(registry_proto, project, tags)
218+
)
219+
220+
182221
@registry_proto_cache_with_tags
183222
def list_feature_views(
184223
registry_proto: RegistryProto, project: str, tags: Optional[dict[str, str]]

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

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -585,7 +585,26 @@ def apply_materialization(
585585
self.commit()
586586
return
587587

588-
raise FeatureViewNotFoundException(feature_view.name, project)
588+
def list_all_feature_views(
589+
self,
590+
project: str,
591+
allow_cache: bool = False,
592+
tags: Optional[dict[str, str]] = None,
593+
) -> List[BaseFeatureView]:
594+
registry_proto = self._get_registry_proto(
595+
project=project, allow_cache=allow_cache
596+
)
597+
return proto_registry_utils.list_all_feature_views(
598+
registry_proto, project, tags
599+
)
600+
601+
def get_any_feature_view(
602+
self, name: str, project: str, allow_cache: bool = False
603+
) -> BaseFeatureView:
604+
registry_proto = self._get_registry_proto(
605+
project=project, allow_cache=allow_cache
606+
)
607+
return proto_registry_utils.get_any_feature_view(registry_proto, name, project)
589608

590609
def list_feature_views(
591610
self,

0 commit comments

Comments
 (0)