feat(snowflake): support timestamp bucketing

ibis-project · Oct 22, 2023 · a95ffa9 · a95ffa9
1 parent 094881b
commit a95ffa9
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 13 deletions.
diff --git a/ibis/backends/snowflake/registry.py b/ibis/backends/snowflake/registry.py
@@ -257,6 +257,24 @@ def _map_get(t, op):
     return sa.cast(expr, sqla_type)
 
 
+def _timestamp_bucket(t, op):
+    if op.offset is not None:
+        raise com.UnsupportedOperationError(
+            "`offset` is not supported in the Snowflake backend for timestamp bucketing"
+        )
+
+    interval = op.interval
+
+    if not isinstance(interval, ops.Literal):
+        raise com.UnsupportedOperationError(
+            f"Interval must be a literal for the Snowflake backend, got {type(interval)}"
+        )
+
+    return sa.func.time_slice(
+        t.translate(op.arg), interval.value, interval.dtype.unit.name
+    )
+
+
 _TIMESTAMP_UNITS_TO_SCALE = {"s": 0, "ms": 3, "us": 6, "ns": 9}
 
 _SF_POS_INF = sa.func.to_double("Inf")
@@ -471,6 +489,7 @@ def _map_get(t, op):
         ops.TimestampDelta: fixed_arity(
             lambda part, left, right: sa.func.timestampdiff(part, right, left), 3
         ),
+        ops.TimestampBucket: _timestamp_bucket,
     }
 )
 

diff --git a/ibis/backends/tests/test_temporal.py b/ibis/backends/tests/test_temporal.py
@@ -2526,7 +2526,6 @@ def test_delta(con, start, end, unit, expected):
         "oracle",
         "pandas",
         "pyspark",
-        "snowflake",
         "sqlite",
         "trino",
     ],
@@ -2543,16 +2542,24 @@ def test_delta(con, start, end, unit, expected):
         param(
             {"milliseconds": 50},
             "50ms",
-            marks=pytest.mark.notimpl(
-                ["clickhouse"],
-                raises=com.UnsupportedOperationError,
-                reason="clickhouse doesn't support sub-second interval precision",
-            ),
+            marks=[
+                pytest.mark.notimpl(
+                    ["clickhouse"],
+                    raises=com.UnsupportedOperationError,
+                    reason="backend doesn't support sub-second interval precision",
+                ),
+                pytest.mark.notimpl(
+                    ["snowflake"],
+                    raises=sa.exc.ProgrammingError,
+                    reason="snowflake doesn't support sub-second interval precision",
+                ),
+            ],
+            id="milliseconds",
         ),
-        ({"seconds": 2}, "2s"),
-        ({"minutes": 5}, "300s"),
-        ({"hours": 2}, "2h"),
-        ({"days": 2}, "2D"),
+        param({"seconds": 2}, "2s", id="seconds"),
+        param({"minutes": 5}, "300s", id="minutes"),
+        param({"hours": 2}, "2h", id="hours"),
+        param({"days": 2}, "2D", id="days"),
     ],
 )
 def test_timestamp_bucket(backend, kws, pd_freq):
@@ -2573,7 +2580,6 @@ def test_timestamp_bucket(backend, kws, pd_freq):
         "oracle",
         "pandas",
         "pyspark",
-        "snowflake",
         "sqlite",
         "trino",
     ],
@@ -2585,11 +2591,11 @@ def test_timestamp_bucket(backend, kws, pd_freq):
     reason="Druid tests load timestamp_col as a string currently",
 )
 @pytest.mark.notimpl(
-    ["clickhouse", "mssql"],
+    ["clickhouse", "mssql", "snowflake"],
     reason="offset arg not supported",
     raises=com.UnsupportedOperationError,
 )
-@pytest.mark.parametrize("offset_mins", [2, -2])
+@pytest.mark.parametrize("offset_mins", [2, -2], ids=["pos", "neg"])
 def test_timestamp_bucket_offset(backend, offset_mins):
     ts = backend.functional_alltypes.timestamp_col.name("ts")
     expr = ts.bucket(minutes=5, offset=ibis.interval(minutes=offset_mins)).name("ts")