chore(druid): Standardizing time grain transformations (#17050)

* chore(druid): Standardizing time grain transformations * Update druid_tests.py * Update druid_tests.py Co-authored-by: John Bodley <john.bodley@airbnb.com>
apache · Oct 12, 2021 · 5e85f48 · 5e85f48
1 parent c2e1ab6
commit 5e85f48
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 11 deletions.
diff --git a/superset/db_engine_specs/druid.py b/superset/db_engine_specs/druid.py
@@ -41,26 +41,26 @@ class DruidEngineSpec(BaseEngineSpec):
 
     _time_grain_expressions = {
         None: "{col}",
-        "PT1S": "FLOOR({col} TO SECOND)",
+        "PT1S": "TIME_FLOOR({col}, 'PT1S')",
         "PT5S": "TIME_FLOOR({col}, 'PT5S')",
         "PT30S": "TIME_FLOOR({col}, 'PT30S')",
-        "PT1M": "FLOOR({col} TO MINUTE)",
+        "PT1M": "TIME_FLOOR({col}, 'PT1M')",
         "PT5M": "TIME_FLOOR({col}, 'PT5M')",
         "PT10M": "TIME_FLOOR({col}, 'PT10M')",
         "PT15M": "TIME_FLOOR({col}, 'PT15M')",
         "PT0.5H": "TIME_FLOOR({col}, 'PT30M')",
-        "PT1H": "FLOOR({col} TO HOUR)",
+        "PT1H": "TIME_FLOOR({col}, 'PT1H')",
         "PT6H": "TIME_FLOOR({col}, 'PT6H')",
-        "P1D": "FLOOR({col} TO DAY)",
-        "P1W": "FLOOR({col} TO WEEK)",
-        "P1M": "FLOOR({col} TO MONTH)",
-        "P0.25Y": "FLOOR({col} TO QUARTER)",
-        "P1Y": "FLOOR({col} TO YEAR)",
+        "P1D": "TIME_FLOOR({col}, 'P1D')",
+        "P1W": "TIME_FLOOR({col}, 'P1W')",
+        "P1M": "TIME_FLOOR({col}, 'P1M')",
+        "P0.25Y": "TIME_FLOOR({col}, 'P0.25Y')",
+        "P1Y": "TIME_FLOOR({col}, 'P1Y')",
         "P1W/1970-01-03T00:00:00Z": (
-            "TIMESTAMPADD(DAY, 5, FLOOR(TIMESTAMPADD(DAY, 1, {col}) TO WEEK))"
+            "TIME_SHIFT(TIME_FLOOR(TIME_SHIFT({col}, 'P1D', 1), 'P1W'), 'P1D', 5)"
         ),
         "1969-12-28T00:00:00Z/P1W": (
-            "TIMESTAMPADD(DAY, -1, FLOOR(TIMESTAMPADD(DAY, 1, {col}) TO WEEK))"
+            "TIME_SHIFT(TIME_FLOOR(TIME_SHIFT({col}, 'P1D', 1), 'P1W'), 'P1D', -1)"
         ),
     }
 

diff --git a/tests/integration_tests/db_engine_specs/druid_tests.py b/tests/integration_tests/db_engine_specs/druid_tests.py
@@ -50,8 +50,10 @@ def test_timegrain_expressions(self):
         col = "__time"
         sqla_col = column(col)
         test_cases = {
-            "PT1S": f"FLOOR({col} TO SECOND)",
+            "PT1S": f"TIME_FLOOR({col}, 'PT1S')",
             "PT5M": f"TIME_FLOOR({col}, 'PT5M')",
+            "P1W/1970-01-03T00:00:00Z": f"TIME_SHIFT(TIME_FLOOR(TIME_SHIFT({col}, 'P1D', 1), 'P1W'), 'P1D', 5)",
+            "1969-12-28T00:00:00Z/P1W": f"TIME_SHIFT(TIME_FLOOR(TIME_SHIFT({col}, 'P1D', 1), 'P1W'), 'P1D', -1)",
         }
         for grain, expected in test_cases.items():
             actual = DruidEngineSpec.get_timestamp_expr(