Skip to content

Commit

Permalink
feat: handle temporal columns in group bys (#16795)
Browse files Browse the repository at this point in the history
* feat: handle temporal columns in group bys

* Rebase
  • Loading branch information
betodealmeida authored Sep 23, 2021
1 parent 6921d94 commit 76f0408
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 3 deletions.
17 changes: 14 additions & 3 deletions superset/connectors/sqla/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
Union,
)

import dateutil.parser
import pandas as pd
import sqlalchemy as sa
import sqlparse
Expand Down Expand Up @@ -1416,15 +1417,25 @@ def _get_series_orderby(
)
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)
Expand Down
14 changes: 14 additions & 0 deletions superset/db_engine_specs/druid.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,3 +101,17 @@ def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]:
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})"

0 comments on commit 76f0408

Please sign in to comment.