Add improved typed casting to BigQuery (#8331)

This commit is contained in:
Beto Dealmeida 2019-10-21 09:19:24 -07:00 committed by GitHub
parent 4cfd6b3ce1
commit cca689b712
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 21 additions and 0 deletions

View File

@ -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
}