Skip to content

Commit

Permalink
feat(pyarrow): support Arrow PyCapsule interface on ibis.Table objects
Browse files Browse the repository at this point in the history
  • Loading branch information
jcrist committed May 7, 2024
1 parent 19a973d commit 1a262b9
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 0 deletions.
15 changes: 15 additions & 0 deletions ibis/backends/tests/test_export.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import pyarrow as pa
import pyarrow.csv as pcsv
import pytest
from packaging.version import parse as vparse
from pytest import param

import ibis
Expand Down Expand Up @@ -44,6 +45,20 @@
limit_no_limit = limit + no_limit


@pytest.mark.skipif(
vparse(pa.__version__) < vparse("14"), reason="pyarrow >= 14 required"
)
def test_table___arrow_c_stream__(awards_players):
res = pa.table(awards_players)
sol = awards_players.to_pyarrow()
assert res.equals(sol)

# With explicit schema
schema = awards_players.schema().to_pyarrow()
res = pa.table(awards_players, schema=schema)
assert res.equals(sol)


@pytest.mark.parametrize("limit", limit_no_limit)
def test_table_to_pyarrow_batches(limit, awards_players):
with awards_players.to_pyarrow_batches(limit=limit) as batch_reader:
Expand Down
3 changes: 3 additions & 0 deletions ibis/expr/types/relations.py
Original file line number Diff line number Diff line change
Expand Up @@ -188,6 +188,9 @@ def __dataframe__(self, nan_as_null: bool = False, allow_copy: bool = True):

return IbisDataFrame(self, nan_as_null=nan_as_null, allow_copy=allow_copy)

def __arrow_c_stream__(self, requested_schema: object | None = None) -> object:
return self.to_pyarrow().__arrow_c_stream__(requested_schema)

def __pyarrow_result__(
self, table: pa.Table, data_mapper: type[PyArrowData] | None = None
) -> pa.Table:
Expand Down

0 comments on commit 1a262b9

Please sign in to comment.