feat: add support for Ibis 6.x

ibis-project · Aug 2, 2023 · fcfc595 · fcfc595
1 parent 5c6f4fa
commit fcfc595
Show file tree

Hide file tree

Showing 5 changed files with 432 additions and 43 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -94,11 +94,11 @@ jobs:
         ibis-version:
           - "4.1"
           - "5.1"
+          - "6.0"
         include:
          - os: windows-latest
            python-version: "3.11"
-           ibis-version: "5.1"
-
+           ibis-version: "6.0"
 
     steps:
       - uses: actions/checkout@v3
@@ -133,6 +133,7 @@ jobs:
         ibis-version:
           - "4.1"
           - "5.1"
+          - "6.0"
 
     steps:
       - uses: actions/checkout@v3

diff --git a/ibis_substrait/compiler/translate.py b/ibis_substrait/compiler/translate.py
@@ -134,21 +134,49 @@ def _timestamp(dtype: dt.Timestamp) -> stt.Type:
     return stt.Type(timestamp=stt.Type.Timestamp(nullability=nullability))
 
 
-@translate.register
-def _interval(dtype: dt.Interval) -> stt.Type:
-    unit = dtype.unit
-    nullability = _nullability(dtype)
+if IBIS_GTE_6:
+
+    @translate.register
+    def _interval(dtype: dt.Interval) -> stt.Type:
+        unit = dtype.unit.name
+        nullability = _nullability(dtype)
+
+        if unit == "YEAR":
+            return stt.Type(
+                interval_year=stt.Type.IntervalYear(nullability=nullability)
+            )
+        elif unit == "MONTH":
+            return stt.Type(
+                interval_year=stt.Type.IntervalYear(nullability=nullability)
+            )
+        elif unit == "DAY":
+            return stt.Type(interval_day=stt.Type.IntervalDay(nullability=nullability))
+        elif unit == "SECOND":
+            return stt.Type(interval_day=stt.Type.IntervalDay(nullability=nullability))
 
-    if unit == "Y":
-        return stt.Type(interval_year=stt.Type.IntervalYear(nullability=nullability))
-    elif unit == "M":
-        return stt.Type(interval_year=stt.Type.IntervalYear(nullability=nullability))
-    elif unit == "D":
-        return stt.Type(interval_day=stt.Type.IntervalDay(nullability=nullability))
-    elif unit == "s":
-        return stt.Type(interval_day=stt.Type.IntervalDay(nullability=nullability))
+        raise ValueError(f"unsupported substrait unit: {unit!r}")
+
+else:
+
+    @translate.register
+    def _interval(dtype: dt.Interval) -> stt.Type:
+        unit = dtype.unit
+        nullability = _nullability(dtype)
+
+        if unit == "Y":
+            return stt.Type(
+                interval_year=stt.Type.IntervalYear(nullability=nullability)
+            )
+        elif unit == "M":
+            return stt.Type(
+                interval_year=stt.Type.IntervalYear(nullability=nullability)
+            )
+        elif unit == "D":
+            return stt.Type(interval_day=stt.Type.IntervalDay(nullability=nullability))
+        elif unit == "s":
+            return stt.Type(interval_day=stt.Type.IntervalDay(nullability=nullability))
 
-    raise ValueError(f"unsupported substrait unit: {unit!r}")
+        raise ValueError(f"unsupported substrait unit: {unit!r}")
 
 
 @translate.register

diff --git a/ibis_substrait/tests/compiler/test_extensions.py b/ibis_substrait/tests/compiler/test_extensions.py
@@ -260,7 +260,10 @@ def test_extension_boolean(compiler, left, right, bin_op, exp_func, exp_uri):
 
 
 def test_extension_udf_compile(compiler):
-    from ibis.udf.vectorized import elementwise
+    try:
+        from ibis.udf.vectorized import elementwise
+    except ImportError:
+        from ibis.legacy.udf.vectorized import elementwise
 
     pc = None
 

diff --git a/ibis_substrait/tests/integration/test_pyarrow.py b/ibis_substrait/tests/integration/test_pyarrow.py
@@ -1,6 +1,5 @@
 import ibis
 import pytest
-from ibis.udf.vectorized import elementwise
 from packaging.version import parse as vparse
 
 from ibis_substrait.compiler.core import SubstraitCompiler
@@ -9,6 +8,11 @@
 pc = pytest.importorskip("pyarrow.compute")
 pa_substrait = pytest.importorskip("pyarrow.substrait")
 
+try:
+    from ibis.udf.vectorized import elementwise
+except ImportError:
+    from ibis.legacy.udf.vectorized import elementwise
+
 
 arrow12 = pytest.mark.skipif(
     vparse(pa.__version__) <= vparse("11.0.0"),
@@ -103,7 +107,14 @@ def register_pyarrow_udf(udf, registry=None):
         """
         import inspect
 
-        from ibis.backends.pyarrow.datatypes import to_pyarrow_type
+        try:
+            # Ibis 6.x
+            from ibis.formats.pyarrow import PyArrowType
+
+            to_pyarrow_type = PyArrowType.from_ibis
+        except ImportError:
+            # Ibis 4.x, 5.x
+            from ibis.backends.pyarrow.datatypes import to_pyarrow_type
 
         if registry is None:
             registry = pc.function_registry()