wehub-resource-sync
305 lines
9.7 KiB
Python
305 lines
9.7 KiB
Python
from collections import UserDict
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data._internal.tensor_extensions.arrow import (
|
|
create_arrow_fixed_shape_tensor_type,
|
|
)
|
|
from ray.data.context import DataContext
|
|
from ray.data.dataset import Schema
|
|
from ray.data.tests.conftest import * # noqa
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
|
|
def test_strict_read_schemas(ray_start_regular_shared_2_cpus):
|
|
ds = ray.data.range(1)
|
|
assert ds.take()[0] == {"id": 0}
|
|
|
|
ds = ray.data.range_tensor(1)
|
|
assert ds.take()[0] == {"data": np.array([0])}
|
|
|
|
ds = ray.data.from_items([1])
|
|
assert ds.take()[0] == {"item": 1}
|
|
|
|
ds = ray.data.from_items([object()])
|
|
assert isinstance(ds.take()[0]["item"], object)
|
|
|
|
ds = ray.data.read_numpy("example://mnist_subset.npy")
|
|
assert "data" in ds.take()[0]
|
|
|
|
ds = ray.data.from_numpy(np.ones((100, 10)))
|
|
assert "data" in ds.take()[0]
|
|
|
|
ds = ray.data.from_numpy_refs(ray.put(np.ones((100, 10))))
|
|
assert "data" in ds.take()[0]
|
|
|
|
ds = ray.data.read_binary_files("example://image-datasets/simple")
|
|
assert "bytes" in ds.take()[0]
|
|
|
|
ds = ray.data.read_images("example://image-datasets/simple")
|
|
assert "image" in ds.take()[0]
|
|
|
|
ds = ray.data.read_text("example://sms_spam_collection_subset.txt")
|
|
assert "text" in ds.take()[0]
|
|
|
|
|
|
def test_strict_map_output(ray_start_regular_shared_2_cpus):
|
|
ds = ray.data.range(1)
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.map(lambda x: 0, max_retries=0).materialize()
|
|
ds.map(lambda x: {"id": 0}).materialize()
|
|
ds.map(lambda x: UserDict({"id": 0})).materialize()
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.map_batches(lambda x: np.array([0]), max_retries=0).materialize()
|
|
ds.map_batches(lambda x: {"id": [0]}).materialize()
|
|
ds.map_batches(lambda x: UserDict({"id": [0]})).materialize()
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.map(lambda x: np.ones(10), max_retries=0).materialize()
|
|
ds.map(lambda x: {"x": np.ones(10)}).materialize()
|
|
ds.map(lambda x: UserDict({"x": np.ones(10)})).materialize()
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.map_batches(lambda x: np.ones(10), max_retries=0).materialize()
|
|
ds.map_batches(lambda x: {"x": np.ones(10)}).materialize()
|
|
ds.map_batches(lambda x: UserDict({"x": np.ones(10)})).materialize()
|
|
|
|
# Not allowed in normal mode either.
|
|
with pytest.raises(ValueError):
|
|
ds.map_batches(lambda x: object(), max_retries=0).materialize()
|
|
with pytest.raises(ValueError):
|
|
ds.map_batches(lambda x: {"x": object()}, max_retries=0).materialize()
|
|
ds.map_batches(lambda x: {"x": [object()]}).materialize()
|
|
ds.map_batches(lambda x: UserDict({"x": [object()]})).materialize()
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.map(lambda x: object(), max_retries=0).materialize()
|
|
ds.map(lambda x: {"x": object()}).materialize()
|
|
ds.map(lambda x: UserDict({"x": object()})).materialize()
|
|
|
|
|
|
def test_strict_convert_map_output(ray_start_regular_shared_2_cpus):
|
|
ds = ray.data.range(1).map_batches(lambda x: {"id": [0, 1, 2, 3]}).materialize()
|
|
assert ds.take_batch()["id"].tolist() == [0, 1, 2, 3]
|
|
|
|
with pytest.raises(ValueError):
|
|
# Strings not converted into array.
|
|
ray.data.range(1).map_batches(
|
|
lambda x: {"id": "string"}, max_retries=0
|
|
).materialize()
|
|
|
|
class UserObj:
|
|
def __eq__(self, other):
|
|
return isinstance(other, UserObj)
|
|
|
|
ds = (
|
|
ray.data.range(1)
|
|
.map_batches(lambda x: {"id": [0, 1, 2, UserObj()]})
|
|
.materialize()
|
|
)
|
|
assert ds.take_batch()["id"].tolist() == [0, 1, 2, UserObj()]
|
|
|
|
|
|
def test_strict_convert_map_groups(ray_start_regular_shared_2_cpus):
|
|
ds = ray.data.read_csv("example://iris.csv")
|
|
|
|
def process_group(group):
|
|
variety = group["variety"][0]
|
|
count = len(group["variety"])
|
|
|
|
# Test implicit list->array conversion here.
|
|
return {
|
|
"variety": [variety],
|
|
"count": [count],
|
|
}
|
|
|
|
ds = ds.groupby("variety").map_groups(process_group)
|
|
ds.show()
|
|
|
|
|
|
def test_strict_default_batch_format(ray_start_regular_shared_2_cpus):
|
|
ds = ray.data.range(1)
|
|
|
|
@ray.remote
|
|
class Queue:
|
|
def __init__(self):
|
|
self.item = None
|
|
|
|
def put(self, item):
|
|
old = self.item
|
|
self.item = item
|
|
return old
|
|
|
|
q = Queue.remote()
|
|
|
|
assert isinstance(next(iter(ds.iter_batches()))["id"], np.ndarray)
|
|
assert isinstance(ds.take_batch()["id"], np.ndarray)
|
|
|
|
def f(x):
|
|
ray.get(q.put.remote(x))
|
|
return x
|
|
|
|
ds.map_batches(f).materialize()
|
|
batch = ray.get(q.put.remote(None))
|
|
assert isinstance(batch["id"], np.ndarray), batch
|
|
|
|
|
|
@pytest.mark.parametrize("shape", [(10,), (10, 2)])
|
|
def test_strict_tensor_support(
|
|
ray_start_regular_shared_2_cpus, restore_data_context, shape
|
|
):
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
ds = ray.data.from_items([np.ones(shape), np.ones(shape)])
|
|
assert np.array_equal(ds.take()[0]["item"], np.ones(shape))
|
|
|
|
ds = ds.map(lambda x: {"item": x["item"] * 2})
|
|
assert np.array_equal(ds.take()[0]["item"], 2 * np.ones(shape))
|
|
|
|
ds = ds.map_batches(lambda x: {"item": x["item"] * 2})
|
|
assert np.array_equal(ds.take()[0]["item"], 4 * np.ones(shape))
|
|
|
|
|
|
def test_strict_value_repr(ray_start_regular_shared_2_cpus):
|
|
ds = ray.data.from_items([{"__value__": np.ones(10)}])
|
|
|
|
ds = ds.map_batches(lambda x: {"__value__": x["__value__"] * 2})
|
|
ds = ds.map(lambda x: {"x": x["__value__"] * 2})
|
|
assert np.array_equal(ds.take()[0]["x"], 4 * np.ones(10))
|
|
assert np.array_equal(ds.take_batch()["x"][0], 4 * np.ones(10))
|
|
|
|
|
|
def test_strict_object_support(ray_start_regular_shared_2_cpus):
|
|
ds = ray.data.from_items([{"x": 2}, {"x": object()}])
|
|
ds.map_batches(lambda x: x, batch_format="numpy").materialize()
|
|
|
|
|
|
def test_strict_compute(ray_start_regular_shared_2_cpus):
|
|
with pytest.raises(ValueError):
|
|
ray.data.range(10).map(lambda x: x, compute="actors").show()
|
|
with pytest.raises(ValueError):
|
|
ray.data.range(10).map(lambda x: x, compute="tasks").show()
|
|
|
|
|
|
def test_strict_schema(ray_start_regular_shared_2_cpus, tensor_format_context):
|
|
import pyarrow as pa
|
|
|
|
from ray.data.extensions.object_extension import (
|
|
ArrowPythonObjectType,
|
|
)
|
|
|
|
ds = ray.data.from_items([{"x": 2}])
|
|
schema = ds.schema()
|
|
assert isinstance(schema.base_schema, pa.lib.Schema)
|
|
assert schema.names == ["x"]
|
|
assert schema.types == [pa.int64()]
|
|
|
|
ds = ray.data.from_items([{"x": 2, "y": [1, 2]}])
|
|
schema = ds.schema()
|
|
assert isinstance(schema.base_schema, pa.lib.Schema)
|
|
assert schema.names == ["x", "y"]
|
|
assert schema.types == [pa.int64(), pa.list_(pa.int64())]
|
|
|
|
ds = ray.data.from_items([{"x": 2, "y": object(), "z": [1, 2]}])
|
|
schema = ds.schema()
|
|
assert isinstance(schema.base_schema, pa.lib.Schema)
|
|
assert schema.names == ["x", "y", "z"]
|
|
assert schema.types == [
|
|
pa.int64(),
|
|
ArrowPythonObjectType(),
|
|
pa.list_(pa.int64()),
|
|
]
|
|
|
|
ds = ray.data.from_numpy(np.ones((100, 10)))
|
|
schema = ds.schema()
|
|
assert isinstance(schema.base_schema, pa.lib.Schema)
|
|
assert schema.names == ["data"]
|
|
|
|
expected_arrow_ext_type = create_arrow_fixed_shape_tensor_type(
|
|
shape=(10,), dtype=pa.float64()
|
|
)
|
|
|
|
assert schema.types == [expected_arrow_ext_type]
|
|
|
|
def _id(batch):
|
|
assert isinstance(batch, pd.DataFrame)
|
|
return batch
|
|
|
|
schema = ds.map_batches(_id, batch_format="pandas").schema()
|
|
|
|
assert isinstance(schema.base_schema, pa.lib.Schema)
|
|
assert schema.names == ["data"]
|
|
# NOTE: Schema by default returns Arrow types
|
|
assert schema.types == [expected_arrow_ext_type]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_dtype, expected_arrow_type",
|
|
[
|
|
(pd.ArrowDtype(pa.int32()), pa.int32()),
|
|
(np.dtype("int64"), pa.int64()),
|
|
# Integer nullable types
|
|
(pd.Int8Dtype(), pa.int8()),
|
|
(pd.Int16Dtype(), pa.int16()),
|
|
(pd.Int32Dtype(), pa.int32()),
|
|
(pd.Int64Dtype(), pa.int64()),
|
|
(pd.UInt8Dtype(), pa.uint8()),
|
|
(pd.UInt16Dtype(), pa.uint16()),
|
|
(pd.UInt32Dtype(), pa.uint32()),
|
|
(pd.UInt64Dtype(), pa.uint64()),
|
|
# Float nullable types
|
|
(pd.Float32Dtype(), pa.float32()),
|
|
(pd.Float64Dtype(), pa.float64()),
|
|
# Boolean nullable type
|
|
(pd.BooleanDtype(), pa.bool_()),
|
|
# String type (default storage)
|
|
(pd.StringDtype(), pa.string()),
|
|
# String type with explicit pyarrow storage
|
|
(pd.StringDtype(storage="pyarrow"), pa.string()),
|
|
# String type with python storage
|
|
(pd.StringDtype(storage="python"), pa.string()),
|
|
],
|
|
)
|
|
def test_schema_types_property(input_dtype, expected_arrow_type):
|
|
"""
|
|
Tests that the Schema.types property correctly converts pandas and numpy
|
|
dtypes to pyarrow types, including BaseMaskedDtype subclasses.
|
|
"""
|
|
from ray.data._internal.pandas_block import PandasBlockSchema
|
|
|
|
schema = Schema(PandasBlockSchema(names=["a"], types=[input_dtype]))
|
|
assert schema.types == [expected_arrow_type]
|
|
|
|
|
|
def test_use_raw_dicts(ray_start_regular_shared_2_cpus):
|
|
assert type(ray.data.range(10).take(1)[0]) is dict
|
|
assert type(ray.data.from_items([1]).take(1)[0]) is dict
|
|
|
|
def checker(x):
|
|
assert type(x) is dict
|
|
return x
|
|
|
|
ray.data.range(10).map(checker).show()
|
|
|
|
|
|
def test_strict_require_batch_size_for_gpu():
|
|
ray.shutdown()
|
|
ray.init(num_cpus=4, num_gpus=1)
|
|
ds = ray.data.range(1)
|
|
with pytest.raises(ValueError):
|
|
ds.map_batches(lambda x: x, num_gpus=1)
|
|
|
|
ds.map_batches(lambda x: x, num_gpus=0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|