wehub-resource-sync
594 lines
19 KiB
Python
594 lines
19 KiB
Python
import math
|
|
import sys
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data.block import BlockAccessor
|
|
from ray.data.context import DataContext
|
|
from ray.data.tests.conftest import * # noqa
|
|
from ray.data.tests.util import extract_values
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
|
|
def test_iter_rows(ray_start_regular_shared):
|
|
# Test simple rows.
|
|
n = 10
|
|
ds = ray.data.range(n)
|
|
for row, k in zip(ds.iter_rows(), range(n)):
|
|
assert row == {"id": k}
|
|
|
|
# Test tabular rows.
|
|
t1 = pa.Table.from_pydict({"one": [1, 2, 3], "two": [2, 3, 4]})
|
|
t2 = pa.Table.from_pydict({"one": [4, 5, 6], "two": [5, 6, 7]})
|
|
t3 = pa.Table.from_pydict({"one": [7, 8, 9], "two": [8, 9, 10]})
|
|
t4 = pa.Table.from_pydict({"one": [10, 11, 12], "two": [11, 12, 13]})
|
|
ts = [t1, t2, t3, t4]
|
|
t = pa.concat_tables(ts)
|
|
ds = ray.data.from_arrow(ts)
|
|
|
|
def to_pylist(table):
|
|
pydict = table.to_pydict()
|
|
names = table.schema.names
|
|
pylist = [
|
|
{column: pydict[column][row] for column in names}
|
|
for row in range(table.num_rows)
|
|
]
|
|
return pylist
|
|
|
|
# Default ArrowRows.
|
|
for row, t_row in zip(ds.iter_rows(), to_pylist(t)):
|
|
assert isinstance(row, dict)
|
|
assert row == t_row
|
|
|
|
# PandasRows after conversion.
|
|
pandas_ds = ds.map_batches(lambda x: x, batch_format="pandas")
|
|
df = t.to_pandas()
|
|
for row, (index, df_row) in zip(pandas_ds.iter_rows(), df.iterrows()):
|
|
assert isinstance(row, dict)
|
|
assert row == df_row.to_dict()
|
|
|
|
|
|
def test_iter_batches_basic(ray_start_regular_shared):
|
|
df1 = pd.DataFrame({"one": [1, 2, 3], "two": [2, 3, 4]})
|
|
df2 = pd.DataFrame({"one": [4, 5, 6], "two": [5, 6, 7]})
|
|
df3 = pd.DataFrame({"one": [7, 8, 9], "two": [8, 9, 10]})
|
|
df4 = pd.DataFrame({"one": [10, 11, 12], "two": [11, 12, 13]})
|
|
dfs = [df1, df2, df3, df4]
|
|
ds = ray.data.from_blocks(dfs)
|
|
|
|
# Default.
|
|
for batch, df in zip(ds.iter_batches(batch_size=None, batch_format="pandas"), dfs):
|
|
assert isinstance(batch, pd.DataFrame)
|
|
assert batch.equals(df)
|
|
|
|
# pyarrow.Table format.
|
|
for batch, df in zip(ds.iter_batches(batch_size=None, batch_format="pyarrow"), dfs):
|
|
assert isinstance(batch, pa.Table)
|
|
assert batch.equals(pa.Table.from_pandas(df))
|
|
|
|
# NumPy format.
|
|
for batch, df in zip(ds.iter_batches(batch_size=None, batch_format="numpy"), dfs):
|
|
assert isinstance(batch, dict)
|
|
assert list(batch.keys()) == ["one", "two"]
|
|
assert all(isinstance(col, np.ndarray) for col in batch.values())
|
|
pd.testing.assert_frame_equal(pd.DataFrame(batch), df)
|
|
|
|
# Test NumPy format on Arrow blocks.
|
|
ds2 = ds.map_batches(lambda b: b, batch_size=None, batch_format="pyarrow")
|
|
for batch, df in zip(ds2.iter_batches(batch_size=None, batch_format="numpy"), dfs):
|
|
assert isinstance(batch, dict)
|
|
assert list(batch.keys()) == ["one", "two"]
|
|
assert all(isinstance(col, np.ndarray) for col in batch.values())
|
|
pd.testing.assert_frame_equal(pd.DataFrame(batch), df)
|
|
|
|
# Default format -> numpy.
|
|
for batch, df in zip(ds.iter_batches(batch_size=None, batch_format="default"), dfs):
|
|
assert isinstance(batch, dict)
|
|
assert list(batch.keys()) == ["one", "two"]
|
|
assert all(isinstance(col, np.ndarray) for col in batch.values())
|
|
pd.testing.assert_frame_equal(pd.DataFrame(batch), df)
|
|
|
|
# Batch size.
|
|
batch_size = 2
|
|
batches = list(ds.iter_batches(batch_size=batch_size, batch_format="pandas"))
|
|
assert all(len(batch) == batch_size for batch in batches)
|
|
assert len(batches) == math.ceil(
|
|
(len(df1) + len(df2) + len(df3) + len(df4)) / batch_size
|
|
)
|
|
assert pd.concat(batches, ignore_index=True).equals(
|
|
pd.concat(dfs, ignore_index=True)
|
|
)
|
|
|
|
# Batch size larger than block.
|
|
batch_size = 4
|
|
batches = list(ds.iter_batches(batch_size=batch_size, batch_format="pandas"))
|
|
assert all(len(batch) == batch_size for batch in batches)
|
|
assert len(batches) == math.ceil(
|
|
(len(df1) + len(df2) + len(df3) + len(df4)) / batch_size
|
|
)
|
|
assert pd.concat(batches, ignore_index=True).equals(
|
|
pd.concat(dfs, ignore_index=True)
|
|
)
|
|
|
|
# Batch size larger than dataset.
|
|
batch_size = 15
|
|
batches = list(ds.iter_batches(batch_size=batch_size, batch_format="pandas"))
|
|
assert all(len(batch) == ds.count() for batch in batches)
|
|
assert len(batches) == 1
|
|
assert pd.concat(batches, ignore_index=True).equals(
|
|
pd.concat(dfs, ignore_index=True)
|
|
)
|
|
|
|
# Batch size drop partial.
|
|
batch_size = 5
|
|
batches = list(
|
|
ds.iter_batches(batch_size=batch_size, drop_last=True, batch_format="pandas")
|
|
)
|
|
assert all(len(batch) == batch_size for batch in batches)
|
|
assert len(batches) == (len(df1) + len(df2) + len(df3) + len(df4)) // batch_size
|
|
assert pd.concat(batches, ignore_index=True).equals(
|
|
pd.concat(dfs, ignore_index=True)[:10]
|
|
)
|
|
|
|
# Batch size don't drop partial.
|
|
batch_size = 5
|
|
batches = list(
|
|
ds.iter_batches(batch_size=batch_size, drop_last=False, batch_format="pandas")
|
|
)
|
|
assert all(len(batch) == batch_size for batch in batches[:-1])
|
|
assert len(batches[-1]) == (len(df1) + len(df2) + len(df3) + len(df4)) % batch_size
|
|
assert len(batches) == math.ceil(
|
|
(len(df1) + len(df2) + len(df3) + len(df4)) / batch_size
|
|
)
|
|
assert pd.concat(batches, ignore_index=True).equals(
|
|
pd.concat(dfs, ignore_index=True)
|
|
)
|
|
|
|
# Prefetch.
|
|
batches = list(
|
|
ds.iter_batches(prefetch_batches=1, batch_size=None, batch_format="pandas")
|
|
)
|
|
assert len(batches) == len(dfs)
|
|
for batch, df in zip(batches, dfs):
|
|
assert isinstance(batch, pd.DataFrame)
|
|
assert batch.equals(df)
|
|
|
|
batch_size = 2
|
|
old_preserve_order = ds.context.execution_options.preserve_order
|
|
try:
|
|
ds.context.execution_options.preserve_order = True
|
|
batches = list(
|
|
ds.iter_batches(
|
|
prefetch_batches=2, batch_size=batch_size, batch_format="pandas"
|
|
)
|
|
)
|
|
assert all(len(batch) == batch_size for batch in batches)
|
|
assert len(batches) == math.ceil(
|
|
(len(df1) + len(df2) + len(df3) + len(df4)) / batch_size
|
|
)
|
|
assert pd.concat(batches, ignore_index=True).equals(
|
|
pd.concat(dfs, ignore_index=True)
|
|
)
|
|
|
|
# Prefetch more than number of blocks.
|
|
batches = list(
|
|
ds.iter_batches(
|
|
prefetch_batches=len(dfs), batch_size=None, batch_format="pandas"
|
|
)
|
|
)
|
|
assert len(batches) == len(dfs)
|
|
for batch, df in zip(batches, dfs):
|
|
assert isinstance(batch, pd.DataFrame)
|
|
assert batch.equals(df)
|
|
finally:
|
|
ds.context.execution_options.preserve_order = old_preserve_order
|
|
|
|
# Prefetch with ray.wait.
|
|
context = DataContext.get_current()
|
|
old_config = context.actor_prefetcher_enabled
|
|
try:
|
|
context.actor_prefetcher_enabled = False
|
|
batches = list(
|
|
ds.iter_batches(prefetch_batches=1, batch_size=None, batch_format="pandas")
|
|
)
|
|
assert len(batches) == len(dfs)
|
|
for batch, df in zip(batches, dfs):
|
|
assert isinstance(batch, pd.DataFrame)
|
|
assert batch.equals(df)
|
|
finally:
|
|
context.actor_prefetcher_enabled = old_config
|
|
|
|
|
|
def test_iter_batches_empty_block(ray_start_regular_shared):
|
|
ds = ray.data.range(1).repartition(10)
|
|
assert str(list(ds.iter_batches(batch_size=None))) == "[{'id': array([0])}]"
|
|
assert (
|
|
str(list(ds.iter_batches(batch_size=1, local_shuffle_buffer_size=1)))
|
|
== "[{'id': array([0])}]"
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("ds_format", ["arrow", "pandas"])
|
|
def test_iter_batches_local_shuffle(shutdown_only, ds_format):
|
|
# Input validation.
|
|
# Batch size must be given for local shuffle.
|
|
with pytest.raises(ValueError):
|
|
list(
|
|
ray.data.range(100).iter_batches(
|
|
batch_size=None, local_shuffle_buffer_size=10
|
|
)
|
|
)
|
|
|
|
def range(n, parallelism=200):
|
|
if ds_format == "arrow":
|
|
ds = ray.data.range(n, override_num_blocks=parallelism)
|
|
elif ds_format == "pandas":
|
|
ds = ray.data.range(n, override_num_blocks=parallelism).map_batches(
|
|
lambda df: df, batch_size=None, batch_format="pandas"
|
|
)
|
|
return ds
|
|
|
|
def to_row_dicts(batch):
|
|
if isinstance(batch, pd.DataFrame):
|
|
return batch.to_dict(orient="records")
|
|
return [{"id": v} for v in batch["id"]]
|
|
|
|
def unbatch(batches):
|
|
return [r for batch in batches for r in to_row_dicts(batch)]
|
|
|
|
def sort(r):
|
|
return sorted(r, key=lambda v: v["id"])
|
|
|
|
base = range(100).take_all()
|
|
|
|
# Local shuffle.
|
|
r1 = unbatch(
|
|
range(100, parallelism=10).iter_batches(
|
|
batch_size=3,
|
|
local_shuffle_buffer_size=25,
|
|
)
|
|
)
|
|
r2 = unbatch(
|
|
range(100, parallelism=10).iter_batches(
|
|
batch_size=3,
|
|
local_shuffle_buffer_size=25,
|
|
)
|
|
)
|
|
# Check randomness of shuffle.
|
|
assert r1 != r2, (r1, r2)
|
|
assert r1 != base
|
|
assert r2 != base
|
|
# Check content.
|
|
assert sort(r1) == sort(base)
|
|
assert sort(r2) == sort(base)
|
|
|
|
# Set seed.
|
|
r1 = unbatch(
|
|
range(100, parallelism=10).iter_batches(
|
|
batch_size=3,
|
|
local_shuffle_buffer_size=25,
|
|
local_shuffle_seed=0,
|
|
)
|
|
)
|
|
r2 = unbatch(
|
|
range(100, parallelism=10).iter_batches(
|
|
batch_size=3,
|
|
local_shuffle_buffer_size=25,
|
|
local_shuffle_seed=0,
|
|
)
|
|
)
|
|
# Check randomness of shuffle.
|
|
assert r1 == r2, (r1, r2)
|
|
assert r1 != base
|
|
# Check content.
|
|
assert sort(r1) == sort(base)
|
|
|
|
# Single block.
|
|
r1 = unbatch(
|
|
range(100, parallelism=1).iter_batches(
|
|
batch_size=3,
|
|
local_shuffle_buffer_size=25,
|
|
)
|
|
)
|
|
r2 = unbatch(
|
|
range(100, parallelism=1).iter_batches(
|
|
batch_size=3,
|
|
local_shuffle_buffer_size=25,
|
|
)
|
|
)
|
|
# Check randomness of shuffle.
|
|
assert r1 != r2, (r1, r2)
|
|
assert r1 != base
|
|
assert r2 != base
|
|
# Check content.
|
|
assert sort(r1) == sort(base)
|
|
assert sort(r2) == sort(base)
|
|
|
|
# Single-row blocks.
|
|
r1 = unbatch(
|
|
range(100, parallelism=100).iter_batches(
|
|
batch_size=3,
|
|
local_shuffle_buffer_size=25,
|
|
)
|
|
)
|
|
r2 = unbatch(
|
|
range(100, parallelism=100).iter_batches(
|
|
batch_size=3,
|
|
local_shuffle_buffer_size=25,
|
|
)
|
|
)
|
|
# Check randomness of shuffle.
|
|
assert r1 != r2, (r1, r2)
|
|
assert r1 != base
|
|
assert r2 != base
|
|
# Check content.
|
|
assert sort(r1) == sort(base)
|
|
assert sort(r2) == sort(base)
|
|
|
|
# Buffer larger than dataset.
|
|
r1 = unbatch(
|
|
range(100, parallelism=10).iter_batches(
|
|
batch_size=3,
|
|
local_shuffle_buffer_size=200,
|
|
)
|
|
)
|
|
r2 = unbatch(
|
|
range(100, parallelism=10).iter_batches(
|
|
batch_size=3,
|
|
local_shuffle_buffer_size=200,
|
|
)
|
|
)
|
|
# Check randomness of shuffle.
|
|
assert r1 != r2, (r1, r2)
|
|
assert r1 != base
|
|
assert r2 != base
|
|
# Check content.
|
|
assert sort(r1) == sort(base)
|
|
assert sort(r2) == sort(base)
|
|
|
|
# Batch size larger than block.
|
|
r1 = unbatch(
|
|
range(100, parallelism=20).iter_batches(
|
|
batch_size=12,
|
|
local_shuffle_buffer_size=25,
|
|
)
|
|
)
|
|
r2 = unbatch(
|
|
range(100, parallelism=20).iter_batches(
|
|
batch_size=12,
|
|
local_shuffle_buffer_size=25,
|
|
)
|
|
)
|
|
# Check randomness of shuffle.
|
|
assert r1 != r2, (r1, r2)
|
|
assert r1 != base
|
|
assert r2 != base
|
|
# Check content.
|
|
assert sort(r1) == sort(base)
|
|
assert sort(r2) == sort(base)
|
|
|
|
# Batch size larger than dataset.
|
|
r1 = unbatch(
|
|
range(100, parallelism=10).iter_batches(
|
|
batch_size=200,
|
|
local_shuffle_buffer_size=400,
|
|
)
|
|
)
|
|
r2 = unbatch(
|
|
range(100, parallelism=10).iter_batches(
|
|
batch_size=200,
|
|
local_shuffle_buffer_size=400,
|
|
)
|
|
)
|
|
# Check randomness of shuffle.
|
|
assert r1 != r2, (r1, r2)
|
|
assert r1 != base
|
|
assert r2 != base
|
|
# Check content.
|
|
assert sort(r1) == sort(base)
|
|
assert sort(r2) == sort(base)
|
|
|
|
# Drop partial batches.
|
|
r1 = unbatch(
|
|
range(100, parallelism=10).iter_batches(
|
|
batch_size=7,
|
|
local_shuffle_buffer_size=21,
|
|
drop_last=True,
|
|
)
|
|
)
|
|
r2 = unbatch(
|
|
range(100, parallelism=10).iter_batches(
|
|
batch_size=7,
|
|
local_shuffle_buffer_size=21,
|
|
drop_last=True,
|
|
)
|
|
)
|
|
# Check randomness of shuffle.
|
|
assert r1 != r2, (r1, r2)
|
|
assert r1 != base
|
|
assert r2 != base
|
|
# Check content.
|
|
# Check that partial batches were dropped.
|
|
assert len(r1) % 7 == 0
|
|
assert len(r2) % 7 == 0
|
|
tmp_base = base
|
|
if ds_format in ("arrow", "pandas"):
|
|
r1 = [tuple(r.items()) for r in r1]
|
|
r2 = [tuple(r.items()) for r in r2]
|
|
tmp_base = [tuple(r.items()) for r in base]
|
|
assert set(r1) <= set(tmp_base)
|
|
assert set(r2) <= set(tmp_base)
|
|
|
|
# Test empty dataset.
|
|
ds = ray.data.from_items([])
|
|
r1 = unbatch(ds.iter_batches(batch_size=2, local_shuffle_buffer_size=10))
|
|
assert len(r1) == 0
|
|
assert r1 == ds.take()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"block_sizes,batch_size,drop_last",
|
|
[
|
|
# Single block, batch smaller than block, keep partial
|
|
([10], 3, False),
|
|
# Single block, batch smaller than block, drop partial
|
|
([10], 3, True),
|
|
# Single block, exact division
|
|
([10], 5, False),
|
|
# Multiple equal-sized blocks, batch doesn't divide evenly, keep partial
|
|
([5, 5, 5], 7, False),
|
|
# Multiple equal-sized blocks, batch doesn't divide evenly, drop partial
|
|
([5, 5, 5], 7, True),
|
|
# Multiple unequal-sized blocks, keep partial
|
|
([1, 5, 10], 4, False),
|
|
# Multiple unequal-sized blocks, drop partial
|
|
([1, 5, 10], 4, True),
|
|
# Edge case: batch_size = 1
|
|
([5, 3, 7], 1, False),
|
|
# Edge case: batch larger than total rows
|
|
([2, 3, 4], 100, False),
|
|
# Exact division across multiple blocks
|
|
([6, 12, 18], 6, False),
|
|
],
|
|
)
|
|
def test_iter_batches_grid(
|
|
ray_start_regular_shared,
|
|
block_sizes,
|
|
batch_size,
|
|
drop_last,
|
|
):
|
|
# Tests slicing, batch combining, and partial batch dropping logic over
|
|
# specific dataset, batching, and dropping configurations.
|
|
# Create the dataset with the given block sizes.
|
|
dfs = []
|
|
running_size = 0
|
|
for block_size in block_sizes:
|
|
dfs.append(
|
|
pd.DataFrame(
|
|
{"value": list(range(running_size, running_size + block_size))}
|
|
)
|
|
)
|
|
running_size += block_size
|
|
num_rows = running_size
|
|
ds = ray.data.from_blocks(dfs)
|
|
|
|
batches = list(
|
|
ds.iter_batches(
|
|
batch_size=batch_size,
|
|
drop_last=drop_last,
|
|
batch_format="pandas",
|
|
)
|
|
)
|
|
if num_rows % batch_size == 0 or not drop_last:
|
|
# Number of batches should be equal to
|
|
# num_rows / batch_size, rounded up.
|
|
assert len(batches) == math.ceil(num_rows / batch_size)
|
|
# Concatenated batches should equal the DataFrame
|
|
# representation of the entire dataset.
|
|
assert pd.concat(batches, ignore_index=True).equals(ds.to_pandas())
|
|
else:
|
|
# Number of batches should be equal to
|
|
# num_rows / batch_size, rounded down.
|
|
assert len(batches) == num_rows // batch_size
|
|
# Concatenated batches should equal the DataFrame
|
|
# representation of the dataset with the partial batch
|
|
# remainder sliced off.
|
|
assert pd.concat(batches, ignore_index=True).equals(
|
|
ds.to_pandas()[: batch_size * (num_rows // batch_size)]
|
|
)
|
|
if num_rows % batch_size == 0 or drop_last:
|
|
assert all(len(batch) == batch_size for batch in batches)
|
|
else:
|
|
assert all(len(batch) == batch_size for batch in batches[:-1])
|
|
assert len(batches[-1]) == num_rows % batch_size
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.version_info >= (3, 12), reason="No tensorflow for Python 3.12+"
|
|
)
|
|
def test_iter_tf_batches_emits_deprecation_warning(ray_start_regular_shared):
|
|
with pytest.warns(DeprecationWarning):
|
|
ray.data.range(1).iter_tf_batches()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.version_info >= (3, 12), reason="No tensorflow for Python 3.12+"
|
|
)
|
|
def test_iter_tf_batches(ray_start_regular_shared):
|
|
df1 = pd.DataFrame(
|
|
{"one": [1, 2, 3], "two": [1.0, 2.0, 3.0], "label": [1.0, 2.0, 3.0]}
|
|
)
|
|
df2 = pd.DataFrame(
|
|
{"one": [4, 5, 6], "two": [4.0, 5.0, 6.0], "label": [4.0, 5.0, 6.0]}
|
|
)
|
|
df3 = pd.DataFrame({"one": [7, 8], "two": [7.0, 8.0], "label": [7.0, 8.0]})
|
|
df = pd.concat([df1, df2, df3])
|
|
ds = ray.data.from_pandas([df1, df2, df3])
|
|
|
|
num_epochs = 2
|
|
for _ in range(num_epochs):
|
|
iterations = []
|
|
for batch in ds.iter_tf_batches(batch_size=3):
|
|
iterations.append(
|
|
np.stack((batch["one"], batch["two"], batch["label"]), axis=1)
|
|
)
|
|
combined_iterations = np.concatenate(iterations)
|
|
np.testing.assert_array_equal(
|
|
np.sort(df.values, axis=0), np.sort(combined_iterations, axis=0)
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.version_info >= (3, 12), reason="No tensorflow for Python 3.12+"
|
|
)
|
|
def test_iter_tf_batches_tensor_ds(ray_start_regular_shared):
|
|
arr1 = np.arange(12).reshape((3, 2, 2))
|
|
arr2 = np.arange(12, 24).reshape((3, 2, 2))
|
|
arr = np.concatenate((arr1, arr2))
|
|
ds = ray.data.from_numpy([arr1, arr2])
|
|
|
|
num_epochs = 2
|
|
for _ in range(num_epochs):
|
|
iterations = []
|
|
for batch in ds.iter_tf_batches(batch_size=2):
|
|
iterations.append(batch["data"])
|
|
combined_iterations = np.concatenate(iterations)
|
|
np.testing.assert_array_equal(
|
|
np.sort(arr, axis=0), np.sort(combined_iterations, axis=0)
|
|
)
|
|
|
|
|
|
def test_get_internal_block_refs(ray_start_regular_shared):
|
|
blocks = ray.data.range(10, override_num_blocks=10).get_internal_block_refs()
|
|
assert len(blocks) == 10
|
|
out = []
|
|
for b in ray.get(blocks):
|
|
out.extend(extract_values("id", BlockAccessor.for_block(b).iter_rows(True)))
|
|
out = sorted(out)
|
|
assert out == list(range(10)), out
|
|
|
|
|
|
def test_iter_internal_ref_bundles(ray_start_regular_shared):
|
|
n = 10
|
|
ds = ray.data.range(n, override_num_blocks=n)
|
|
iter_ref_bundles = ds.iter_internal_ref_bundles()
|
|
|
|
out = []
|
|
ref_bundle_count = 0
|
|
for ref_bundle in iter_ref_bundles:
|
|
for entry in ref_bundle.blocks:
|
|
b = ray.get(entry.ref)
|
|
out.extend(extract_values("id", BlockAccessor.for_block(b).iter_rows(True)))
|
|
ref_bundle_count += 1
|
|
out = sorted(out)
|
|
assert ref_bundle_count == n
|
|
assert out == list(range(n)), out
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", __file__]))
|