1818
1919import feast .proto .core .FeatureProto ;
2020import feast .proto .core .FeatureViewProto ;
21+ import feast .proto .core .OnDemandFeatureViewProto ;
2122import feast .proto .core .RegistryProto ;
2223import feast .proto .serving .ServingAPIProto ;
2324import feast .serving .exception .SpecRetrievalException ;
2425import java .nio .file .Files ;
2526import java .nio .file .Path ;
27+ import java .util .List ;
28+ import java .util .Map ;
29+ import java .util .function .Function ;
30+ import java .util .stream .Collectors ;
2631
2732public class LocalRegistryRepo implements RegistryRepository {
2833 private final RegistryProto .Registry registry ;
34+ private Map <String , FeatureViewProto .FeatureViewSpec > featureViewNameToSpec ;
35+ private Map <String , OnDemandFeatureViewProto .OnDemandFeatureViewSpec >
36+ onDemandFeatureViewNameToSpec ;
2937
3038 public LocalRegistryRepo (Path localRegistryPath ) {
3139 if (!localRegistryPath .toFile ().exists ()) {
@@ -38,6 +46,26 @@ public LocalRegistryRepo(Path localRegistryPath) {
3846 } catch (final Exception e ) {
3947 throw new RuntimeException (e );
4048 }
49+
50+ final RegistryProto .Registry registry = this .getRegistry ();
51+ List <FeatureViewProto .FeatureViewSpec > featureViewSpecs =
52+ registry .getFeatureViewsList ().stream ()
53+ .map (fv -> fv .getSpec ())
54+ .collect (Collectors .toList ());
55+ featureViewNameToSpec =
56+ featureViewSpecs .stream ()
57+ .collect (
58+ Collectors .toMap (FeatureViewProto .FeatureViewSpec ::getName , Function .identity ()));
59+ List <OnDemandFeatureViewProto .OnDemandFeatureViewSpec > onDemandFeatureViewSpecs =
60+ registry .getOnDemandFeatureViewsList ().stream ()
61+ .map (odfv -> odfv .getSpec ())
62+ .collect (Collectors .toList ());
63+ onDemandFeatureViewNameToSpec =
64+ onDemandFeatureViewSpecs .stream ()
65+ .collect (
66+ Collectors .toMap (
67+ OnDemandFeatureViewProto .OnDemandFeatureViewSpec ::getName ,
68+ Function .identity ()));
4169 }
4270
4371 @ Override
@@ -48,15 +76,12 @@ public RegistryProto.Registry getRegistry() {
4876 @ Override
4977 public FeatureViewProto .FeatureViewSpec getFeatureViewSpec (
5078 String projectName , ServingAPIProto .FeatureReferenceV2 featureReference ) {
51- final RegistryProto .Registry registry = this .getRegistry ();
52- for (final FeatureViewProto .FeatureView featureView : registry .getFeatureViewsList ()) {
53- if (featureView .getSpec ().getName ().equals (featureReference .getFeatureTable ())) {
54- return featureView .getSpec ();
55- }
79+ String featureViewName = featureReference .getFeatureTable ();
80+ if (featureViewNameToSpec .containsKey (featureViewName )) {
81+ return featureViewNameToSpec .get (featureViewName );
5682 }
5783 throw new SpecRetrievalException (
58- String .format (
59- "Unable to find feature view with name: %s" , featureReference .getFeatureTable ()));
84+ String .format ("Unable to find feature view with name: %s" , featureViewName ));
6085 }
6186
6287 @ Override
@@ -75,4 +100,22 @@ public FeatureProto.FeatureSpecV2 getFeatureSpec(
75100 "Unable to find feature with name: %s in feature view: %s" ,
76101 featureReference .getName (), featureReference .getFeatureTable ()));
77102 }
103+
104+ @ Override
105+ public OnDemandFeatureViewProto .OnDemandFeatureViewSpec getOnDemandFeatureViewSpec (
106+ String projectName , ServingAPIProto .FeatureReferenceV2 featureReference ) {
107+ String onDemandFeatureViewName = featureReference .getFeatureTable ();
108+ if (onDemandFeatureViewNameToSpec .containsKey (onDemandFeatureViewName )) {
109+ return onDemandFeatureViewNameToSpec .get (onDemandFeatureViewName );
110+ }
111+ throw new SpecRetrievalException (
112+ String .format (
113+ "Unable to find on demand feature view with name: %s" , onDemandFeatureViewName ));
114+ }
115+
116+ @ Override
117+ public boolean isOnDemandFeatureReference (ServingAPIProto .FeatureReferenceV2 featureReference ) {
118+ String onDemandFeatureViewName = featureReference .getFeatureTable ();
119+ return onDemandFeatureViewNameToSpec .containsKey (onDemandFeatureViewName );
120+ }
78121}
0 commit comments