fix(pyarrow): support accepting pyarrow dictionary types as inputs

jcrist · Feb 8, 2024 · 7695ae1 · 7695ae1
1 parent ea40f57
commit 7695ae1
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/ibis/formats/pyarrow.py b/ibis/formats/pyarrow.py
@@ -129,6 +129,8 @@ def to_ibis(cls, typ: pa.DataType, nullable=True) -> dt.DataType:
             return dt.Map(key_dtype, value_dtype, nullable=nullable)
         elif isinstance(typ, JSONType):
             return dt.JSON()
+        elif pa.types.is_dictionary(typ):
+            return cls.to_ibis(typ.value_type)
         else:
             return _from_pyarrow_types[typ](nullable=nullable)
 

diff --git a/ibis/formats/tests/test_pyarrow.py b/ibis/formats/tests/test_pyarrow.py
@@ -122,6 +122,14 @@ def test_dtype_from_nullable_list_type(value_nullable, list_nullable):
     assert restored_type.value_field.nullable is value_nullable
 
 
+@pytest.mark.parametrize("value_type", [pa.string(), pa.date64()])
+def test_dtype_from_dictionary_type(value_type):
+    dict_type = pa.dictionary(pa.int32(), value_type)
+    ibis_type = PyArrowType.to_ibis(dict_type)
+    expected = PyArrowType.to_ibis(value_type)
+    assert ibis_type == expected
+
+
 @pytest.mark.parametrize(
     ("ibis_type", "arrow_type"),
     [