|
7 | 7 | import pandas.util._test_decorators as td |
8 | 8 |
|
9 | 9 | import pandas as pd |
| 10 | +import pandas._testing as tm |
10 | 11 |
|
11 | 12 | pa = pytest.importorskip("pyarrow") |
12 | 13 |
|
@@ -47,3 +48,42 @@ def test_dataframe_to_arrow(): |
47 | 48 | table = pa.RecordBatchReader.from_stream(df, schema=schema).read_all() |
48 | 49 | expected = expected.cast(schema) |
49 | 50 | assert table.equals(expected) |
| 51 | + |
| 52 | + |
| 53 | +class ArrowArrayWrapper: |
| 54 | + def __init__(self, batch): |
| 55 | + self.array = batch |
| 56 | + |
| 57 | + def __arrow_c_array__(self, requested_schema=None): |
| 58 | + return self.array.__arrow_c_array__(requested_schema) |
| 59 | + |
| 60 | + |
| 61 | +class ArrowStreamWrapper: |
| 62 | + def __init__(self, table): |
| 63 | + self.stream = table |
| 64 | + |
| 65 | + def __arrow_c_stream__(self, requested_schema=None): |
| 66 | + return self.stream.__arrow_c_stream__(requested_schema) |
| 67 | + |
| 68 | + |
| 69 | +@td.skip_if_no("pyarrow", min_version="14.0") |
| 70 | +def test_dataframe_from_arrow(): |
| 71 | + # objects with __arrow_c_stream__ |
| 72 | + table = pa.table({"a": [1, 2, 3], "b": ["a", "b", "c"]}) |
| 73 | + |
| 74 | + result = pd.DataFrame.from_arrow(table) |
| 75 | + expected = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}) |
| 76 | + tm.assert_frame_equal(result, expected) |
| 77 | + |
| 78 | + # not only pyarrow object are supported |
| 79 | + result = pd.DataFrame.from_arrow(ArrowStreamWrapper(table)) |
| 80 | + tm.assert_frame_equal(result, expected) |
| 81 | + |
| 82 | + # objects with __arrow_c_array__ |
| 83 | + batch = pa.record_batch([[1, 2, 3], ["a", "b", "c"]], names=["a", "b"]) |
| 84 | + |
| 85 | + result = pd.DataFrame.from_arrow(table) |
| 86 | + tm.assert_frame_equal(result, expected) |
| 87 | + |
| 88 | + result = pd.DataFrame.from_arrow(ArrowArrayWrapper(batch)) |
| 89 | + tm.assert_frame_equal(result, expected) |
0 commit comments