From 73d1c86016f13f30878fbc101b89281921e8469b Mon Sep 17 00:00:00 2001
From: Jim Crist-Harif <jcristharif@gmail.com>
Date: Wed, 18 Oct 2023 17:35:26 -0500
Subject: [PATCH] feat(mssql): add support for timestamp `bucket`

---
 ibis/backends/mssql/registry.py      | 30 +++++++++++++++++++++++++---
 ibis/backends/tests/test_temporal.py |  4 +---
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/ibis/backends/mssql/registry.py b/ibis/backends/mssql/registry.py
index d780bbddfa157..a6a4eb41354b2 100644
--- a/ibis/backends/mssql/registry.py
+++ b/ibis/backends/mssql/registry.py
@@ -83,7 +83,7 @@ def _timestamp_from_unix(x, unit="s"):
     raise com.UnsupportedOperationError(f"{unit!r} unit is not supported!")
 
 
-_truncate_precisions = {
+_interval_units = {
     "us": "microsecond",
     "ms": "millisecond",
     "s": "second",
@@ -100,10 +100,33 @@ def _timestamp_from_unix(x, unit="s"):
 def _timestamp_truncate(t, op):
     arg = t.translate(op.arg)
     unit = op.unit.short
-    if unit not in _truncate_precisions:
+    if unit not in _interval_units:
         raise com.UnsupportedOperationError(f"Unsupported truncate unit {op.unit!r}")
 
-    return sa.func.datetrunc(sa.text(_truncate_precisions[unit]), arg)
+    return sa.func.datetrunc(sa.text(_interval_units[unit]), arg)
+
+
+def _timestamp_bucket(t, op):
+    unit = op.interval.dtype.unit.short
+    if not isinstance(op.interval, ops.Literal):
+        raise com.UnsupportedOperationError(
+            "Only literal interval values are supported"
+        )
+    if unit == "us" or unit not in _interval_units:
+        raise com.UnsupportedOperationError(
+            f"Unsupported bucket interval {op.interval!r}"
+        )
+    if op.offset is not None:
+        raise com.UnsupportedOperationError(
+            "Timestamp bucket with offset is not supported"
+        )
+
+    part = sa.literal_column(_interval_units[unit])
+    value = sa.literal_column(str(op.interval.value))
+    arg = t.translate(op.arg)
+    origin = sa.literal_column("CAST('1970-01-01' AS DATETIME2)")
+
+    return sa.func.DATE_BUCKET(part, value, arg, origin)
 
 
 def _temporal_delta(t, op):
@@ -210,6 +233,7 @@ def _not(t, op):
         ),
         ops.TimestampTruncate: _timestamp_truncate,
         ops.DateTruncate: _timestamp_truncate,
+        ops.TimestampBucket: _timestamp_bucket,
         ops.Hash: unary(sa.func.checksum),
         ops.ExtractMicrosecond: fixed_arity(
             lambda arg: sa.func.datepart(sa.literal_column("microsecond"), arg), 1
diff --git a/ibis/backends/tests/test_temporal.py b/ibis/backends/tests/test_temporal.py
index ace89a9c8f8ba..61310e0f1b4a1 100644
--- a/ibis/backends/tests/test_temporal.py
+++ b/ibis/backends/tests/test_temporal.py
@@ -2507,7 +2507,6 @@ def test_delta(con, start, end, unit, expected):
         "druid",
         "flink",
         "impala",
-        "mssql",
         "mysql",
         "oracle",
         "pandas",
@@ -2552,7 +2551,6 @@ def test_timestamp_bucket(backend, kws, pd_freq):
         "druid",
         "flink",
         "impala",
-        "mssql",
         "mysql",
         "oracle",
         "pandas",
@@ -2565,7 +2563,7 @@ def test_timestamp_bucket(backend, kws, pd_freq):
     raises=com.OperationNotDefinedError,
 )
 @pytest.mark.notimpl(
-    ["clickhouse"],
+    ["clickhouse", "mssql"],
     reason="offset arg not supported",
     raises=com.UnsupportedOperationError,
 )