refactor(exasol): add custom TimestampTruncate (#8590)

ibis-project · Mar 8, 2024 · 66f12c3 · 66f12c3
1 parent de4e988
commit 66f12c3
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 6 deletions.
diff --git a/ibis/backends/exasol/compiler.py b/ibis/backends/exasol/compiler.py
@@ -85,7 +85,6 @@ class ExasolCompiler(SQLGlotCompiler):
             ops.TimestampDelta,
             ops.TimestampDiff,
             ops.TimestampSub,
-            ops.TimestampTruncate,
             ops.TypeOf,
             ops.Unnest,
             ops.Variance,
@@ -165,8 +164,23 @@ def visit_CountDistinctStar(self, op, *, arg, where):
             "COUNT(DISTINCT *) is not supported in Exasol"
         )
 
+    def visit_TimestampTruncate(self, op, *, arg, unit):
+        short_name = unit.short
+        unit_mapping = {"W": "IW"}
+        unsupported = {"ms", "us"}
+
+        if short_name in unsupported:
+            raise com.UnsupportedOperationError(
+                f"Unsupported truncate unit {short_name}"
+            )
+
+        if short_name not in unit_mapping:
+            return super().visit_TimestampTruncate(op, arg=arg, unit=unit)
+
+        return self.f.date_trunc(unit_mapping[short_name], arg)
+
     def visit_DateTruncate(self, op, *, arg, unit):
-        return super().visit_TimestampTruncate(op, arg=arg, unit=unit)
+        return self.visit_TimestampTruncate(op, arg=arg, unit=unit)
 
     def visit_DateDelta(self, op, *, part, left, right):
         # Note: general delta handling could be done based on part (unit),

diff --git a/ibis/backends/tests/test_temporal.py b/ibis/backends/tests/test_temporal.py
@@ -342,7 +342,7 @@ def test_timestamp_extract_week_of_year(backend, alltypes, df):
             "ms",
             marks=[
                 pytest.mark.notimpl(
-                    ["clickhouse", "mysql", "sqlite", "datafusion"],
+                    ["clickhouse", "mysql", "sqlite", "datafusion", "exasol"],
                     raises=com.UnsupportedOperationError,
                 ),
                 pytest.mark.broken(
@@ -356,7 +356,7 @@ def test_timestamp_extract_week_of_year(backend, alltypes, df):
             "us",
             marks=[
                 pytest.mark.notimpl(
-                    ["clickhouse", "mysql", "sqlite", "trino", "datafusion"],
+                    ["clickhouse", "mysql", "sqlite", "trino", "datafusion", "exasol"],
                     raises=com.UnsupportedOperationError,
                 ),
                 pytest.mark.broken(
@@ -389,6 +389,7 @@ def test_timestamp_extract_week_of_year(backend, alltypes, df):
                         "trino",
                         "mssql",
                         "datafusion",
+                        "exasol",
                     ],
                     raises=com.UnsupportedOperationError,
                 ),
@@ -411,7 +412,6 @@ def test_timestamp_extract_week_of_year(backend, alltypes, df):
     raises=AttributeError,
     reason="AttributeError: 'StringColumn' object has no attribute 'truncate'",
 )
-@pytest.mark.notimpl(["exasol"], raises=com.OperationNotDefinedError)
 def test_timestamp_truncate(backend, alltypes, df, unit):
     expr = alltypes.timestamp_col.truncate(unit).name("tmp")
 
@@ -439,7 +439,7 @@ def test_timestamp_truncate(backend, alltypes, df, unit):
             marks=[
                 pytest.mark.notyet(["mysql"], raises=com.UnsupportedOperationError),
                 pytest.mark.broken(
-                    ["flink", "exasol"],
+                    ["flink"],
                     raises=AssertionError,
                     reason="Implemented, but behavior doesn't match other backends",
                 ),