feat: handle temporal columns in group bys (#16795)

* feat: handle temporal columns in group bys

* Rebase
This commit is contained in:
Beto Dealmeida 2021-09-23 15:18:17 -07:00 committed by GitHub
parent 6921d94861
commit 76f0408932
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 28 additions and 3 deletions

View File

@ -35,6 +35,7 @@ from typing import (
Union,
)
import dateutil.parser
import pandas as pd
import sqlalchemy as sa
import sqlparse
@ -1416,15 +1417,25 @@ class SqlaTable(Model, BaseDatasource): # pylint: disable=too-many-public-metho
)
return ob
@staticmethod
def _get_top_groups(
df: pd.DataFrame, dimensions: List[str], groupby_exprs: Dict[str, Any],
self, df: pd.DataFrame, dimensions: List[str], groupby_exprs: Dict[str, Any],
) -> ColumnElement:
column_map = {column.column_name: column for column in self.columns}
groups = []
for _unused, row in df.iterrows():
group = []
for dimension in dimensions:
group.append(groupby_exprs[dimension] == row[dimension])
value = row[dimension]
# Some databases like Druid will return timestamps as strings, but
# do not perform automatic casting when comparing these strings to
# a timestamp. For cases like this we convert the value from a
# string into a timestamp.
if column_map[dimension].is_temporal and isinstance(value, str):
dttm = dateutil.parser.parse(value)
value = text(self.db_engine_spec.convert_dttm("TIMESTAMP", dttm))
group.append(groupby_exprs[dimension] == value)
groups.append(and_(*group))
return or_(*groups)

View File

@ -101,3 +101,17 @@ class DruidEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method
if tt in (utils.TemporalType.DATETIME, utils.TemporalType.TIMESTAMP):
return f"""TIME_PARSE('{dttm.isoformat(timespec="seconds")}')"""
return None
@classmethod
def epoch_to_dttm(cls) -> str:
"""
Convert from number of seconds since the epoch to a timestamp.
"""
return "MILLIS_TO_TIMESTAMP({col} * 1000)"
@classmethod
def epoch_ms_to_dttm(cls) -> str:
"""
Convert from number of milliseconds since the epoch to a timestamp.
"""
return "MILLIS_TO_TIMESTAMP({col})"