feat(trino): support years and months in datetime arithmetic

ibis-project · Aug 9, 2024 · 1133973 · 1133973
1 parent 75f594d
commit 1133973
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 16 deletions.
diff --git a/ibis/backends/sql/compilers/trino.py b/ibis/backends/sql/compilers/trino.py
@@ -442,17 +442,20 @@ def visit_TemporalDelta(self, op, *, part, left, right):
 
     def _make_interval(self, arg, unit):
         short = unit.short
-        if short in ("Y", "Q", "M", "W"):
+        if short in ("Q", "W"):
             raise com.UnsupportedOperationError(f"Interval unit {unit!r} not supported")
+
+        if isinstance(arg, sge.Literal):
+            # force strings in interval literals because trino requires it
+            arg.args["is_string"] = True
+            return super()._make_interval(arg, unit)
+
+        elif short in ("Y", "M"):
+            return arg * super()._make_interval(sge.convert("1"), unit)
         elif short in ("D", "h", "m", "s", "ms", "us"):
-            if isinstance(arg, sge.Literal):
-                # force strings in interval literals because trino requires it
-                arg.args["is_string"] = True
-                return super()._make_interval(arg, unit)
-            else:
-                return self.f.parse_duration(
-                    self.f.concat(self.cast(arg, dt.string), short.lower())
-                )
+            return self.f.parse_duration(
+                self.f.concat(self.cast(arg, dt.string), short.lower())
+            )
         else:
             raise com.UnsupportedOperationError(
                 f"Interval unit {unit.name!r} not supported"

diff --git a/ibis/backends/tests/test_temporal.py b/ibis/backends/tests/test_temporal.py
@@ -504,11 +504,6 @@ def test_date_truncate(backend, alltypes, df, unit):
                     raises=TypeError,
                     reason="duration() got an unexpected keyword argument 'months'",
                 ),
-                pytest.mark.notyet(
-                    ["trino"],
-                    raises=com.UnsupportedOperationError,
-                    reason="month not implemented",
-                ),
                 pytest.mark.notyet(
                     ["oracle"],
                     raises=OracleInterfaceError,
@@ -624,7 +619,6 @@ def convert_to_offset(offset, displacement_type=displacement_type):
         param(
             "Y",
             marks=[
-                pytest.mark.notyet(["trino"], raises=com.UnsupportedOperationError),
                 pytest.mark.notyet(
                     ["polars"], raises=TypeError, reason="not supported by polars"
                 ),
@@ -635,7 +629,6 @@ def convert_to_offset(offset, displacement_type=displacement_type):
         param(
             "M",
             marks=[
-                pytest.mark.notyet(["trino"], raises=com.UnsupportedOperationError),
                 pytest.mark.notyet(
                     ["polars"], raises=TypeError, reason="not supported by polars"
                 ),