[Feature Store] Fix alias handling in OnlineVectorService (#2073)

mlrun · Jun 25, 2022 · 1741729 · 1741729
1 parent 49a79be
commit 1741729
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 9 deletions.
diff --git a/mlrun/feature_store/api.py b/mlrun/feature_store/api.py
@@ -265,9 +265,6 @@ def get_online_feature_service(
 
     # todo: support remote service (using remote nuclio/mlrun function if run_config)
 
-    for old_name in service.vector.get_feature_aliases().keys():
-        if old_name in service.vector.status.features.keys():
-            del service.vector.status.features[old_name]
     return service
 
 

diff --git a/mlrun/feature_store/feature_vector.py b/mlrun/feature_store/feature_vector.py
@@ -488,6 +488,11 @@ def get(self, entity_rows: List[Union[dict, list]], as_list=False):
         for row in entity_rows:
             futures.append(self._controller.emit(row, return_awaitable_result=True))
 
+        requested_columns = list(self.vector.status.features.keys())
+        aliases = self.vector.get_feature_aliases()
+        for i, column in enumerate(requested_columns):
+            requested_columns[i] = aliases.get(column, column)
+
         for future in futures:
             result = future.await_result()
             data = result.body
@@ -497,7 +502,6 @@ def get(self, entity_rows: List[Union[dict, list]], as_list=False):
             if not data:
                 data = None
             else:
-                requested_columns = self.vector.status.features.keys()
                 actual_columns = data.keys()
                 for column in requested_columns:
                     if (
@@ -513,12 +517,9 @@ def get(self, entity_rows: List[Union[dict, list]], as_list=False):
                         data[name] = self._impute_values.get(name, v)
 
             if as_list and data:
-                keys = set(self.vector.status.features.keys()).union(
-                    set(self.vector.get_feature_aliases().values())
-                )
                 data = [
                     data.get(key, None)
-                    for key in keys
+                    for key in requested_columns
                     if key != self.vector.status.label_column
                 ]
             results.append(data)

diff --git a/tests/system/feature_store/test_feature_store.py b/tests/system/feature_store/test_feature_store.py
@@ -2735,13 +2735,15 @@ def test_alias_change(self):
             resp = service.get([{"ticker": "AAPL"}])
             assert resp == [
                 {
+                    "new_alias_for_total_ask": 0.0,
                     "bids_min_1h": math.inf,
                     "bids_max_1h": -math.inf,
-                    "new_alias_for_total_ask": 0.0,
                     "name": "Apple Inc",
                     "exchange": "NASDAQ",
                 }
             ]
+            resp = service.get([{"ticker": "AAPL"}], as_list=True)
+            assert resp == [[0.0, math.inf, -math.inf, "Apple Inc", "NASDAQ"]]
         finally:
             service.close()