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

---
 ibis/backends/polars/compiler.py     | 28 ++++++++++++++++++++++++++++
 ibis/backends/tests/test_temporal.py |  2 --
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/ibis/backends/polars/compiler.py b/ibis/backends/polars/compiler.py
index 727806acb118..f92af900f3e0 100644
--- a/ibis/backends/polars/compiler.py
+++ b/ibis/backends/polars/compiler.py
@@ -751,6 +751,34 @@ def temporal_truncate(op, **kw):
     return arg.dt.truncate(unit, "-1w")
 
 
+def _compile_literal_interval(op):
+    if not isinstance(op, ops.Literal):
+        raise com.UnsupportedOperationError(
+            "Only literal interval values are supported"
+        )
+
+    if op.dtype.unit.short == "M":
+        suffix = "mo"
+    else:
+        suffix = op.dtype.unit.short.lower()
+
+    return f"{op.value}{suffix}"
+
+
+@translate.register(ops.TimestampBucket)
+def timestamp_bucket(op, **kw):
+    arg = translate(op.arg, **kw)
+    interval = _compile_literal_interval(op.interval)
+    if op.offset is not None:
+        offset = _compile_literal_interval(op.offset)
+        neg_offset = offset[1:] if offset.startswith("-") else f"-{offset}"
+        arg = arg.dt.offset_by(neg_offset)
+    else:
+        offset = None
+    res = arg.dt.truncate(interval, offset)
+    return res
+
+
 @translate.register(ops.DateFromYMD)
 def date_from_ymd(op, **kw):
     return pl.date(
diff --git a/ibis/backends/tests/test_temporal.py b/ibis/backends/tests/test_temporal.py
index 3e0903436a69..99ab454529dc 100644
--- a/ibis/backends/tests/test_temporal.py
+++ b/ibis/backends/tests/test_temporal.py
@@ -2525,7 +2525,6 @@ def test_delta(con, start, end, unit, expected):
         "mysql",
         "oracle",
         "pandas",
-        "polars",
         "pyspark",
         "snowflake",
         "sqlite",
@@ -2573,7 +2572,6 @@ def test_timestamp_bucket(backend, kws, pd_freq):
         "mysql",
         "oracle",
         "pandas",
-        "polars",
         "pyspark",
         "snowflake",
         "sqlite",