Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(LogSerialization): dataframe timestamp serialization #751

Merged
merged 3 commits into from
Nov 15, 2023
Merged

fix(LogSerialization): dataframe timestamp serialization #751

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
377b80e
fix(json_serializer): datetime convert to json serialization
ArslanSaleem Nov 14, 2023
d82ecd3
refactor(QueryTracker): add function to convert dataframe to dict
ArslanSaleem Nov 14, 2023
cd87101
chore(tests): remove leftover print statements
ArslanSaleem Nov 14, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 8 additions & 10 deletions pandasai/helpers/query_exec_tracker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,10 @@ def start_new_track(self):
self._query_info = {}
self._func_exec_count: dict = defaultdict(int)

def convert_dataframe_to_dict(self, df):
json_data = json.loads(df.to_json(orient="split", date_format="iso"))
return {"headers": json_data["columns"], "rows": json_data["data"]}

def add_dataframes(self, dfs: List) -> None:
"""
Add used dataframes for the query to query exec tracker
Expand All @@ -96,9 +100,7 @@ def add_dataframes(self, dfs: List) -> None:
"""
for df in dfs:
head = df.head_df
self._dataframes.append(
{"headers": head.columns.tolist(), "rows": head.values.tolist()}
)
self._dataframes.append(self.convert_dataframe_to_dict(head))

def add_step(self, step: dict) -> None:
"""
Expand Down Expand Up @@ -200,13 +202,9 @@ def _format_response(self, result: ResponseType) -> ResponseType:
ResponseType: formatted response output
"""
if result["type"] == "dataframe":
return {
"type": result["type"],
"value": {
"headers": result["value"].columns.tolist(),
"rows": result["value"].values.tolist(),
},
}
df_dict = self.convert_dataframe_to_dict(result["value"])
return {"type": result["type"], "value": df_dict}

elif result["type"] == "plot":
with open(result["value"], "rb") as image_file:
image_data = image_file.read()
Expand Down
30 changes: 30 additions & 0 deletions tests/test_query_tracker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import json
import os
import time
from typing import Optional
Expand All @@ -9,6 +10,7 @@
from pandasai.llm.fake import FakeLLM
from pandasai.smart_dataframe import SmartDataframe
from unittest import TestCase
from datetime import datetime, timedelta


assert_almost_equal = TestCase().assertAlmostEqual
Expand Down Expand Up @@ -120,6 +122,34 @@ def test_format_response_dataframe(
assert len(formatted_response["value"]["headers"]) == 3
assert len(formatted_response["value"]["rows"]) == 10

def test_format_response_dataframe_with_datetime_field(
self, tracker: QueryExecTracker, sample_df: pd.DataFrame
):
# Add a date column with random dates for demonstration
start_date = datetime(2023, 1, 1)
date_range = [start_date + timedelta(days=x) for x in range(len(sample_df))]

sample_df["date"] = date_range

print("Sample", sample_df)
column_dtypes = sample_df.dtypes

# Display the result
print(column_dtypes)
# Create a sample ResponseType for a dataframe
response = {"type": "dataframe", "value": sample_df}

# Format the response using _format_response
formatted_response = tracker._format_response(response)

# Validate dataframe json serialization
json.dumps(formatted_response)

# Check if the response is formatted correctly
assert formatted_response["type"] == "dataframe"
assert len(formatted_response["value"]["headers"]) == 4
assert len(formatted_response["value"]["rows"]) == 10
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new test case test_format_response_dataframe_with_datetime_field is designed to ensure that the _format_response method can handle DataFrames with datetime fields correctly. This is a crucial test since datetime serialization can often be problematic.

However, there are a couple of print statements (lines 134 and 138) that seem to be used for debugging purposes. These should be removed before merging the code to maintain a clean and professional codebase.

-        print("Sample", sample_df)
-        column_dtypes = sample_df.dtypes
-        print(column_dtypes)
Commitable suggestion

[!IMPORTANT]
Carefully review the code before committing. Make sure it correctly replaces the highlighted code, has no missing lines and indentaion issues.

Suggested change
assert len(formatted_response["value"]["headers"]) == 3
assert len(formatted_response["value"]["rows"]) == 10
def test_format_response_dataframe_with_datetime_field(
self, tracker: QueryExecTracker, sample_df: pd.DataFrame
):
# Add a date column with random dates for demonstration
start_date = datetime(2023, 1, 1)
date_range = [start_date + timedelta(days=x) for x in range(len(sample_df))]
sample_df["date"] = date_range
print("Sample", sample_df)
column_dtypes = sample_df.dtypes
# Display the result
print(column_dtypes)
# Create a sample ResponseType for a dataframe
response = {"type": "dataframe", "value": sample_df}
# Format the response using _format_response
formatted_response = tracker._format_response(response)
# Validate dataframe json serialization
json.dumps(formatted_response)
# Check if the response is formatted correctly
assert formatted_response["type"] == "dataframe"
assert len(formatted_response["value"]["headers"]) == 4
assert len(formatted_response["value"]["rows"]) == 10
assert len(formatted_response["value"]["headers"]) == 3
assert len(formatted_response["value"]["rows"]) == 10
def test_format_response_dataframe_with_datetime_field(
self, tracker: QueryExecTracker, sample_df: pd.DataFrame
):
# Add a date column with random dates for demonstration
start_date = datetime(2023, 1, 1)
date_range = [start_date + timedelta(days=x) for x in range(len(sample_df))]
sample_df["date"] = date_range
# Create a sample ResponseType for a dataframe
response = {"type": "dataframe", "value": sample_df}
# Format the response using _format_response
formatted_response = tracker._format_response(response)
# Validate dataframe json serialization
json.dumps(formatted_response)
# Check if the response is formatted correctly
assert formatted_response["type"] == "dataframe"
assert len(formatted_response["value"]["headers"]) == 4
assert len(formatted_response["value"]["rows"]) == 10


def test_format_response_other_type(self, tracker: QueryExecTracker):
# Create a sample ResponseType for a non-dataframe response
response = {
Expand Down