feat(clickhouse): support ms/us/ns truncate units

ibis-project · Aug 9, 2024 · 9881edb · 9881edb
1 parent 1133973
commit 9881edb
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 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
@@ -1,2 +1,2 @@
 SELECT
-  toDate(parseDateTimeBestEffort('2009-05-17T12:34:56')) AS "TimestampTruncate(datetime.datetime(2009, 5, 17, 12, 34, 56), DAY)"
+  toStartOfDay(parseDateTimeBestEffort('2009-05-17T12:34:56')) AS "TimestampTruncate(datetime.datetime(2009, 5, 17, 12, 34, 56), DAY)"
diff --git 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
@@ -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",