def test_read_batches_larger_than_a_table_added():
    """Add a single table composed of 10 one row ten_batches_each_with_one_record. Then read-out two batches of 4
    and verify that no more batches can be read"""
    ten_batches_each_with_one_record = [_new_record_batch([i]) for i in range(10)]
    table_0_10 = pa.Table.from_batches(ten_batches_each_with_one_record)

    batcher = BatchingTableQueue(4)
    batcher.put(table_0_10)

    assert not batcher.empty()
    next_batch = batcher.get()

    assert 4 == next_batch.num_rows
    np.testing.assert_equal(next_batch.column(0).to_pylist(), list(range(0, 4)))
    np.testing.assert_equal(next_batch.column(1).to_pylist(), list(range(0, 4)))

    assert not batcher.empty()
    next_batch = batcher.get()

    assert 4 == next_batch.num_rows
    np.testing.assert_equal(next_batch.column(0).to_pylist(), list(range(4, 8)))
    np.testing.assert_equal(next_batch.column(1).to_pylist(), list(range(4, 8)))

    assert batcher.empty()
Add a single table composed of 10 one row ten_batches_each_with_one_record. Then read-out two batches of 4
    and verify that no more batches can be read
petastorm/pyarrow_helpers/tests/test_batch_buffer.py