Merge pull request #1733 from matteocacciola/fix/serialize-dataframe-columns

serialization of columns added into the definition of the table

Author: ArslanSaleem

CONTRIBUTING.md

@@ -65,7 +65,25 @@ make spell_fix
 We use `pytest` to test our code. You can run the tests by running the following command:
 
 ```bash
-make tests
+make test_all
+```
+
+If you prefer, you can run only the core tests with the command:
+
+```bash
+make test_core
+```
+
+or the test of extensions with the command:
+
+```bash
+make test_extensions
+```
+
+You can also run the tests with coverage by running the following command:
+
+```bash
+make test-coverage
 ```
 
 Make sure that all tests pass before submitting a pull request.

pandasai/agent/base.py

@@ -175,6 +175,7 @@ def generate_code_with_retries(self, query: str) -> Any:
                     self._state.logger.log(
                         f"Retrying Code Generation ({attempts}/{max_retries})..."
                     )
+            return None
 
     def execute_with_retries(self, code: str) -> Any:
         """Execute the code with retry logic."""
@@ -194,6 +195,7 @@ def execute_with_retries(self, code: str) -> Any:
                     f"Retrying execution ({attempts}/{max_retries})..."
                 )
                 code = self._regenerate_code_after_error(code, e)
+        return None
 
     def train(
         self,

pandasai/data_loader/semantic_layer_schema.py

@@ -47,9 +47,7 @@ class Column(BaseModel):
     name: str = Field(..., description="Name of the column.")
     type: Optional[str] = Field(None, description="Data type of the column.")
     description: Optional[str] = Field(None, description="Description of the column")
-    expression: Optional[str] = Field(
-        None, description="Aggregation expression (avg, min, max, sum)"
-    )
+    expression: Optional[str] = Field(None, description="Aggregation expression (avg, min, max, sum)")
     alias: Optional[str] = Field(None, description="Alias for the column")
 
     @field_validator("type")
@@ -63,7 +61,9 @@ def is_column_type_supported(cls, type: str) -> str:
 
     @field_validator("expression")
     @classmethod
-    def is_expression_valid(cls, expr: str) -> str:
+    def is_expression_valid(cls, expr: str) -> Optional[str]:
+        if expr is None:
+            return expr
         try:
             parse_one(expr)
             return expr
@@ -73,12 +73,8 @@ def is_expression_valid(cls, expr: str) -> str:
 
 class Relation(BaseModel):
     name: Optional[str] = Field(None, description="Name of the relationship.")
-    description: Optional[str] = Field(
-        None, description="Description of the relationship."
-    )
-    from_: str = Field(
-        ..., alias="from", description="Source column for the relationship."
-    )
+    description: Optional[str] = Field(None, description="Description of the relationship.")
+    from_: str = Field(..., alias="from", description="Source column for the relationship.")
     to: str = Field(..., description="Target column for the relationship.")
 
 

pandasai/helpers/dataframe_serializer.py

@@ -28,6 +28,10 @@ def serialize(cls, df: "DataFrame", dialect: str = "postgres") -> str:
         if df.schema.description is not None:
             dataframe_info += f' description="{df.schema.description}"'
 
+        if df.schema.columns:
+            columns = [column.model_dump() for column in df.schema.columns]
+            dataframe_info += f' columns="{json.dumps(columns, ensure_ascii=False)}"'
+
         dataframe_info += f' dimensions="{df.rows_count}x{df.columns_count}">'
 
         # Truncate long values

tests/unit_tests/helpers/test_dataframe_serializer.py

@@ -1,5 +1,3 @@
-import pandas as pd
-
 from pandasai.helpers.dataframe_serializer import DataframeSerializer
 
 
@@ -8,7 +6,7 @@ def test_serialize_with_name_and_description(self, sample_df):
         """Test serialization with name and description attributes."""
 
         result = DataframeSerializer.serialize(sample_df)
-        expected = """<table dialect="postgres" table_name="table_6c30b42101939c7bdf95f4c1052d615c" dimensions="3x2">
+        expected = """<table dialect="postgres" table_name="table_6c30b42101939c7bdf95f4c1052d615c" columns="[{"name": "A", "type": "integer", "description": null, "expression": null, "alias": null}, {"name": "B", "type": "integer", "description": null, "expression": null, "alias": null}]" dimensions="3x2">
 A,B
 1,4
 2,5
@@ -21,7 +19,7 @@ def test_serialize_with_name_and_description_with_dialect(self, sample_df):
         """Test serialization with name and description attributes."""
 
         result = DataframeSerializer.serialize(sample_df, dialect="mysql")
-        expected = """<table dialect="mysql" table_name="table_6c30b42101939c7bdf95f4c1052d615c" dimensions="3x2">
+        expected = """<table dialect="mysql" table_name="table_6c30b42101939c7bdf95f4c1052d615c" columns="[{"name": "A", "type": "integer", "description": null, "expression": null, "alias": null}, {"name": "B", "type": "integer", "description": null, "expression": null, "alias": null}]" dimensions="3x2">
 A,B
 1,4
 2,5
@@ -44,7 +42,7 @@ def test_serialize_with_dataframe_long_strings(self, sample_df):
         truncated_text = long_text[: DataframeSerializer.MAX_COLUMN_TEXT_LENGTH] + "…"
 
         # Expected output
-        expected = f"""<table dialect="mysql" table_name="table_6c30b42101939c7bdf95f4c1052d615c" dimensions="3x2">
+        expected = f"""<table dialect="mysql" table_name="table_6c30b42101939c7bdf95f4c1052d615c" columns="[{{"name": "A", "type": "integer", "description": null, "expression": null, "alias": null}}, {{"name": "B", "type": "integer", "description": null, "expression": null, "alias": null}}]" dimensions="3x2">
 A,B
 {truncated_text},4
 2,5