From fd61f2c34eafcc50d7a8de652222554b45bd2ef6 Mon Sep 17 00:00:00 2001
From: Phillip Cloud <417981+cpcloud@users.noreply.github.com>
Date: Tue, 30 Jul 2024 07:17:07 -0400
Subject: [PATCH] feat(bigquery): support timestamp bucket

---
 ibis/backends/sql/compilers/bigquery.py | 16 +++++++++++++++-
 ibis/backends/tests/test_temporal.py    | 13 +------------
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/ibis/backends/sql/compilers/bigquery.py b/ibis/backends/sql/compilers/bigquery.py
index b37a91d7d72b..6ab1a1a311b0 100644
--- a/ibis/backends/sql/compilers/bigquery.py
+++ b/ibis/backends/sql/compilers/bigquery.py
@@ -45,7 +45,6 @@ class BigQueryCompiler(SQLGlotCompiler):
         ops.MultiQuantile,
         ops.RegexSplit,
         ops.RowID,
-        ops.TimestampBucket,
         ops.TimestampDiff,
     )
 
@@ -760,3 +759,18 @@ def visit_TableUnnest(
             .from_(parent)
             .join(unnest, join_type="CROSS" if not keep_empty else "LEFT")
         )
+
+    def visit_TimestampBucket(self, op, *, arg, interval, offset):
+        arg_dtype = op.arg.dtype
+        if arg_dtype.timezone is not None:
+            funcname = "timestamp"
+        else:
+            funcname = "datetime"
+
+        func = self.f[f"{funcname}_bucket"]
+
+        origin = sge.convert("1970-01-01")
+        if offset is not None:
+            origin = self.f.anon[f"{funcname}_add"](origin, offset)
+
+        return func(arg, interval, origin)
diff --git a/ibis/backends/tests/test_temporal.py b/ibis/backends/tests/test_temporal.py
index a0835ce5ffbd..0450badb91b2 100644
--- a/ibis/backends/tests/test_temporal.py
+++ b/ibis/backends/tests/test_temporal.py
@@ -2041,17 +2041,7 @@ def test_delta(con, start, end, unit, expected):
 
 
 @pytest.mark.notimpl(
-    [
-        "bigquery",
-        "dask",
-        "impala",
-        "mysql",
-        "pandas",
-        "pyspark",
-        "sqlite",
-        "trino",
-        "druid",
-    ],
+    ["dask", "impala", "mysql", "pandas", "pyspark", "sqlite", "trino", "druid"],
     raises=com.OperationNotDefinedError,
 )
 @pytest.mark.parametrize(
@@ -2154,7 +2144,6 @@ def test_timestamp_bucket(backend, kws, pd_freq):
 
 @pytest.mark.notimpl(
     [
-        "bigquery",
         "dask",
         "datafusion",
         "impala",