From 4f1d735d602be4e31fb3ca2e5e6d0277ee32445c Mon Sep 17 00:00:00 2001 From: Iaroslav Zeigerman Date: Thu, 11 Jul 2024 20:18:20 -0600 Subject: [PATCH] Fix: Avoid type coercion when converting pandas dataframes to agate tables in the dbt adapter (#2897) --- sqlmesh/dbt/util.py | 6 +++--- tests/dbt/test_util.py | 12 ++++++++++++ 2 files changed, 15 insertions(+), 3 deletions(-) create mode 100644 tests/dbt/test_util.py diff --git a/sqlmesh/dbt/util.py b/sqlmesh/dbt/util.py index a06f098d6..8fc6c6ecd 100644 --- a/sqlmesh/dbt/util.py +++ b/sqlmesh/dbt/util.py @@ -13,9 +13,9 @@ def _get_dbt_version() -> t.Tuple[int, int]: DBT_VERSION = _get_dbt_version() if DBT_VERSION < (1, 8): - from dbt.clients.agate_helper import table_from_data, empty_table, as_matrix # type: ignore # noqa: F401 + from dbt.clients.agate_helper import table_from_data_flat, empty_table, as_matrix # type: ignore # noqa: F401 else: - from dbt_common.clients.agate_helper import table_from_data, empty_table, as_matrix # type: ignore # noqa: F401 + from dbt_common.clients.agate_helper import table_from_data_flat, empty_table, as_matrix # type: ignore # noqa: F401 def pandas_to_agate(df: pd.DataFrame) -> agate.Table: @@ -23,4 +23,4 @@ def pandas_to_agate(df: pd.DataFrame) -> agate.Table: Converts a Pandas DataFrame to an Agate Table """ - return table_from_data(df.to_dict(orient="records"), df.columns.tolist()) + return table_from_data_flat(df.to_dict(orient="records"), df.columns.tolist()) diff --git a/tests/dbt/test_util.py b/tests/dbt/test_util.py new file mode 100644 index 000000000..67c8ba8e8 --- /dev/null +++ b/tests/dbt/test_util.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +import pandas as pd +from sqlmesh.dbt.util import pandas_to_agate + + +def test_pandas_to_agate_type_coercion(): + df = pd.DataFrame({"data": ["_2024_01_01", "_2024_01_02", "_2024_01_03"]}) + agate_rows = pandas_to_agate(df).rows + + values = [v[0] for v in agate_rows.values()] + assert values == ["_2024_01_01", "_2024_01_02", "_2024_01_03"]