11import copy
22import itertools
3+ import logging
34import os
5+ import threading
46import typing
57import warnings
68from collections import Counter , defaultdict
@@ -1380,6 +1382,96 @@ def _get_online_request_context(
13801382 )
13811383
13821384
1385+ _feature_resolution_cache : Dict [tuple , tuple ] = {}
1386+ _feature_resolution_cache_lock = threading .Lock ()
1387+ _feature_resolution_registry_ts : Optional [datetime ] = None
1388+
1389+ _logger = logging .getLogger (__name__ )
1390+
1391+
1392+ def _get_cached_request_context (
1393+ registry ,
1394+ project : str ,
1395+ features : Union [List [str ], "FeatureService" ],
1396+ full_feature_names : bool ,
1397+ ):
1398+ """Return the output of _get_online_request_context plus resolved ODFV
1399+ entity join keys, using a cache that is invalidated whenever the
1400+ registry refreshes or its cache TTL expires."""
1401+ from feast .feature_service import FeatureService as _FS
1402+
1403+ global _feature_resolution_cache , _feature_resolution_registry_ts
1404+
1405+ registry_ts = getattr (registry , "cached_registry_proto_created" , None )
1406+
1407+ if isinstance (features , _FS ):
1408+ features_key : tuple = ("__fs__" , features .name )
1409+ else :
1410+ features_key = tuple (features )
1411+
1412+ cache_key = (features_key , project , full_feature_names )
1413+
1414+ is_cache_valid = getattr (registry , "is_cache_valid" , None )
1415+ registry_cache_valid = is_cache_valid () if callable (is_cache_valid ) else False
1416+
1417+ if registry_ts is not None and registry_cache_valid :
1418+ with _feature_resolution_cache_lock :
1419+ if registry_ts != _feature_resolution_registry_ts :
1420+ _feature_resolution_cache .clear ()
1421+ _feature_resolution_registry_ts = registry_ts
1422+ _logger .debug ("Feature resolution cache cleared (registry refreshed)" )
1423+ else :
1424+ cached = _feature_resolution_cache .get (cache_key )
1425+ if cached is not None :
1426+ return cached
1427+
1428+ ctx = _get_online_request_context (registry , project , features , full_feature_names )
1429+
1430+ (
1431+ feature_refs ,
1432+ requested_on_demand_feature_views ,
1433+ entity_name_to_join_key_map ,
1434+ entity_type_map ,
1435+ join_keys_set ,
1436+ grouped_refs ,
1437+ requested_result_row_names ,
1438+ needed_request_data ,
1439+ entityless_case ,
1440+ ) = ctx
1441+
1442+ odfv_join_keys : set = set ()
1443+ for on_demand_feature_view in requested_on_demand_feature_views :
1444+ entities_for_odfv = getattr (on_demand_feature_view , "entities" , [])
1445+ if len (entities_for_odfv ) > 0 and isinstance (entities_for_odfv [0 ], str ):
1446+ entities_for_odfv = [
1447+ registry .get_entity (entity_name , project , allow_cache = True )
1448+ for entity_name in entities_for_odfv
1449+ ]
1450+ odfv_join_keys .update (entities_for_odfv )
1451+
1452+ result = (
1453+ feature_refs ,
1454+ requested_on_demand_feature_views ,
1455+ entity_name_to_join_key_map ,
1456+ entity_type_map ,
1457+ join_keys_set | odfv_join_keys ,
1458+ grouped_refs ,
1459+ frozenset (requested_result_row_names ),
1460+ needed_request_data ,
1461+ entityless_case ,
1462+ )
1463+
1464+ registry_ts_after = getattr (registry , "cached_registry_proto_created" , None )
1465+ if registry_ts_after is not None :
1466+ with _feature_resolution_cache_lock :
1467+ if registry_ts_after != _feature_resolution_registry_ts :
1468+ _feature_resolution_cache .clear ()
1469+ _feature_resolution_registry_ts = registry_ts_after
1470+ _feature_resolution_cache [cache_key ] = result
1471+
1472+ return result
1473+
1474+
13831475def _prepare_entities_to_read_from_online_store (
13841476 registry ,
13851477 project ,
@@ -1399,10 +1491,13 @@ def _prepare_entities_to_read_from_online_store(
13991491 entity_type_map ,
14001492 join_keys_set ,
14011493 grouped_refs ,
1402- requested_result_row_names ,
1494+ requested_result_row_names_frozen ,
14031495 needed_request_data ,
14041496 entityless_case ,
1405- ) = _get_online_request_context (registry , project , features , full_feature_names )
1497+ ) = _get_cached_request_context (registry , project , features , full_feature_names )
1498+
1499+ # Mutable copy — downstream code adds join keys to this set.
1500+ requested_result_row_names = set (requested_result_row_names_frozen )
14061501
14071502 # Extract Sequence from RepeatedValue Protobuf.
14081503 entity_value_lists : Dict [str , Union [List [Any ], List [ValueProto ]]] = {
@@ -1424,23 +1519,6 @@ def _prepare_entities_to_read_from_online_store(
14241519
14251520 num_rows = _validate_entity_values (entity_proto_values )
14261521
1427- odfv_entities : List [Entity ] = []
1428- request_source_keys : List [str ] = []
1429- for on_demand_feature_view in requested_on_demand_feature_views :
1430- entities_for_odfv = getattr (on_demand_feature_view , "entities" , [])
1431- if len (entities_for_odfv ) > 0 and isinstance (entities_for_odfv [0 ], str ):
1432- entities_for_odfv = [
1433- registry .get_entity (entity_name , project , allow_cache = True )
1434- for entity_name in entities_for_odfv
1435- ]
1436- odfv_entities .extend (entities_for_odfv )
1437- for source in on_demand_feature_view .source_request_sources :
1438- source_schema = on_demand_feature_view .source_request_sources [source ].schema
1439- for column in source_schema :
1440- request_source_keys .append (column .name )
1441-
1442- join_keys_set .update (set (odfv_entities ))
1443-
14441522 join_key_values : Dict [str , List [ValueProto ]] = {}
14451523 request_data_features : Dict [str , List [ValueProto ]] = {}
14461524 # Entity rows may be either entities or request data.
0 commit comments