fix(chart): fix time comparison error (#13)

* fix: use dttm_col name to normalize dttm_col to fix time comparison error
pinterest · Feb 1, 2023 · b3dacc1 · b3dacc1
1 parent b7a0d61
commit b3dacc1
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 8 deletions.
diff --git a/superset/common/query_context_processor.py b/superset/common/query_context_processor.py
@@ -218,6 +218,7 @@ def get_query_result(self, query_object: QueryObject) -> QueryResult:
     def normalize_df(self, df: pd.DataFrame, query_object: QueryObject) -> pd.DataFrame:
         datasource = self._qc_datasource
         timestamp_format = None
+        dttm_col = None
         if datasource.type == "table":
             dttm_col = datasource.get_column(query_object.granularity)
             if dttm_col:
@@ -228,6 +229,7 @@ def normalize_df(self, df: pd.DataFrame, query_object: QueryObject) -> pd.DataFr
             timestamp_format=timestamp_format,
             offset=datasource.offset,
             time_shift=query_object.time_shift,
+            dttm_col_name=(dttm_col.column_name if dttm_col else DTTM_ALIAS),
         )
 
         if self.enforce_numerical_metrics:

diff --git a/superset/utils/core.py b/superset/utils/core.py
@@ -1809,28 +1809,29 @@ def normalize_dttm_col(
     timestamp_format: Optional[str],
     offset: int,
     time_shift: Optional[timedelta],
+    dttm_col_name: Optional[str] = DTTM_ALIAS,
 ) -> None:
-    if DTTM_ALIAS not in df.columns:
+    if dttm_col_name not in df.columns:
         return
     if timestamp_format in ("epoch_s", "epoch_ms"):
-        dttm_col = df[DTTM_ALIAS]
+        dttm_col = df[dttm_col_name]
         if is_numeric_dtype(dttm_col):
             # Column is formatted as a numeric value
             unit = timestamp_format.replace("epoch_", "")
-            df[DTTM_ALIAS] = pd.to_datetime(
+            df[dttm_col_name] = pd.to_datetime(
                 dttm_col, utc=False, unit=unit, origin="unix", errors="coerce"
             )
         else:
             # Column has already been formatted as a timestamp.
-            df[DTTM_ALIAS] = dttm_col.apply(pd.Timestamp)
+            df[dttm_col_name] = dttm_col.apply(pd.Timestamp)
     else:
-        df[DTTM_ALIAS] = pd.to_datetime(
-            df[DTTM_ALIAS], utc=False, format=timestamp_format, errors="coerce"
+        df[dttm_col_name] = pd.to_datetime(
+            df[dttm_col_name], utc=False, format=timestamp_format, errors="coerce"
         )
     if offset:
-        df[DTTM_ALIAS] += timedelta(hours=offset)
+        df[dttm_col_name] += timedelta(hours=offset)
     if time_shift is not None:
-        df[DTTM_ALIAS] += time_shift
+        df[dttm_col_name] += time_shift
 
 
 def parse_boolean_string(bool_str: Optional[str]) -> bool:

diff --git a/superset/viz.py b/superset/viz.py
@@ -280,6 +280,7 @@ def get_df(self, query_obj: Optional[QueryObjectDict] = None) -> pd.DataFrame:
         self.error_msg = ""
 
         timestamp_format = None
+        granularity_col = None
         if self.datasource.type == "table":
             granularity_col = self.datasource.get_column(query_obj["granularity"])
             if granularity_col:
@@ -304,6 +305,7 @@ def get_df(self, query_obj: Optional[QueryObjectDict] = None) -> pd.DataFrame:
                 timestamp_format=timestamp_format,
                 offset=self.datasource.offset,
                 time_shift=self.time_shift,
+                dttm_col_name=(granularity_col.column_name if granularity_col else DTTM_ALIAS),
             )
 
             if self.enforce_numerical_metrics: