[Feature store] Fix trying to read data when it is None (#945)

mlrun · May 19, 2021 · e0c1129 · e0c1129
1 parent a38b6af
commit e0c1129
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/mlrun/feature_store/feature_vector.py b/mlrun/feature_store/feature_vector.py
@@ -327,11 +327,12 @@ def get(self, entity_rows: List[dict], as_list=False):
                     del data[key]
             if not data:
                 data = None
-            requested_columns = self.vector.status.features.keys()
-            actual_columns = data.keys()
-            for column in requested_columns:
-                if column not in actual_columns:
-                    data[column] = None
+            else:
+                requested_columns = self.vector.status.features.keys()
+                actual_columns = data.keys()
+                for column in requested_columns:
+                    if column not in actual_columns:
+                        data[column] = None
             if as_list:
                 data = [
                     result.body[key]

diff --git a/tests/system/feature_store/test_feature_store.py b/tests/system/feature_store/test_feature_store.py
@@ -164,7 +164,7 @@ def test_ingest_and_query(self):
         )  # (*) returns 2 features, label adds 1 feature
         self._get_offline_vector(features, features_size)
 
-        self._logger.debug("Get offline feature vector")
+        self._logger.debug("Get online feature vector")
         self._get_online_features(features, features_size)
 
     def test_feature_set_db(self):