From 9881edb0c84a919c1d76ae8af745406f76cc6bea Mon Sep 17 00:00:00 2001
From: Phillip Cloud <417981+cpcloud@users.noreply.github.com>
Date: Fri, 9 Aug 2024 06:05:15 -0400
Subject: [PATCH] feat(clickhouse): support ms/us/ns truncate units

---
 .../test_timestamp_truncate/d/out.sql         |  2 +-
 ibis/backends/sql/compilers/clickhouse.py     | 29 +++++++++----------
 ibis/backends/tests/test_temporal.py          |  5 ++--
 3 files changed, 17 insertions(+), 19 deletions(-)

diff --git a/ibis/backends/clickhouse/tests/snapshots/test_functions/test_timestamp_truncate/d/out.sql b/ibis/backends/clickhouse/tests/snapshots/test_functions/test_timestamp_truncate/d/out.sql
index f7597fcd2c21..570242c29545 100644
--- a/ibis/backends/clickhouse/tests/snapshots/test_functions/test_timestamp_truncate/d/out.sql
+++ b/ibis/backends/clickhouse/tests/snapshots/test_functions/test_timestamp_truncate/d/out.sql
@@ -1,2 +1,2 @@
 SELECT
-  toDate(parseDateTimeBestEffort('2009-05-17T12:34:56')) AS "TimestampTruncate(datetime.datetime(2009, 5, 17, 12, 34, 56), DAY)"
\ No newline at end of file
+  toStartOfDay(parseDateTimeBestEffort('2009-05-17T12:34:56')) AS "TimestampTruncate(datetime.datetime(2009, 5, 17, 12, 34, 56), DAY)"
\ No newline at end of file
diff --git a/ibis/backends/sql/compilers/clickhouse.py b/ibis/backends/sql/compilers/clickhouse.py
index d1f3db494f2e..5ea7d7fab26d 100644
--- a/ibis/backends/sql/compilers/clickhouse.py
+++ b/ibis/backends/sql/compilers/clickhouse.py
@@ -372,24 +372,23 @@ def visit_TimestampFromUNIX(self, op, *, arg, unit):
         return self.f.toDateTime(arg)
 
     def visit_TimestampTruncate(self, op, *, arg, unit):
-        converters = {
-            "Y": "toStartOfYear",
-            "Q": "toStartOfQuarter",
-            "M": "toStartOfMonth",
-            "W": "toMonday",
-            "D": "toDate",
-            "h": "toStartOfHour",
-            "m": "toStartOfMinute",
-            "s": "toDateTime",
-        }
+        if (short := unit.short) == "W":
+            func = "toMonday"
+        else:
+            func = f"toStartOf{unit.singular.capitalize()}"
 
-        unit = unit.short
-        if (converter := converters.get(unit)) is None:
-            raise com.UnsupportedOperationError(f"Unsupported truncate unit {unit}")
+        if short in ("s", "ms", "us", "ns"):
+            arg = self.f.toDateTime64(arg, op.arg.dtype.scale or 0)
+        return self.f[func](arg)
 
-        return self.f[converter](arg)
+    visit_TimeTruncate = visit_TimestampTruncate
 
-    visit_TimeTruncate = visit_DateTruncate = visit_TimestampTruncate
+    def visit_DateTruncate(self, op, *, arg, unit):
+        if unit.short == "W":
+            func = "toMonday"
+        else:
+            func = f"toStartOf{unit.singular.capitalize()}"
+        return self.f[func](arg)
 
     def visit_TimestampBucket(self, op, *, arg, interval, offset):
         if offset is not None:
diff --git a/ibis/backends/tests/test_temporal.py b/ibis/backends/tests/test_temporal.py
index 4c8cec0b903c..0509d4816adc 100644
--- a/ibis/backends/tests/test_temporal.py
+++ b/ibis/backends/tests/test_temporal.py
@@ -359,7 +359,7 @@ def test_timestamp_extract_week_of_year(backend, alltypes, df):
             "ms",
             marks=[
                 pytest.mark.notimpl(
-                    ["clickhouse", "mysql", "sqlite", "datafusion", "exasol"],
+                    ["mysql", "sqlite", "datafusion", "exasol"],
                     raises=com.UnsupportedOperationError,
                 ),
                 pytest.mark.notimpl(["druid"], raises=PyDruidProgrammingError),
@@ -370,7 +370,7 @@ def test_timestamp_extract_week_of_year(backend, alltypes, df):
             "us",
             marks=[
                 pytest.mark.notimpl(
-                    ["clickhouse", "mysql", "sqlite", "trino", "datafusion", "exasol"],
+                    ["mysql", "sqlite", "trino", "datafusion", "exasol"],
                     raises=com.UnsupportedOperationError,
                 ),
                 pytest.mark.notyet(
@@ -388,7 +388,6 @@ def test_timestamp_extract_week_of_year(backend, alltypes, df):
                 pytest.mark.notimpl(
                     [
                         "bigquery",
-                        "clickhouse",
                         "duckdb",
                         "impala",
                         "mysql",