From 7695ae1a0d568a184fb4faea472d010182e43c4c Mon Sep 17 00:00:00 2001 From: Jim Crist-Harif Date: Wed, 7 Feb 2024 23:09:41 -0600 Subject: [PATCH] fix(pyarrow): support accepting pyarrow dictionary types as inputs --- ibis/formats/pyarrow.py | 2 ++ ibis/formats/tests/test_pyarrow.py | 8 ++++++++ 2 files changed, 10 insertions(+) diff --git a/ibis/formats/pyarrow.py b/ibis/formats/pyarrow.py index 993ce20fe703..9acd3366a477 100644 --- 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 index c94ce18866f2..bf8a191e4c5d 100644 --- 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"), [