def test_two_tables_of_10_added_reading_5_batches_of_4():
    """Add two tables to batcher and read a batch that covers parts of both tables"""
    table_0_9 = pa.Table.from_batches([_new_record_batch(range(0, 10))])
    table_10_19 = pa.Table.from_batches([_new_record_batch(range(10, 20))])

    batcher = BatchingTableQueue(4)
    assert batcher.empty()

    batcher.put(table_0_9)
    batcher.put(table_10_19)

    for i in range(5):
        assert not batcher.empty()
        next_batch = batcher.get()

        assert (i != 4) == (not batcher.empty())

        assert 4 == next_batch.num_rows
        expected_values = list(range(i * 4, i * 4 + 4))
        np.testing.assert_equal(next_batch.column(0).to_pylist(), expected_values)
        np.testing.assert_equal(next_batch.column(1).to_pylist(), expected_values)
Add two tables to batcher and read a batch that covers parts of both tables
petastorm/pyarrow_helpers/tests/test_batch_buffer.py