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( compat_column_data(next_batch.column(0)).to_pylist(), expected_values) np.testing.assert_equal( compat_column_data(next_batch.column(1)).to_pylist(), expected_values)
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()
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( compat_column_data(next_batch.column(0)).to_pylist(), list(range(0, 4))) np.testing.assert_equal( compat_column_data(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( compat_column_data(next_batch.column(0)).to_pylist(), list(range(4, 8))) np.testing.assert_equal( compat_column_data(next_batch.column(1)).to_pylist(), list(range(4, 8))) assert batcher.empty()
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 compat_column_data(read_batch.columns[0]): assert value == read_seq read_seq += 1 assert read_seq > 0
def test_batch_size_of_one(): """Try if BatchingTableQueue can be used to retrieve row-by-row data (batch size of 1)""" batches = [_new_record_batch([i]) for i in range(3)] table = pa.Table.from_batches(batches) batcher = BatchingTableQueue(1) batcher.put(table) for _ in range(3): assert not batcher.empty() next_batch = batcher.get() assert 1 == next_batch.num_rows assert batcher.empty()