Skip to content

Commit

Permalink
feat(pyarrow): support conversion to pyarrow map and struct types
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud committed Jan 20, 2023
1 parent 64be7b1 commit 54a4557
Showing 1 changed file with 16 additions and 8 deletions.
24 changes: 16 additions & 8 deletions ibis/backends/pyarrow/datatypes.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,28 +30,36 @@


@functools.singledispatch
def to_pyarrow_type(dtype):
def to_pyarrow_type(dtype: dt.DataType):
return _to_pyarrow_types[dtype.__class__]


@to_pyarrow_type.register(dt.Array)
def from_ibis_array(dtype):
return pa.list_(to_pyarrow_type(dtype.value_type))


@to_pyarrow_type.register(dt.Set)
def from_ibis_set(dtype):
def from_ibis_collection(dtype: dt.Array | dt.Set):
return pa.list_(to_pyarrow_type(dtype.value_type))


@to_pyarrow_type.register(dt.Interval)
def from_ibis_interval(dtype):
@to_pyarrow_type.register
def from_ibis_interval(dtype: dt.Interval):
try:
return pa.duration(dtype.unit)
except ValueError:
raise com.IbisTypeError(f"Unsupported interval unit: {dtype.unit}")


@to_pyarrow_type.register
def from_ibis_struct(dtype: dt.Struct):
return pa.struct(
pa.field(name, to_pyarrow_type(typ)) for name, typ in dtype.pairs.items()
)


@to_pyarrow_type.register
def from_ibis_map(dtype: dt.Map):
return pa.map_(to_pyarrow_type(dtype.key_type), to_pyarrow_type(dtype.value_type))


_to_ibis_dtypes = {
pa.int8(): dt.Int8,
pa.int16(): dt.Int16,
Expand Down

0 comments on commit 54a4557

Please sign in to comment.