[Feature store] Ignore errors in drop columns in case entity is not s…

…et (#821)
mlrun · Mar 21, 2021 · 074d660 · 074d660
1 parent 5fedf06
commit 074d660
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/mlrun/feature_store/retrieval/local_merger.py b/mlrun/feature_store/retrieval/local_merger.py
@@ -72,7 +72,7 @@ def append_index(key):
                 if field not in index_columns:
                     index_columns.append(field)
 
-            self._result_df.drop(columns=index_columns, inplace=True)
+            self._result_df.drop(columns=index_columns, inplace=True, errors="ignore")
 
         if target:
             is_persistent_vector = self.vector.metadata.name is not None

diff --git a/tests/system/feature_store/test_feature_store.py b/tests/system/feature_store/test_feature_store.py
@@ -110,6 +110,11 @@ def _get_offline_vector(self, features, features_size):
         assert df.shape[1] == columns, "unexpected num of returned df columns"
         resp.to_parquet(str(self.results_path / "query.parquet"))
 
+        # check simple api without join with other df
+        resp = fs.get_offline_features(vector)
+        df = resp.to_dataframe()
+        assert df.shape[1] == features_size, "unexpected num of returned df columns"
+
     def _get_online_features(self, features, features_size):
         # test real-time query
         vector = fs.FeatureVector("my-vec", features)