wehub-resource-sync
69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
import sys
|
|
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data._internal.execution.interfaces.ref_bundle import (
|
|
_ref_bundles_iterator_to_block_refs_list,
|
|
)
|
|
from ray.data.tests.conftest import * # noqa
|
|
from ray.data.tests.util import extract_values
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_blocks",
|
|
[
|
|
[pd.DataFrame({"column": ["spam"]}), pd.DataFrame({"column": ["ham", "eggs"]})],
|
|
[
|
|
pa.Table.from_pydict({"column": ["spam"]}),
|
|
pa.Table.from_pydict({"column": ["ham", "eggs"]}),
|
|
],
|
|
],
|
|
)
|
|
def test_from_blocks(input_blocks, ray_start_regular_shared):
|
|
ds = ray.data.from_blocks(input_blocks)
|
|
|
|
bundles = ds.iter_internal_ref_bundles()
|
|
output_blocks = ray.get(_ref_bundles_iterator_to_block_refs_list(bundles))
|
|
assert len(input_blocks) == len(output_blocks)
|
|
assert all(
|
|
input_block.equals(output_block)
|
|
for input_block, output_block in zip(input_blocks, output_blocks)
|
|
)
|
|
|
|
|
|
def test_from_items(ray_start_regular_shared):
|
|
ds = ray.data.from_items(["hello", "world"])
|
|
assert extract_values("item", ds.take()) == ["hello", "world"]
|
|
assert isinstance(next(iter(ds.iter_batches(batch_format=None))), pa.Table)
|
|
|
|
|
|
@pytest.mark.parametrize("parallelism", list(range(1, 21)))
|
|
def test_from_items_parallelism(ray_start_regular_shared, parallelism):
|
|
# Test that specifying parallelism yields the expected number of blocks.
|
|
n = 20
|
|
records = [{"a": i} for i in range(n)]
|
|
ds = ray.data.from_items(records, override_num_blocks=parallelism)
|
|
out = ds.take_all()
|
|
assert out == records
|
|
assert ds._logical_plan.initial_num_blocks() == parallelism
|
|
|
|
|
|
def test_from_items_parallelism_truncated(ray_start_regular_shared):
|
|
# Test that specifying parallelism greater than the number of items is truncated to
|
|
# the number of items.
|
|
n = 10
|
|
parallelism = 20
|
|
records = [{"a": i} for i in range(n)]
|
|
ds = ray.data.from_items(records, override_num_blocks=parallelism)
|
|
out = ds.take_all()
|
|
assert out == records
|
|
assert ds._logical_plan.initial_num_blocks() == n
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", __file__]))
|