Updated schema in rest_api to make returned table format consistent w…

…ith how it is returned when using Haystacks' Document.to_dict(). (#3872)
deepset-ai · Jan 27, 2023 · 8dd3575 · 8dd3575
1 parent e7e3801
commit 8dd3575
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/rest_api/rest_api/schema.py b/rest_api/rest_api/schema.py
@@ -16,7 +16,10 @@
 
 
 BaseConfig.arbitrary_types_allowed = True
-BaseConfig.json_encoders = {np.ndarray: lambda x: x.tolist(), pd.DataFrame: lambda x: x.to_dict(orient="records")}
+BaseConfig.json_encoders = {
+    np.ndarray: lambda x: x.tolist(),
+    pd.DataFrame: lambda x: [x.columns.tolist()] + x.values.tolist(),
+}
 
 
 PrimitiveType = Union[str, int, float, bool]

diff --git a/rest_api/test/test_rest_api.py b/rest_api/test/test_rest_api.py
@@ -436,10 +436,7 @@ def test_query_with_dataframe(client):
         response = client.post(url="/query", json={"query": TEST_QUERY})
         assert 200 == response.status_code
         assert len(response.json()["documents"]) == 1
-        assert response.json()["documents"][0]["content"] == [
-            {"col1": "text_1", "col2": 1},
-            {"col1": "text_2", "col2": 2},
-        ]
+        assert response.json()["documents"][0]["content"] == [["col1", "col2"], ["text_1", 1], ["text_2", 2]]
         assert response.json()["documents"][0]["content_type"] == "table"
         # Ensure `run` was called with the expected parameters
         mocked_pipeline.run.assert_called_with(query=TEST_QUERY, params={}, debug=False)