Skip to content

Commit 9881edb

Browse files
cpcloudjcrist
authored andcommitted
feat(clickhouse): support ms/us/ns truncate units
1 parent 1133973 commit 9881edb

File tree

3 files changed

+17
-19
lines changed

3 files changed

+17
-19
lines changed
Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,2 @@
11
SELECT
2-
toDate(parseDateTimeBestEffort('2009-05-17T12:34:56')) AS "TimestampTruncate(datetime.datetime(2009, 5, 17, 12, 34, 56), DAY)"
2+
toStartOfDay(parseDateTimeBestEffort('2009-05-17T12:34:56')) AS "TimestampTruncate(datetime.datetime(2009, 5, 17, 12, 34, 56), DAY)"

ibis/backends/sql/compilers/clickhouse.py

Lines changed: 14 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -372,24 +372,23 @@ def visit_TimestampFromUNIX(self, op, *, arg, unit):
372372
return self.f.toDateTime(arg)
373373

374374
def visit_TimestampTruncate(self, op, *, arg, unit):
375-
converters = {
376-
"Y": "toStartOfYear",
377-
"Q": "toStartOfQuarter",
378-
"M": "toStartOfMonth",
379-
"W": "toMonday",
380-
"D": "toDate",
381-
"h": "toStartOfHour",
382-
"m": "toStartOfMinute",
383-
"s": "toDateTime",
384-
}
375+
if (short := unit.short) == "W":
376+
func = "toMonday"
377+
else:
378+
func = f"toStartOf{unit.singular.capitalize()}"
385379

386-
unit = unit.short
387-
if (converter := converters.get(unit)) is None:
388-
raise com.UnsupportedOperationError(f"Unsupported truncate unit {unit}")
380+
if short in ("s", "ms", "us", "ns"):
381+
arg = self.f.toDateTime64(arg, op.arg.dtype.scale or 0)
382+
return self.f[func](arg)
389383

390-
return self.f[converter](arg)
384+
visit_TimeTruncate = visit_TimestampTruncate
391385

392-
visit_TimeTruncate = visit_DateTruncate = visit_TimestampTruncate
386+
def visit_DateTruncate(self, op, *, arg, unit):
387+
if unit.short == "W":
388+
func = "toMonday"
389+
else:
390+
func = f"toStartOf{unit.singular.capitalize()}"
391+
return self.f[func](arg)
393392

394393
def visit_TimestampBucket(self, op, *, arg, interval, offset):
395394
if offset is not None:

ibis/backends/tests/test_temporal.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -359,7 +359,7 @@ def test_timestamp_extract_week_of_year(backend, alltypes, df):
359359
"ms",
360360
marks=[
361361
pytest.mark.notimpl(
362-
["clickhouse", "mysql", "sqlite", "datafusion", "exasol"],
362+
["mysql", "sqlite", "datafusion", "exasol"],
363363
raises=com.UnsupportedOperationError,
364364
),
365365
pytest.mark.notimpl(["druid"], raises=PyDruidProgrammingError),
@@ -370,7 +370,7 @@ def test_timestamp_extract_week_of_year(backend, alltypes, df):
370370
"us",
371371
marks=[
372372
pytest.mark.notimpl(
373-
["clickhouse", "mysql", "sqlite", "trino", "datafusion", "exasol"],
373+
["mysql", "sqlite", "trino", "datafusion", "exasol"],
374374
raises=com.UnsupportedOperationError,
375375
),
376376
pytest.mark.notyet(
@@ -388,7 +388,6 @@ def test_timestamp_extract_week_of_year(backend, alltypes, df):
388388
pytest.mark.notimpl(
389389
[
390390
"bigquery",
391-
"clickhouse",
392391
"duckdb",
393392
"impala",
394393
"mysql",

0 commit comments

Comments
 (0)