Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(pyarrow): support __arrow_c_schema__ on ibis.Schema objects #9665

Merged
merged 1 commit into from
Jul 22, 2024
Merged

feat(pyarrow): support __arrow_c_schema__ on ibis.Schema objects #9665

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions ibis/expr/schema.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -183,6 +183,9 @@ def to_pyarrow(self):

return PyArrowSchema.from_ibis(self)

def __arrow_c_schema__(self):
return self.to_pyarrow().__arrow_c_schema__()

def to_polars(self):
"""Return the equivalent polars schema."""
from ibis.formats.polars import PolarsSchema
Expand Down
7 changes: 7 additions & 0 deletions ibis/expr/tests/test_schema.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -393,6 +393,13 @@ def test_schema_from_to_pyarrow_schema():
assert restored_schema == pyarrow_schema


def test_schema___arrow_c_schema__():
pytest.importorskip("pyarrow")
schema = sch.Schema({"a": dt.int64, "b": dt.string, "c": dt.boolean})
# smoketest, since no way to create schema from capsule in current pyarrow
assert schema.__arrow_c_schema__() is not None
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

pa.schema will use the __arrow_c_schema__ method for recent pyarrow (v15+ I think).

I sometimes use a wrapper class in tests to ensure it sees only the __arrow_c_schema__ method:

class CSchemaWrapper:
    def __init__(obj):
        self.obj = obj

    def __arrow_c_schema__(self):
        return self.obj.__arrow_c_schema__()

then check

pa.schema(CSchemaWrapper(schema))

Copy link

@corleyma corleyma Jul 23, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fwiw, I tried this (pa.schema(ibis_schema)) with the ibis 9.2.0 and pyarrow 17.0 and alas, no dice:

File [~/test-hook/lib/python3.11/site-packages/pyarrow/types.pxi:5360](https://notebooks.freenome.net/user/mcorley/lab/tree/test-hook/lib/python3.11/site-packages/pyarrow/types.pxi#line=5359), in pyarrow.lib.schema()

File [~/test-hook/lib/python3.11/site-packages/pyarrow/types.pxi:3541](https://notebooks.freenome.net/user/mcorley/lab/tree/test-hook/lib/python3.11/site-packages/pyarrow/types.pxi#line=3540), in pyarrow.lib.field()

File [~/test-hook/lib/python3.11/site-packages/pyarrow/types.pxi:5301](https://notebooks.freenome.net/user/mcorley/lab/tree/test-hook/lib/python3.11/site-packages/pyarrow/types.pxi#line=5300), in pyarrow.lib.ensure_type()

TypeError: DataType expected, got <class 'ibis.expr.datatypes.core.String'>

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There's something odd going on with our base classes and whatever mechanism that pyarrow is using to invoke the method.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Arg, pyarrow first checks whether the input is a Mapping and then checks whether __arrow_c_schema__ exists 😞

Copy link
Member

@cpcloud cpcloud Jul 23, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seems like this functionality is broken in pyarrow: an object can't be both a mapping and implement the __arrow_c_schema__ interface.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Opened up apache/arrow#43388 to track this.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the meantime I'll add an xfailing test that will start to pass once this is addressed.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#9671



@pytest.mark.parametrize("lazy", [False, True])
def test_schema_infer_polars_dataframe(lazy):
pl = pytest.importorskip("polars")
Expand Down