From cca689b712efebd0cc806c20dbb0150046bfe3ff Mon Sep 17 00:00:00 2001 From: Beto Dealmeida Date: Mon, 21 Oct 2019 09:19:24 -0700 Subject: [PATCH] Add improved typed casting to BigQuery (#8331) --- superset/db_engine_specs/bigquery.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/superset/db_engine_specs/bigquery.py b/superset/db_engine_specs/bigquery.py index bcb9c2f81..4b31ac505 100644 --- a/superset/db_engine_specs/bigquery.py +++ b/superset/db_engine_specs/bigquery.py @@ -25,6 +25,21 @@ from sqlalchemy import literal_column from superset.db_engine_specs.base import BaseEngineSpec +pandas_dtype_map = { + "STRING": "object", + "BOOLEAN": "bool", + "INTEGER": "Int64", + "FLOAT": "float64", + "TIMESTAMP": "datetime64[ns]", + "DATETIME": "datetime64[ns]", + "DATE": "object", + "BYTES": "object", + "TIME": "object", + "RECORD": "object", + "NUMERIC": "object", +} + + class BigQueryEngineSpec(BaseEngineSpec): """Engine spec for Google's BigQuery @@ -183,3 +198,9 @@ class BigQueryEngineSpec(BaseEngineSpec): if key in kwargs: gbq_kwargs[key] = kwargs[key] pandas_gbq.to_gbq(df, **gbq_kwargs) + + @classmethod + def get_pandas_dtype(cls, cursor_description: List[tuple]) -> Dict[str, str]: + return { + col[0]: pandas_dtype_map.get(col[1], "object") for col in cursor_description + }