Skip to content

Commit

Permalink
fix(pyarrow): support accepting pyarrow dictionary types as inputs (i…
Browse files Browse the repository at this point in the history
  • Loading branch information
jcrist authored and cpcloud committed Feb 12, 2024
1 parent bd874d8 commit cc54f13
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 0 deletions.
2 changes: 2 additions & 0 deletions ibis/formats/pyarrow.py
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,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)

Expand Down
8 changes: 8 additions & 0 deletions ibis/formats/tests/test_pyarrow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"),
[
Expand Down

0 comments on commit cc54f13

Please sign in to comment.