Fix get_online_features return schema (#1455)

jklegar · web-flow · commit a65800c7076b · 2021-04-12T10:19:30.000-07:00
Signed-off-by: Jacob Klegar &lt;jacob@tecton.ai&gt;
diff --git a/sdk/python/feast/online_response.py b/sdk/python/feast/online_response.py
@@ -51,7 +51,7 @@ def to_dict(self) -> Dict[str, Any]:
         """
         Converts GetOnlineFeaturesResponse features into a dictionary form.
         """
-        fields = [k for row in self.field_values for k, _ in row.fields.items()]
+        fields = [k for row in self.field_values for k, _ in row.statuses.items()]
         features_dict: Dict[str, List[Any]] = {k: list() for k in fields}
 
         for row in self.field_values:
diff --git a/sdk/python/tests/test_online_retrieval.py b/sdk/python/tests/test_online_retrieval.py
@@ -107,6 +107,14 @@ def test_online() -> None:
         assert result["customer_profile__name"] == ["John", "John"]
         assert result["customer_driver_combined__trips"] == [7, 7]
 
+        # Ensure features are still in result when keys not found
+        result = store.get_online_features(
+            feature_refs=["customer_driver_combined:trips"],
+            entity_rows=[{"driver": 0, "customer": 0}],
+        ).to_dict()
+
+        assert "customer_driver_combined__trips" in result
+
         # invalid table reference
         with pytest.raises(ValueError):
             store.get_online_features(
