Visitar URL original
fix: Fix shared SQL registry crash - avoid unnecessary UDF deserializ… · feast-dev/feast@3693501 · GitHub
Skip to content

Commit 3693501

Browse files
ntkatholentkathole
committed
fix: Fix shared SQL registry crash - avoid unnecessary UDF deserialization in proto cache building
Signed-off-by: ntkathole <nikhilkathole2683@gmail.com>
1 parent 7279c75 commit 3693501

12 files changed

Lines changed: 417 additions & 89 deletions

File tree

‎sdk/python/feast/feature_view.py‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,7 @@ class FeatureView(BaseFeatureView):
113113
materialization_intervals: List[Tuple[datetime, datetime]]
114114
mode: Optional[Union["TransformationMode", str]]
115115
enable_validation: bool
116+
_raw_feature_transformation_proto: Optional[Message] = None
116117

117118
def __init__(
118119
self,
@@ -481,7 +482,9 @@ def to_proto_spec(
481482
]
482483

483484
feature_transformation_proto = None
484-
if hasattr(self, "feature_transformation") and self.feature_transformation:
485+
if hasattr(self, "_raw_feature_transformation_proto"):
486+
feature_transformation_proto = self._raw_feature_transformation_proto
487+
elif hasattr(self, "feature_transformation") and self.feature_transformation:
485488
feature_transformation_proto = transformation_to_proto(
486489
self.feature_transformation
487490
)
@@ -636,8 +639,14 @@ def _from_proto_internal(
636639
source=source_views if source_views else batch_source, # type: ignore[arg-type]
637640
sink_source=batch_source if source_views else None,
638641
mode=mode,
639-
feature_transformation=transformation,
642+
feature_transformation=transformation
643+
if not skip_udf
644+
else feature_transformation_proto, # type: ignore[arg-type]
640645
)
646+
if skip_udf:
647+
feature_view._raw_feature_transformation_proto = (
648+
feature_transformation_proto
649+
)
641650
else:
642651
mode_from_spec = (
643652
feature_view_proto.spec.mode if feature_view_proto.spec.mode else None

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -368,6 +368,7 @@ def list_stream_feature_views(
368368
project: str,
369369
allow_cache: bool = False,
370370
tags: Optional[dict[str, str]] = None,
371+
skip_udf: bool = False,
371372
) -> List[StreamFeatureView]:
372373
"""
373374
Retrieve a list of stream feature views from the registry
@@ -376,6 +377,7 @@ def list_stream_feature_views(
376377
project: Filter stream feature views based on project name
377378
allow_cache: Whether to allow returning stream feature views from a cached registry
378379
tags: Filter by tags
380+
skip_udf: Skip deserializing UDFs (for metadata-only operations)
379381
380382
Returns:
381383
List of stream feature views
@@ -407,6 +409,7 @@ def list_on_demand_feature_views(
407409
project: str,
408410
allow_cache: bool = False,
409411
tags: Optional[dict[str, str]] = None,
412+
skip_udf: bool = False,
410413
) -> List[OnDemandFeatureView]:
411414
"""
412415
Retrieve a list of on demand feature views from the registry
@@ -415,6 +418,7 @@ def list_on_demand_feature_views(
415418
project: Filter on demand feature views based on project name
416419
allow_cache: Whether to allow returning on demand feature views from a cached registry
417420
tags: Filter by tags
421+
skip_udf: Skip deserializing UDFs (for metadata-only operations)
418422
419423
Returns:
420424
List of on demand feature views
@@ -446,6 +450,7 @@ def list_feature_views(
446450
project: str,
447451
allow_cache: bool = False,
448452
tags: Optional[dict[str, str]] = None,
453+
skip_udf: bool = False,
449454
) -> List[FeatureView]:
450455
"""
451456
Retrieve a list of feature views from the registry
@@ -454,6 +459,7 @@ def list_feature_views(
454459
allow_cache: Allow returning feature views from the cached registry
455460
project: Filter feature views based on project name
456461
tags: Filter by tags
462+
skip_udf: Skip deserializing UDFs (for metadata-only operations)
457463
458464
Returns:
459465
List of feature views
@@ -484,6 +490,7 @@ def list_all_feature_views(
484490
project: str,
485491
allow_cache: bool = False,
486492
tags: Optional[dict[str, str]] = None,
493+
skip_udf: bool = False,
487494
) -> List[BaseFeatureView]:
488495
"""
489496
Retrieve a list of feature views of all types from the registry
@@ -492,6 +499,7 @@ def list_all_feature_views(
492499
allow_cache: Allow returning feature views from the cached registry
493500
project: Filter feature views based on project name
494501
tags: Filter by tags
502+
skip_udf: Skip deserializing UDFs (for metadata-only operations)
495503
496504
Returns:
497505
List of feature views

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

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -130,13 +130,14 @@ def list_all_feature_views(
130130
project: str,
131131
allow_cache: bool = False,
132132
tags: Optional[dict[str, str]] = None,
133+
skip_udf: bool = False,
133134
) -> List[BaseFeatureView]:
134135
if allow_cache:
135136
self._refresh_cached_registry_if_necessary()
136137
return proto_registry_utils.list_all_feature_views(
137-
self.cached_registry_proto, project, tags
138+
self.cached_registry_proto, project, tags, skip_udf=skip_udf
138139
)
139-
return self._list_all_feature_views(project, tags)
140+
return self._list_all_feature_views(project, tags, skip_udf=skip_udf)
140141

141142
@abstractmethod
142143
def _get_feature_view(self, name: str, project: str) -> FeatureView:
@@ -163,13 +164,14 @@ def list_feature_views(
163164
project: str,
164165
allow_cache: bool = False,
165166
tags: Optional[dict[str, str]] = None,
167+
skip_udf: bool = False,
166168
) -> List[FeatureView]:
167169
if allow_cache:
168170
self._refresh_cached_registry_if_necessary()
169171
return proto_registry_utils.list_feature_views(
170-
self.cached_registry_proto, project, tags
172+
self.cached_registry_proto, project, tags, skip_udf=skip_udf
171173
)
172-
return self._list_feature_views(project, tags)
174+
return self._list_feature_views(project, tags, skip_udf=skip_udf)
173175

174176
@abstractmethod
175177
def _get_on_demand_feature_view(
@@ -198,13 +200,14 @@ def list_on_demand_feature_views(
198200
project: str,
199201
allow_cache: bool = False,
200202
tags: Optional[dict[str, str]] = None,
203+
skip_udf: bool = False,
201204
) -> List[OnDemandFeatureView]:
202205
if allow_cache:
203206
self._refresh_cached_registry_if_necessary()
204207
return proto_registry_utils.list_on_demand_feature_views(
205-
self.cached_registry_proto, project, tags
208+
self.cached_registry_proto, project, tags, skip_udf=skip_udf
206209
)
207-
return self._list_on_demand_feature_views(project, tags)
210+
return self._list_on_demand_feature_views(project, tags, skip_udf=skip_udf)
208211

209212
@abstractmethod
210213
def _get_stream_feature_view(self, name: str, project: str) -> StreamFeatureView:
@@ -231,13 +234,14 @@ def list_stream_feature_views(
231234
project: str,
232235
allow_cache: bool = False,
233236
tags: Optional[dict[str, str]] = None,
237+
skip_udf: bool = False,
234238
) -> List[StreamFeatureView]:
235239
if allow_cache:
236240
self._refresh_cached_registry_if_necessary()
237241
return proto_registry_utils.list_stream_feature_views(
238-
self.cached_registry_proto, project, tags
242+
self.cached_registry_proto, project, tags, skip_udf=skip_udf
239243
)
240-
return self._list_stream_feature_views(project, tags)
244+
return self._list_stream_feature_views(project, tags, skip_udf=skip_udf)
241245

242246
@abstractmethod
243247
def _get_feature_service(self, name: str, project: str) -> FeatureService:

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

Lines changed: 32 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -57,15 +57,19 @@ def wrapper(
5757
registry_proto: RegistryProto,
5858
project: str,
5959
tags: Optional[dict[str, str]],
60+
**kwargs,
6061
):
6162
nonlocal cache_key, cache_value
6263

63-
key = tuple([id(registry_proto), registry_proto.version_id, project, tags])
64+
kwargs_key = tuple(sorted(kwargs.items())) if kwargs else ()
65+
key = tuple(
66+
[id(registry_proto), registry_proto.version_id, project, tags, kwargs_key]
67+
)
6468

6569
if key == cache_key:
6670
return cache_value
6771
else:
68-
cache_value = func(registry_proto, project, tags)
72+
cache_value = func(registry_proto, project, tags, **kwargs)
6973
cache_key = key
7074
return cache_value
7175

@@ -279,54 +283,70 @@ def list_feature_services(
279283

280284
@registry_proto_cache_with_tags
281285
def list_all_feature_views(
282-
registry_proto: RegistryProto, project: str, tags: Optional[dict[str, str]]
286+
registry_proto: RegistryProto,
287+
project: str,
288+
tags: Optional[dict[str, str]],
289+
skip_udf: bool = False,
283290
) -> List[BaseFeatureView]:
284291
return (
285-
list_feature_views(registry_proto, project, tags)
286-
+ list_stream_feature_views(registry_proto, project, tags)
287-
+ list_on_demand_feature_views(registry_proto, project, tags)
292+
list_feature_views(registry_proto, project, tags, skip_udf=skip_udf)
293+
+ list_stream_feature_views(registry_proto, project, tags, skip_udf=skip_udf)
294+
+ list_on_demand_feature_views(registry_proto, project, tags, skip_udf=skip_udf)
288295
)
289296

290297

291298
@registry_proto_cache_with_tags
292299
def list_feature_views(
293-
registry_proto: RegistryProto, project: str, tags: Optional[dict[str, str]]
300+
registry_proto: RegistryProto,
301+
project: str,
302+
tags: Optional[dict[str, str]],
303+
skip_udf: bool = False,
294304
) -> List[FeatureView]:
295305
feature_views: List[FeatureView] = []
296306
for feature_view_proto in registry_proto.feature_views:
297307
if feature_view_proto.spec.project == project and utils.has_all_tags(
298308
feature_view_proto.spec.tags, tags
299309
):
300-
feature_views.append(FeatureView.from_proto(feature_view_proto))
310+
feature_views.append(
311+
FeatureView.from_proto(feature_view_proto, skip_udf=skip_udf)
312+
)
301313
return feature_views
302314

303315

304316
@registry_proto_cache_with_tags
305317
def list_stream_feature_views(
306-
registry_proto: RegistryProto, project: str, tags: Optional[dict[str, str]]
318+
registry_proto: RegistryProto,
319+
project: str,
320+
tags: Optional[dict[str, str]],
321+
skip_udf: bool = False,
307322
) -> List[StreamFeatureView]:
308323
stream_feature_views = []
309324
for stream_feature_view in registry_proto.stream_feature_views:
310325
if stream_feature_view.spec.project == project and utils.has_all_tags(
311326
stream_feature_view.spec.tags, tags
312327
):
313328
stream_feature_views.append(
314-
StreamFeatureView.from_proto(stream_feature_view)
329+
StreamFeatureView.from_proto(stream_feature_view, skip_udf=skip_udf)
315330
)
316331
return stream_feature_views
317332

318333

319334
@registry_proto_cache_with_tags
320335
def list_on_demand_feature_views(
321-
registry_proto: RegistryProto, project: str, tags: Optional[dict[str, str]]
336+
registry_proto: RegistryProto,
337+
project: str,
338+
tags: Optional[dict[str, str]],
339+
skip_udf: bool = False,
322340
) -> List[OnDemandFeatureView]:
323341
on_demand_feature_views = []
324342
for on_demand_feature_view in registry_proto.on_demand_feature_views:
325343
if on_demand_feature_view.spec.project == project and utils.has_all_tags(
326344
on_demand_feature_view.spec.tags, tags
327345
):
328346
on_demand_feature_views.append(
329-
OnDemandFeatureView.from_proto(on_demand_feature_view)
347+
OnDemandFeatureView.from_proto(
348+
on_demand_feature_view, skip_udf=skip_udf
349+
)
330350
)
331351
return on_demand_feature_views
332352

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

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -877,25 +877,27 @@ def list_stream_feature_views(
877877
project: str,
878878
allow_cache: bool = False,
879879
tags: Optional[dict[str, str]] = None,
880+
skip_udf: bool = False,
880881
) -> List[StreamFeatureView]:
881882
registry_proto = self._get_registry_proto(
882883
project=project, allow_cache=allow_cache
883884
)
884885
return proto_registry_utils.list_stream_feature_views(
885-
registry_proto, project, tags
886+
registry_proto, project, tags, skip_udf=skip_udf
886887
)
887888

888889
def list_on_demand_feature_views(
889890
self,
890891
project: str,
891892
allow_cache: bool = False,
892893
tags: Optional[dict[str, str]] = None,
894+
skip_udf: bool = False,
893895
) -> List[OnDemandFeatureView]:
894896
registry_proto = self._get_registry_proto(
895897
project=project, allow_cache=allow_cache
896898
)
897899
return proto_registry_utils.list_on_demand_feature_views(
898-
registry_proto, project, tags
900+
registry_proto, project, tags, skip_udf=skip_udf
899901
)
900902

901903
def get_on_demand_feature_view(
@@ -978,12 +980,13 @@ def list_all_feature_views(
978980
project: str,
979981
allow_cache: bool = False,
980982
tags: Optional[dict[str, str]] = None,
983+
skip_udf: bool = False,
981984
) -> List[BaseFeatureView]:
982985
registry_proto = self._get_registry_proto(
983986
project=project, allow_cache=allow_cache
984987
)
985988
return proto_registry_utils.list_all_feature_views(
986-
registry_proto, project, tags
989+
registry_proto, project, tags, skip_udf=skip_udf
987990
)
988991

989992
def get_any_feature_view(
@@ -999,11 +1002,14 @@ def list_feature_views(
9991002
project: str,
10001003
allow_cache: bool = False,
10011004
tags: Optional[dict[str, str]] = None,
1005+
skip_udf: bool = False,
10021006
) -> List[FeatureView]:
10031007
registry_proto = self._get_registry_proto(
10041008
project=project, allow_cache=allow_cache
10051009
)
1006-
return proto_registry_utils.list_feature_views(registry_proto, project, tags)
1010+
return proto_registry_utils.list_feature_views(
1011+
registry_proto, project, tags, skip_udf=skip_udf
1012+
)
10071013

10081014
def get_feature_view(
10091015
self, name: str, project: str, allow_cache: bool = False

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -268,6 +268,7 @@ def list_stream_feature_views(
268268
project: str,
269269
allow_cache: bool = False,
270270
tags: Optional[dict[str, str]] = None,
271+
skip_udf: bool = False,
271272
) -> List[StreamFeatureView]:
272273
request = RegistryServer_pb2.ListStreamFeatureViewsRequest(
273274
project=project, allow_cache=allow_cache, tags=tags
@@ -292,6 +293,7 @@ def list_on_demand_feature_views(
292293
project: str,
293294
allow_cache: bool = False,
294295
tags: Optional[dict[str, str]] = None,
296+
skip_udf: bool = False,
295297
) -> List[OnDemandFeatureView]:
296298
request = RegistryServer_pb2.ListOnDemandFeatureViewsRequest(
297299
project=project, allow_cache=allow_cache, tags=tags
@@ -320,6 +322,7 @@ def list_all_feature_views(
320322
project: str,
321323
allow_cache: bool = False,
322324
tags: Optional[dict[str, str]] = None,
325+
skip_udf: bool = False,
323326
) -> List[BaseFeatureView]:
324327
request = RegistryServer_pb2.ListAllFeatureViewsRequest(
325328
project=project, allow_cache=allow_cache, tags=tags
@@ -347,6 +350,7 @@ def list_feature_views(
347350
project: str,
348351
allow_cache: bool = False,
349352
tags: Optional[dict[str, str]] = None,
353+
skip_udf: bool = False,
350354
) -> List[FeatureView]:
351355
request = RegistryServer_pb2.ListFeatureViewsRequest(
352356
project=project, allow_cache=allow_cache, tags=tags

0 commit comments

Comments
 (0)