|
19 | 19 |
|
20 | 20 | from pandas.util.testing import assert_frame_equal |
21 | 21 | import pandas as pd |
| 22 | +import pytest |
22 | 23 |
|
23 | 24 | import pyarrow as pa |
24 | 25 |
|
@@ -50,6 +51,40 @@ def test_recordbatch_from_to_pandas(): |
50 | 51 | assert_frame_equal(data, result) |
51 | 52 |
|
52 | 53 |
|
| 54 | +def test_recordbatchlist_to_pandas(): |
| 55 | + data1 = pd.DataFrame({ |
| 56 | + 'c1': np.array([1, 1, 2], dtype='uint32'), |
| 57 | + 'c2': np.array([1.0, 2.0, 3.0], dtype='float64'), |
| 58 | + 'c3': [True, None, False], |
| 59 | + 'c4': ['foo', 'bar', None] |
| 60 | + }) |
| 61 | + |
| 62 | + data2 = pd.DataFrame({ |
| 63 | + 'c1': np.array([3, 5], dtype='uint32'), |
| 64 | + 'c2': np.array([4.0, 5.0], dtype='float64'), |
| 65 | + 'c3': [True, True], |
| 66 | + 'c4': ['baz', 'qux'] |
| 67 | + }) |
| 68 | + |
| 69 | + batch1 = pa.RecordBatch.from_pandas(data1) |
| 70 | + batch2 = pa.RecordBatch.from_pandas(data2) |
| 71 | + |
| 72 | + result = pa.dataframe_from_batches([batch1, batch2]) |
| 73 | + data = pd.concat([data1, data2], ignore_index=True) |
| 74 | + assert_frame_equal(data, result) |
| 75 | + |
| 76 | + |
| 77 | +def test_recordbatchlist_schema_equals(): |
| 78 | + data1 = pd.DataFrame({'c1': np.array([1], dtype='uint32')}) |
| 79 | + data2 = pd.DataFrame({'c1': np.array([4.0, 5.0], dtype='float64')}) |
| 80 | + |
| 81 | + batch1 = pa.RecordBatch.from_pandas(data1) |
| 82 | + batch2 = pa.RecordBatch.from_pandas(data2) |
| 83 | + |
| 84 | + with pytest.raises(pa.ArrowException): |
| 85 | + pa.dataframe_from_batches([batch1, batch2]) |
| 86 | + |
| 87 | + |
53 | 88 | def test_table_basics(): |
54 | 89 | data = [ |
55 | 90 | pa.from_pylist(range(5)), |
|
0 commit comments