Skip to content

Commit

Permalink
fix(core): cast missing types in pyarrow consistenly with ibis mapping.
Browse files Browse the repository at this point in the history
  • Loading branch information
double-thinker committed Sep 2, 2024
1 parent 5e2903e commit 63b6ff5
Showing 1 changed file with 22 additions and 2 deletions.
24 changes: 22 additions & 2 deletions ibis/backends/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import ibis.expr.types as ir
from ibis import util
from ibis.common.caching import RefCountedCache
from ibis.formats.pyarrow import PyArrowSchema

if TYPE_CHECKING:
from collections.abc import Iterable, Iterator, Mapping, MutableMapping
Expand Down Expand Up @@ -212,8 +213,20 @@ def to_pyarrow(
self._run_pre_execute_hooks(expr)

table_expr = expr.as_table()
schema = table_expr.schema()
arrow_schema = schema.to_pyarrow()
original_schema = table_expr.schema()

# Convert the original schema back and forth to check biyectivity
# of types mappings pyarrow <-> ibis.
pa_compatible_schema = PyArrowSchema.to_ibis(
PyArrowSchema.from_ibis(original_schema)
)
if original_schema != pa_compatible_schema:
# If the original schema is not compatible with PyArrow, we cast
# server side to match the types that PyArrow expects in Ibis.
table_expr = table_expr.cast(pa_compatible_schema)

arrow_schema = pa_compatible_schema.to_pyarrow()

with self.to_pyarrow_batches(
table_expr, params=params, limit=limit, **kwargs
) as reader:
Expand Down Expand Up @@ -258,6 +271,13 @@ def to_polars(
import polars as pl

table = self.to_pyarrow(expr.as_table(), params=params, limit=limit, **kwargs)
pa_compatible_schema = PyArrowSchema.to_ibis(
PyArrowSchema.from_ibis(expr.schema())
)

if expr.schema() != pa_compatible_schema:
expr = expr.cast(pa_compatible_schema)

return expr.__polars_result__(pl.from_arrow(table))

@util.experimental
Expand Down

0 comments on commit 63b6ff5

Please sign in to comment.