def test_single_table_of_10_rows_added_and_2_batches_of_4_read():
    """Add a single table composed of a single batch with 0..9 rows into batcher. Then read two batches of 4
    and verify that no more batches can be read"""

    # Table with two columns. Each column with 0..9 sequence
    one_batch_of_10_records = [_new_record_batch(range(0, 10))]
    table_0_10 = pa.Table.from_batches(one_batch_of_10_records)

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

    # Load 10 rows into batcher
    batcher.put(table_0_10)

    # Get first batch of 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(0, 4)))
    np.testing.assert_equal(next_batch.column(1).to_pylist(), list(range(0, 4)))

    # Get second batch of 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)))

    # No more batches available
    assert batcher.empty()
Add a single table composed of a single batch with 0..9 rows into batcher. Then read two batches of 4
    and verify that no more batches can be read
petastorm/pyarrow_helpers/tests/test_batch_buffer.py