wehub-resource-sync
97 lines
3.4 KiB
Python
97 lines
3.4 KiB
Python
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
from ray.data._internal.arrow_block import ArrowBlockAccessor
|
|
from ray.data._internal.block_batching.block_batching import batch_blocks
|
|
|
|
|
|
def block_generator(num_rows: int, num_blocks: int):
|
|
for i in range(num_blocks):
|
|
yield pa.table({"foo": list(range(i * num_rows, (i + 1) * num_rows))})
|
|
|
|
|
|
class TestBatchBlocks:
|
|
"""Tests for batch_blocks function."""
|
|
|
|
@pytest.mark.parametrize("batch_format", ["pandas", "numpy", "pyarrow"])
|
|
def test_basic(self, batch_format):
|
|
"""Test that batch_blocks yields all data in the requested format."""
|
|
blocks = block_generator(num_rows=3, num_blocks=2)
|
|
batches = list(batch_blocks(blocks, batch_format=batch_format))
|
|
|
|
assert len(batches) == 2
|
|
|
|
if batch_format == "pandas":
|
|
assert isinstance(batches[0], pd.DataFrame)
|
|
assert isinstance(batches[1], pd.DataFrame)
|
|
pd.testing.assert_frame_equal(
|
|
batches[0],
|
|
ArrowBlockAccessor(pa.table({"foo": [0, 1, 2]})).to_pandas(),
|
|
)
|
|
pd.testing.assert_frame_equal(
|
|
batches[1],
|
|
ArrowBlockAccessor(pa.table({"foo": [3, 4, 5]})).to_pandas(),
|
|
)
|
|
elif batch_format == "numpy":
|
|
assert isinstance(batches[0], dict)
|
|
assert isinstance(batches[1], dict)
|
|
np.testing.assert_array_equal(batches[0]["foo"], np.array([0, 1, 2]))
|
|
np.testing.assert_array_equal(batches[1]["foo"], np.array([3, 4, 5]))
|
|
elif batch_format == "pyarrow":
|
|
assert batches == [
|
|
pa.table({"foo": [0, 1, 2]}),
|
|
pa.table({"foo": [3, 4, 5]}),
|
|
]
|
|
else:
|
|
pytest.fail(f"Unsupported batch format {batch_format}")
|
|
|
|
@pytest.mark.parametrize(
|
|
"batch_size,drop_last,expected_values",
|
|
[
|
|
# 6 rows, batch_size=2: yields 3 full batches
|
|
(2, False, [[0, 1], [2, 3], [4, 5]]),
|
|
# 6 rows, batch_size=4: yields 1 full + 1 partial batch
|
|
(4, False, [[0, 1, 2, 3], [4, 5]]),
|
|
# 6 rows, batch_size=4, drop_last: drops partial batch
|
|
(4, True, [[0, 1, 2, 3]]),
|
|
# 6 rows, batch_size=10, drop_last: no batches (all dropped)
|
|
(10, True, []),
|
|
],
|
|
)
|
|
def test_batch_size(self, batch_size, drop_last, expected_values):
|
|
"""Test batch_size and drop_last parameters."""
|
|
blocks = block_generator(num_rows=3, num_blocks=2)
|
|
batches = list(
|
|
batch_blocks(
|
|
blocks,
|
|
batch_size=batch_size,
|
|
drop_last=drop_last,
|
|
batch_format="numpy",
|
|
)
|
|
)
|
|
|
|
assert len(batches) == len(expected_values)
|
|
for batch, expected in zip(batches, expected_values):
|
|
np.testing.assert_array_equal(batch["foo"], np.array(expected))
|
|
|
|
def test_collate_fn(self):
|
|
"""Test that collate_fn transforms batches."""
|
|
|
|
def double_values(batch):
|
|
return {"foo": [x * 2 for x in batch["foo"].tolist()]}
|
|
|
|
blocks = block_generator(num_rows=3, num_blocks=2)
|
|
batches = list(batch_blocks(blocks, collate_fn=double_values))
|
|
|
|
assert len(batches) == 2
|
|
assert batches[0]["foo"] == [0, 2, 4]
|
|
assert batches[1]["foo"] == [6, 8, 10]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|