feature(SqlLoader): transformations in SqlLoader (#1569)

Co-authored-by: scaliseraoul-sinaptik <raoul@sinaptik.ai>

Author: scaliseraoul

pandasai/data_loader/loader.py

@@ -1,5 +1,6 @@
 import os
 
+import pandas as pd
 import yaml
 
 from pandasai.dataframe.base import DataFrame
@@ -12,6 +13,7 @@
 )
 from .query_builder import QueryBuilder
 from .semantic_layer_schema import SemanticLayerSchema
+from .transformation_manager import TransformationManager
 from .view_query_builder import ViewQueryBuilder
 
 
@@ -72,16 +74,12 @@ def load(self) -> DataFrame:
         """
         raise MethodNotImplementedError("Loader not instantiated")
 
-    def _build_dataset(
-        self, schema: SemanticLayerSchema, dataset_path: str
-    ) -> DataFrame:
-        self.schema = schema
-        self.dataset_path = dataset_path
-        is_view = schema.view
+    def _apply_transformations(self, df: pd.DataFrame) -> pd.DataFrame:
+        if not self.schema.transformations:
+            return df
 
-        self.query_builder = (
-            ViewQueryBuilder(schema) if is_view else QueryBuilder(schema)
-        )
+        transformation_manager = TransformationManager(df)
+        return transformation_manager.apply_transformations(self.schema.transformations)
 
     def _get_abs_dataset_path(self):
         return os.path.join(find_project_root(), "datasets", self.dataset_path)

pandasai/data_loader/local_loader.py

@@ -69,10 +69,3 @@ def _filter_columns(self, df: pd.DataFrame) -> pd.DataFrame:
         df_columns = df.columns.tolist()
         columns_to_keep = [col for col in df_columns if col in schema_columns]
         return df[columns_to_keep]
-
-    def _apply_transformations(self, df: pd.DataFrame) -> pd.DataFrame:
-        if not self.schema.transformations:
-            return df
-
-        transformation_manager = TransformationManager(df)
-        return transformation_manager.apply_transformations(self.schema.transformations)

pandasai/data_loader/sql_loader.py

@@ -24,7 +24,6 @@ def __init__(self, schema: SemanticLayerSchema, dataset_path: str):
         self.query_builder: QueryBuilder = QueryBuilder(schema)
 
     def load(self) -> VirtualDataFrame:
-        self.query_builder = QueryBuilder(self.schema)
         return VirtualDataFrame(
             schema=self.schema,
             data_loader=SQLDatasetLoader(self.schema, self.dataset_path),
@@ -37,9 +36,11 @@ def execute_query(self, query: str, params: Optional[list] = None) -> pd.DataFra
 
         formatted_query = self.query_builder.format_query(query)
         load_function = self._get_loader_function(source_type)
-
         try:
-            return load_function(connection_info, formatted_query, params)
+            dataframe: pd.DataFrame = load_function(
+                connection_info, formatted_query, params
+            )
+            return self._apply_transformations(dataframe)
         except Exception as e:
             raise RuntimeError(
                 f"Failed to execute query for '{source_type}' with: {formatted_query}"

pandasai/data_loader/transformation_manager.py

@@ -1,9 +1,9 @@
 from typing import Any, List, Optional, Union
 
-import numpy as np
 import pandas as pd
 
 from ..exceptions import UnsupportedTransformation
+from .semantic_layer_schema import Transformation
 
 
 class TransformationManager:
@@ -268,12 +268,12 @@ def format_date(self, column: str, date_format: str) -> "TransformationManager":
             TransformationManager: Self for method chaining
 
         Example:
-            >>> df = pd.DataFrame({"date": ["2024-01-01 12:30:45"]})
+            >>> df = pd.DataFrame({"date": ["2025-01-01 12:30:45"]})
             >>> manager = TransformationManager(df)
             >>> result = manager.format_date("date", "%Y-%m-%d").df
             >>> print(result)
                      date
-            0  2024-01-01
+            0  2025-01-01
         """
         self.df[column] = self.df[column].dt.strftime(date_format)
         return self
@@ -307,28 +307,28 @@ def to_numeric(
         return self
 
     def to_datetime(
-        self, column: str, format: Optional[str] = None, errors: str = "coerce"
+        self, column: str, _format: Optional[str] = None, errors: str = "coerce"
     ) -> "TransformationManager":
         """Convert values in a column to datetime type.
 
         Args:
             column (str): The column to transform
-            format (Optional[str]): Expected date format of the input
+            _format (Optional[str]): Expected date format of the input
             errors (str): How to handle parsing errors
 
         Returns:
             TransformationManager: Self for method chaining
 
         Example:
-            >>> df = pd.DataFrame({"date": ["2024-01-01", "invalid"]})
+            >>> df = pd.DataFrame({"date": ["2025-01-01", "invalid"]})
             >>> manager = TransformationManager(df)
             >>> result = manager.to_datetime("date", errors="coerce").df
             >>> print(result)
                         date
-            0  2024-01-01
+            0  2025-01-01
             1         NaT
         """
-        self.df[column] = pd.to_datetime(self.df[column], format=format, errors=errors)
+        self.df[column] = pd.to_datetime(self.df[column], format=_format, errors=errors)
         return self
 
     def fill_na(self, column: str, value: Any) -> "TransformationManager":
@@ -884,27 +884,20 @@ def rename(self, column: str, new_name: str) -> "TransformationManager":
         return self
 
     def apply_transformations(
-        self, transformations: Optional[List[dict]] = None
+        self, transformations: List[Transformation]
     ) -> pd.DataFrame:
         """Apply a list of transformations to the DataFrame.
 
         Args:
-            transformations (Optional[List[dict]]): List of transformation configurations
+            transformations List[Transformation]: List of transformation configurations
 
         Returns:
             pd.DataFrame: The transformed DataFrame
         """
-        if not transformations:
-            return self.df
 
         for transformation in transformations:
-            # Handle both dict and object transformations
-            if isinstance(transformation, dict):
-                transformation_type = transformation["type"]
-                params = transformation["params"]
-            else:
-                transformation_type = transformation.type
-                params = transformation.params
+            transformation_type = transformation.type
+            params = transformation.params
 
             handler = self.transformation_handlers.get(transformation_type)
             if not handler:

pandasai/dataframe/virtual_dataframe.py

@@ -1,15 +1,14 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, ClassVar
+from typing import TYPE_CHECKING, Optional
 
 import pandas as pd
 
-from pandasai.data_loader.semantic_layer_schema import SemanticLayerSchema
 from pandasai.dataframe.base import DataFrame
 from pandasai.exceptions import VirtualizationError
 
 if TYPE_CHECKING:
-    from pandasai.data_loader.loader import DatasetLoader
+    from pandasai.data_loader.sql_loader import SQLDatasetLoader
 
 
 class VirtualDataFrame(DataFrame):
@@ -25,7 +24,7 @@ class VirtualDataFrame(DataFrame):
     ]
 
     def __init__(self, *args, **kwargs):
-        self._loader: DatasetLoader = kwargs.pop("data_loader", None)
+        self._loader: Optional[SQLDatasetLoader] = kwargs.pop("data_loader", None)
         if not self._loader:
             raise VirtualizationError("Data loader is required for virtualization!")
         self._head = None

tests/unit_tests/data_loader/test_loader.py

@@ -7,7 +7,6 @@
 from pandasai.data_loader.loader import DatasetLoader
 from pandasai.data_loader.local_loader import LocalDatasetLoader
 from pandasai.data_loader.semantic_layer_schema import SemanticLayerSchema
-from pandasai.data_loader.sql_loader import SQLDatasetLoader
 from pandasai.dataframe.base import DataFrame
 from pandasai.exceptions import InvalidDataSourceType
 
@@ -111,92 +110,6 @@ def test_apply_transformations(self, sample_schema):
         assert result.iloc[0]["email"] != "user1@example.com"
         assert result.iloc[0]["timestamp"].tzname() == "UTC"
 
-    def test_load_mysql_source(self, mysql_schema):
-        """Test loading data from a MySQL source creates a VirtualDataFrame and handles queries correctly."""
-        with patch("os.path.exists", return_value=True), patch(
-            "builtins.open", mock_open(read_data=str(mysql_schema.to_yaml()))
-        ), patch(
-            "pandasai.data_loader.sql_loader.SQLDatasetLoader.execute_query"
-        ) as mock_execute_query:
-            # Mock the query results
-            mock_execute_query.return_value = DataFrame(
-                pd.DataFrame(
-                    {
-                        "email": ["test@example.com"],
-                        "first_name": ["John"],
-                        "timestamp": [pd.Timestamp.now()],
-                    }
-                )
-            )
-
-            loader = SQLDatasetLoader(mysql_schema, "test/users")
-            logging.debug("Loading schema from dataset path: %s", loader)
-            result = loader.load()
-
-            # Test that we get a VirtualDataFrame
-            assert isinstance(result, DataFrame)
-            assert result.schema == mysql_schema
-
-            # Test that load_head() works
-            head_result = result.head()
-            assert isinstance(head_result, DataFrame)
-            assert "email" in head_result.columns
-            assert "first_name" in head_result.columns
-            assert "timestamp" in head_result.columns
-
-            # Verify the SQL query was executed correctly
-            mock_execute_query.assert_called_once_with(
-                "SELECT email, first_name, timestamp FROM users ORDER BY RAND() LIMIT 5"
-            )
-
-            # Test executing a custom query
-            custom_query = "SELECT email FROM users WHERE first_name = 'John'"
-            result.execute_sql_query(custom_query)
-            mock_execute_query.assert_called_with(custom_query)
-
-    def test_build_dataset_mysql_schema(self, mysql_schema):
-        """Test loading data from a MySQL schema directly and creates a VirtualDataFrame and handles queries correctly."""
-        with patch("os.path.exists", return_value=True), patch(
-            "builtins.open", mock_open(read_data=str(mysql_schema.to_yaml()))
-        ), patch(
-            "pandasai.data_loader.sql_loader.SQLDatasetLoader.execute_query"
-        ) as mock_execute_query:
-            # Mock the query results
-            mock_execute_query.return_value = DataFrame(
-                pd.DataFrame(
-                    {
-                        "email": ["test@example.com"],
-                        "first_name": ["John"],
-                        "timestamp": [pd.Timestamp.now()],
-                    }
-                )
-            )
-
-            loader = SQLDatasetLoader(mysql_schema, "test/test")
-            logging.debug("Loading schema from dataset path: %s", loader)
-            result = loader.load()
-
-            # Test that we get a VirtualDataFrame
-            assert isinstance(result, DataFrame)
-            assert result.schema == mysql_schema
-
-            # Test that load_head() works
-            head_result = result.head()
-            assert isinstance(head_result, DataFrame)
-            assert "email" in head_result.columns
-            assert "first_name" in head_result.columns
-            assert "timestamp" in head_result.columns
-
-            # Verify the SQL query was executed correctly
-            mock_execute_query.assert_called_once_with(
-                "SELECT email, first_name, timestamp FROM users ORDER BY RAND() LIMIT 5"
-            )
-
-            # Test executing a custom query
-            custom_query = "SELECT email FROM users WHERE first_name = 'John'"
-            result.execute_sql_query(custom_query)
-            mock_execute_query.assert_called_with(custom_query)
-
     def test_build_dataset_csv_schema(self, sample_schema):
         """Test loading data from a CSV schema directly and creates a VirtualDataFrame and handles queries correctly."""
         with patch("os.path.exists", return_value=True), patch(