def test_random_table_size_and_random_batch_sizes():
    """Add a random number of rows, then read a random number of batches. Repeat multiple times."""
    batch_size = 5
    input_table_size = 50
    read_iter_count = 1000

    batcher = BatchingTableQueue(batch_size)
    write_seq = 0
    read_seq = 0

    for _ in range(read_iter_count):
        next_batch_size = np.random.randint(0, input_table_size)
        new_batch = _new_record_batch(list(range(write_seq, write_seq + next_batch_size)))
        write_seq += next_batch_size

        batcher.put(pa.Table.from_batches([new_batch]))

        next_read = np.random.randint(1, input_table_size // batch_size)
        for _ in range(next_read):
            if not batcher.empty():
                read_batch = batcher.get()
                for value in read_batch.columns[0]:
                    assert value.as_py() == read_seq
                    read_seq += 1

    assert read_seq > 0
Add a random number of rows, then read a random number of batches. Repeat multiple times.
petastorm/pyarrow_helpers/tests/test_batch_buffer.py