test: integration tests for error rate

pandasai/data_loader/loader.py

@@ -19,12 +19,14 @@
 )
 from .query_builder import QueryBuilder
 from .semantic_layer_schema import SemanticLayerSchema
+from .view_query_builder import ViewQueryBuilder
 
 
 class DatasetLoader:
     def __init__(self):
         self.schema: Optional[SemanticLayerSchema] = None
-        self.dataset_path = None
+        self.query_builder: Optional[QueryBuilder] = None
+        self.dataset_path: Optional[str] = None
 
     def load(self, dataset_path: str) -> DataFrame:
         """Load data based on the provided dataset path.
@@ -38,6 +40,11 @@ def load(self, dataset_path: str) -> DataFrame:
         self.dataset_path = dataset_path
         self._load_schema()
 
+        if self.schema.source.view:
+            self.query_builder = ViewQueryBuilder(self.schema)
+        else:
+            self.query_builder = QueryBuilder(self.schema)
+
         source_type = self.schema.source.type
         if source_type in LOCAL_SOURCE_TYPES:
             df = self._load_from_local_source()
@@ -139,13 +146,11 @@ def _load_from_local_source(self) -> pd.DataFrame:
         return self._read_csv_or_parquet(filepath, source_type)
 
     def load_head(self) -> pd.DataFrame:
-        query_builder = QueryBuilder(self.schema)
-        query = query_builder.get_head_query()
+        query = self.query_builder.get_head_query()
         return self.execute_query(query)
 
     def get_row_count(self) -> int:
-        query_builder = QueryBuilder(self.schema)
-        query = query_builder.get_row_count()
+        query = self.query_builder.get_row_count()
         result = self.execute_query(query)
         return result.iloc[0, 0]
 
@@ -154,16 +159,18 @@ def execute_query(self, query: str) -> pd.DataFrame:
         source_type = source.type
         connection_info = source.connection
 
+        formatted_query = self.query_builder.format_query(query)
+
         if not source_type:
             raise ValueError("Source type is missing in the schema.")
 
         load_function = self._get_loader_function(source_type)
 
         try:
-            return load_function(connection_info, query)
+            return load_function(connection_info, formatted_query)
         except Exception as e:
             raise RuntimeError(
-                f"Failed to execute query for source type '{source_type}' with query: {query}"
+                f"Failed to execute query for source type '{source_type}' with query: {formatted_query}"
             ) from e
 
     def _apply_transformations(self, df: pd.DataFrame) -> pd.DataFrame:
@@ -199,5 +206,6 @@ def copy(self) -> "DatasetLoader":
         """
         new_loader = DatasetLoader()
         new_loader.schema = copy.deepcopy(self.schema)
+        new_loader.query_builder = copy.deepcopy(self.query_builder)
         new_loader.dataset_path = self.dataset_path
         return new_loader

pandasai/data_loader/query_builder.py

@@ -1,17 +1,19 @@
 from typing import Any, Dict, List, Union
 
-from pandasai.data_loader.semantic_layer_schema import SemanticLayerSchema
+from pandasai.data_loader.semantic_layer_schema import Relation, SemanticLayerSchema
 
 
 class QueryBuilder:
     def __init__(self, schema: SemanticLayerSchema):
         self.schema = schema
 
+    def format_query(self, query):
+        return query
+
     def build_query(self) -> str:
         columns = self._get_columns()
-        table_name = self._get_table_name()
-        query = f"SELECT {columns} FROM {table_name}"
-
+        query = f"SELECT {columns}"
+        query += self._get_from_statement()
         query += self._add_order_by()
         query += self._add_limit()
 
@@ -23,10 +25,8 @@ def _get_columns(self) -> str:
         else:
             return "*"
 
-    def _get_table_name(self):
-        table_name = self.schema.source.table
-        table_name = table_name.lower()
-        return table_name
+    def _get_from_statement(self):
+        return f" FROM {self.schema.source.table.lower()}"
 
     def _add_order_by(self) -> str:
         if not self.schema.order_by:
@@ -46,13 +46,11 @@ def _add_limit(self, n=None) -> str:
 
     def get_head_query(self, n=5):
         source_type = self.schema.source.type
-        table_name = self._get_table_name()
         columns = self._get_columns()
-
+        query = f"SELECT {columns}"
+        query += self._get_from_statement()
         order_by = "RANDOM()" if source_type in {"sqlite", "postgres"} else "RAND()"
-
-        return f"SELECT {columns} FROM {table_name} ORDER BY {order_by} LIMIT {n}"
+        return f"{query} ORDER BY {order_by} LIMIT {n}"
 
     def get_row_count(self):
-        table_name = self._get_table_name()
-        return f"SELECT COUNT(*) FROM {table_name}"
+        return f"SELECT COUNT(*) {self._get_from_statement()}"

pandasai/data_loader/semantic_layer_schema.py

@@ -168,6 +168,9 @@ def check_columns_relations(self):
                 for column_name in _column_names_in_relations or ()
             }
 
+            if not self.relations:
+                raise ValueError("At least one relation must be defined for view.")
+
             if not all(
                 is_view_column_name(column_name) for column_name in _column_names
             ):

pandasai/data_loader/view_query_builder.py

@@ -0,0 +1,54 @@
+from typing import Any, Dict, List, Union
+
+from pandasai.data_loader.query_builder import QueryBuilder
+from pandasai.data_loader.semantic_layer_schema import Relation, SemanticLayerSchema
+
+
+class ViewQueryBuilder(QueryBuilder):
+    def __init__(self, schema: SemanticLayerSchema):
+        super().__init__(schema)
+
+    def format_query(self, query):
+        return f"{self._get_with_statement()}{query}"
+
+    def build_query(self) -> str:
+        columns = self._get_columns()
+        query = self._get_with_statement()
+        query += f"SELECT {columns}"
+        query += self._get_from_statement()
+        query += self._add_order_by()
+        query += self._add_limit()
+        return query
+
+    def _get_columns(self) -> str:
+        if self.schema.columns:
+            return ", ".join(
+                [f"{col.name.replace('.', '_')}" for col in self.schema.columns]
+            )
+        else:
+            return super()._get_columns()
+
+    def _get_from_statement(self):
+        return f" FROM {self.schema.name}"
+
+    def _get_with_statement(self):
+        relations = self.schema.relations
+        first_table = relations[0].from_.split(".")[0]
+        query = f"WITH {self.schema.name} AS ( SELECT\n"
+
+        if self.schema.columns:
+            query += ", ".join(
+                [
+                    f"{col.name} AS {col.name.replace('.', '_')}"
+                    for col in self.schema.columns
+                ]
+            )
+        else:
+            query += "*"
+
+        query += f"\nFROM {first_table}"
+        for relation in relations:
+            to_table = relation.to.split(".")[0]
+            query += f"\nJOIN {to_table} ON {relation.from_} = {relation.to}"
+        query += ")\n"
+        return query

pandasai/dataframe/virtual_dataframe.py

@@ -36,7 +36,7 @@ def __init__(self, *args, **kwargs):
         if not schema:
             raise VirtualizationError("Schema is required for virtualization!")
 
-        table_name = schema.source.table
+        table_name = schema.source.table or schema.name
 
         description = schema.description