Fix: Convert seed's dataframe columns to the datetime object when their type is explicitly set to date / time (#983)

Author: izeigerman

sqlmesh/core/engine_adapter/bigquery.py

@@ -148,9 +148,11 @@ def _create_table_from_df(
         of the table if `replace` is true.
         """
         assert isinstance(df, pd.DataFrame)
-        table = self.__get_bq_table(table_name, columns_to_types or columns_to_types_from_df(df))
+        if columns_to_types is None:
+            columns_to_types = columns_to_types_from_df(df)
+        table = self.__get_bq_table(table_name, columns_to_types)
         self.client.create_table(table, exists_ok=exists)
-        self.__load_pandas_to_table(table, df, replace=replace)
+        self.__load_pandas_to_table(table, df, columns_to_types, replace=replace)
 
     def _insert_append_pandas_df(
         self,
@@ -162,13 +164,16 @@ def _insert_append_pandas_df(
         """
         Appends to a table from a pandas dataframe. Will create the table if it doesn't exist.
         """
-        table = self.__get_bq_table(table_name, columns_to_types or columns_to_types_from_df(df))
-        self.__load_pandas_to_table(table, df, replace=False)
+        if columns_to_types is None:
+            columns_to_types = columns_to_types_from_df(df)
+        table = self.__get_bq_table(table_name, columns_to_types)
+        self.__load_pandas_to_table(table, df, columns_to_types, replace=False)
 
     def __load_pandas_to_table(
         self,
         table: bigquery.Table,
         df: pd.DataFrame,
+        columns_to_types: t.Dict[str, exp.DataType],
         replace: bool = False,
     ) -> BigQueryQueryResult:
         """
@@ -177,7 +182,7 @@ def __load_pandas_to_table(
         """
         from google.cloud import bigquery
 
-        job_config = bigquery.job.LoadJobConfig()
+        job_config = bigquery.job.LoadJobConfig(schema=self.__get_bq_schema(columns_to_types))
         if replace:
             job_config.write_disposition = bigquery.WriteDisposition.WRITE_TRUNCATE
         result = self.client.load_table_from_dataframe(df, table, job_config=job_config).result()
@@ -258,7 +263,7 @@ def _insert_overwrite_by_condition(
 
             temp_bq_table = self.__get_temp_bq_table(table, columns_to_types)
             self.client.create_table(temp_bq_table, exists_ok=False)
-            result = self.__load_pandas_to_table(temp_bq_table, df, replace=False)
+            result = self.__load_pandas_to_table(temp_bq_table, df, columns_to_types, replace=False)
             if result.errors:
                 raise SQLMeshError(result.errors)
 

sqlmesh/core/model/definition.py

@@ -10,6 +10,7 @@
 from pathlib import Path
 from textwrap import indent
 
+import pandas as pd
 from astor import to_source
 from pydantic import Field
 from sqlglot import diff, exp
@@ -27,6 +28,7 @@
 from sqlmesh.core.model.meta import ModelMeta
 from sqlmesh.core.model.seed import Seed, create_seed
 from sqlmesh.core.renderer import ExpressionRenderer, QueryRenderer
+from sqlmesh.utils import str_to_bool
 from sqlmesh.utils.date import TimeLike, make_inclusive, to_datetime
 from sqlmesh.utils.errors import ConfigError, SQLMeshError, raise_config_error
 from sqlmesh.utils.jinja import JinjaMacroRegistry, extract_macro_references
@@ -814,7 +816,24 @@ def render(
         **kwargs: t.Any,
     ) -> t.Generator[QueryOrDF, None, None]:
         self._ensure_hydrated()
-        yield from self.seed.read(batch_size=self.kind.batch_size)
+
+        date_or_time_columns = []
+        bool_columns = []
+        string_columns = []
+        for name, tpe in (self.columns_to_types_ or {}).items():
+            if tpe.this in exp.DataType.TEMPORAL_TYPES:
+                date_or_time_columns.append(name)
+            elif tpe.is_type("boolean"):
+                bool_columns.append(name)
+            elif tpe.this in exp.DataType.TEXT_TYPES:
+                string_columns.append(name)
+
+        for df in self.seed.read(batch_size=self.kind.batch_size):
+            for column in date_or_time_columns:
+                df[column] = pd.to_datetime(df[column])
+            df[bool_columns] = df[bool_columns].apply(lambda i: str_to_bool(str(i)))
+            df[string_columns] = df[string_columns].astype(str)
+            yield df
 
     def text_diff(self, other: Model) -> str:
         if not isinstance(other, SeedModel):

tests/core/engine_adapter/test_bigquery.py

@@ -127,7 +127,9 @@ def test_replace_query_pandas(mocker: MockerFixture):
     load_result.result.return_value = AttributeDict({"errors": None})
     client_mock.load_table_from_dataframe.return_value = load_result
     df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
-    adapter.replace_query("test_table", df, {"a": "int", "b": "int"})
+    adapter.replace_query(
+        "test_table", df, {"a": exp.DataType.build("int"), "b": exp.DataType.build("int")}
+    )
 
     assert execute_mock.call_args_list == []
     assert client_mock.method_calls[0] == [

tests/core/test_model.py

@@ -1,3 +1,4 @@
+from datetime import date
 from pathlib import Path
 
 import pytest
@@ -588,6 +589,50 @@ def test_seed_pre_statements_only():
     assert not model.post_statements
 
 
+def test_seed_model_custom_types(tmp_path):
+    model_csv_path = (tmp_path / "model.csv").absolute()
+
+    with open(model_csv_path, "w") as fd:
+        fd.write(
+            """key,ds,b_a,b_b,i,i_str
+123,2022-01-01,false,0,321,321
+"""
+        )
+
+    model = create_seed_model(
+        "test_db.test_model",
+        SeedKind(path=str(model_csv_path)),
+        columns={
+            "key": "string",
+            "ds": "date",
+            "b_a": "boolean",
+            "b_b": "boolean",
+            "i": "int",
+            "i_str": "text",
+        },
+    )
+
+    df = next(model.render(context=None))
+
+    assert df["ds"].dtype == "datetime64[ns]"
+    assert df["ds"].iloc[0].date() == date(2022, 1, 1)
+
+    assert df["key"].dtype == "object"
+    assert df["key"].iloc[0] == "123"
+
+    assert df["b_a"].dtype == "bool"
+    assert not df["b_a"].iloc[0]
+
+    assert df["b_b"].dtype == "bool"
+    assert not df["b_b"].iloc[0]
+
+    assert df["i"].dtype == "int64"
+    assert df["i"].iloc[0] == 321
+
+    assert df["i_str"].dtype == "object"
+    assert df["i_str"].iloc[0] == "321"
+
+
 def test_audits():
     expressions = parse(
         """