()
| 1424 | |
| 1425 | @pytest.mark.pandas |
| 1426 | def test_table_to_batches(): |
| 1427 | from pandas.testing import assert_frame_equal |
| 1428 | import pandas as pd |
| 1429 | |
| 1430 | df1 = pd.DataFrame({'a': list(range(10))}) |
| 1431 | df2 = pd.DataFrame({'a': list(range(10, 30))}) |
| 1432 | |
| 1433 | batch1 = pa.RecordBatch.from_pandas(df1, preserve_index=False) |
| 1434 | batch2 = pa.RecordBatch.from_pandas(df2, preserve_index=False) |
| 1435 | |
| 1436 | table = pa.Table.from_batches([batch1, batch2, batch1]) |
| 1437 | |
| 1438 | expected_df = pd.concat([df1, df2, df1], ignore_index=True) |
| 1439 | |
| 1440 | batches = table.to_batches() |
| 1441 | assert len(batches) == 3 |
| 1442 | |
| 1443 | assert_frame_equal(pa.Table.from_batches(batches).to_pandas(), |
| 1444 | expected_df) |
| 1445 | |
| 1446 | batches = table.to_batches(max_chunksize=15) |
| 1447 | assert list(map(len, batches)) == [10, 15, 5, 10] |
| 1448 | |
| 1449 | assert_frame_equal(table.to_pandas(), expected_df) |
| 1450 | assert_frame_equal(pa.Table.from_batches(batches).to_pandas(), |
| 1451 | expected_df) |
| 1452 | |
| 1453 | table_from_iter = pa.Table.from_batches(iter([batch1, batch2, batch1])) |
| 1454 | assert table.equals(table_from_iter) |
| 1455 | |
| 1456 | with pytest.raises(ValueError): |
| 1457 | table.to_batches(max_chunksize=0) |
| 1458 | |
| 1459 | |
| 1460 | @pytest.mark.parametrize( |
nothing calls this directly
no test coverage detected