()
| 1377 | |
| 1378 | @pytest.mark.pandas |
| 1379 | def test_table_to_batches(): |
| 1380 | from pandas.testing import assert_frame_equal |
| 1381 | import pandas as pd |
| 1382 | |
| 1383 | df1 = pd.DataFrame({'a': list(range(10))}) |
| 1384 | df2 = pd.DataFrame({'a': list(range(10, 30))}) |
| 1385 | |
| 1386 | batch1 = pa.RecordBatch.from_pandas(df1, preserve_index=False) |
| 1387 | batch2 = pa.RecordBatch.from_pandas(df2, preserve_index=False) |
| 1388 | |
| 1389 | table = pa.Table.from_batches([batch1, batch2, batch1]) |
| 1390 | |
| 1391 | expected_df = pd.concat([df1, df2, df1], ignore_index=True) |
| 1392 | |
| 1393 | batches = table.to_batches() |
| 1394 | assert len(batches) == 3 |
| 1395 | |
| 1396 | assert_frame_equal(pa.Table.from_batches(batches).to_pandas(), |
| 1397 | expected_df) |
| 1398 | |
| 1399 | batches = table.to_batches(max_chunksize=15) |
| 1400 | assert list(map(len, batches)) == [10, 15, 5, 10] |
| 1401 | |
| 1402 | assert_frame_equal(table.to_pandas(), expected_df) |
| 1403 | assert_frame_equal(pa.Table.from_batches(batches).to_pandas(), |
| 1404 | expected_df) |
| 1405 | |
| 1406 | table_from_iter = pa.Table.from_batches(iter([batch1, batch2, batch1])) |
| 1407 | assert table.equals(table_from_iter) |
| 1408 | |
| 1409 | with pytest.raises(ValueError): |
| 1410 | table.to_batches(max_chunksize=0) |
| 1411 | |
| 1412 | |
| 1413 | @pytest.mark.parametrize( |
nothing calls this directly
no test coverage detected