wehub-resource-sync
515 lines
16 KiB
Python
515 lines
16 KiB
Python
import logging
|
|
import os
|
|
import types
|
|
|
|
import numpy as np
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
import pytest
|
|
from packaging.version import parse as parse_version
|
|
from pytest_lazy_fixtures import lf as lazy_fixture
|
|
|
|
import ray
|
|
import ray.cloudpickle as pickle
|
|
import ray.data
|
|
import ray.train
|
|
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
|
|
from ray.data.extensions.object_extension import (
|
|
ArrowPythonObjectArray,
|
|
)
|
|
from ray.data.extensions.tensor_extension import (
|
|
ArrowTensorArray,
|
|
ArrowVariableShapedTensorArray,
|
|
)
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
|
|
@pytest.fixture
|
|
def null_array():
|
|
return pa.array([])
|
|
|
|
|
|
@pytest.fixture
|
|
def int_array():
|
|
return pa.array(list(range(1000)))
|
|
|
|
|
|
@pytest.fixture
|
|
def int_array_with_nulls():
|
|
return pa.array((list(range(9)) + [None]) * 100)
|
|
|
|
|
|
@pytest.fixture
|
|
def float_array():
|
|
return pa.array([float(i) for i in range(1000)])
|
|
|
|
|
|
@pytest.fixture
|
|
def boolean_array():
|
|
return pa.array([True, False] * 500)
|
|
|
|
|
|
@pytest.fixture
|
|
def string_array():
|
|
return pa.array(["foo", "bar", "bz", None, "quux"] * 200)
|
|
|
|
|
|
@pytest.fixture
|
|
def large_string_array():
|
|
return pa.array(["foo", "bar", "bz", None, "quux"] * 200, type=pa.large_string())
|
|
|
|
|
|
@pytest.fixture
|
|
def binary_array():
|
|
return pa.array([b"foo", b"bar", b"bz", None, b"quux"] * 200)
|
|
|
|
|
|
@pytest.fixture
|
|
def fixed_size_binary_array():
|
|
return pa.array([b"foo", b"bar", b"baz", None, b"qux"] * 200, type=pa.binary(3))
|
|
|
|
|
|
@pytest.fixture
|
|
def large_binary_array():
|
|
return pa.array(
|
|
[b"foo", b"bar", b"bz", None, b"quux"] * 200, type=pa.large_binary()
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def list_array():
|
|
return pa.array(([None] + [list(range(9)) + [None]] * 9) * 100)
|
|
|
|
|
|
@pytest.fixture
|
|
def large_list_array():
|
|
# Large list array with nulls
|
|
return pa.array(
|
|
([None] + [list(range(9)) + [None]] * 9) * 100,
|
|
type=pa.large_list(pa.int64()),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def fixed_size_list_array():
|
|
# Fixed size list array
|
|
return pa.FixedSizeListArray.from_arrays(
|
|
pa.array((list(range(9)) + [None]) * 1000), 10
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def map_array():
|
|
return pa.array(
|
|
[list(zip("abcdefghij", range(10))) for _ in range(1000)],
|
|
type=pa.map_(pa.string(), pa.int64()),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def struct_array():
|
|
# Struct array
|
|
return pa.array({"a": i} for i in range(1000))
|
|
|
|
|
|
@pytest.fixture
|
|
def sparse_union_array():
|
|
return pa.UnionArray.from_sparse(
|
|
pa.array([0, 1] * 500, type=pa.int8()),
|
|
[pa.array(list(range(1000))), pa.array([True, False] * 500)],
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def dense_union_array():
|
|
return pa.UnionArray.from_dense(
|
|
pa.array([0, 1] * 500, type=pa.int8()),
|
|
pa.array(
|
|
[i if i % 2 == 0 else (i % 3) % 2 for i in range(1000)], type=pa.int32()
|
|
),
|
|
[pa.array(list(range(1000))), pa.array([True, False])],
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def dictionary_array():
|
|
return pa.DictionaryArray.from_arrays(
|
|
pa.array((list(range(9)) + [None]) * 100),
|
|
pa.array(["a", "b", "c", "d", "e", "f", "g", "h", "i"]),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def tensor_array():
|
|
return ArrowTensorArray.from_numpy(np.arange(1000 * 4 * 4).reshape((1000, 4, 4)))
|
|
|
|
|
|
@pytest.fixture
|
|
def boolean_tensor_array():
|
|
return ArrowTensorArray.from_numpy(
|
|
np.array(
|
|
[True, False, False, True, False, False, True, True] * 2 * 1000
|
|
).reshape((1000, 4, 4))
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def variable_shaped_tensor_array():
|
|
return ArrowVariableShapedTensorArray.from_numpy(
|
|
np.array(
|
|
[
|
|
np.arange(4).reshape((2, 2)),
|
|
np.arange(4, 13).reshape((3, 3)),
|
|
]
|
|
* 500,
|
|
dtype=object,
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def boolean_variable_shaped_tensor_array():
|
|
return ArrowVariableShapedTensorArray.from_numpy(
|
|
np.array(
|
|
[
|
|
np.array([[True, False], [False, True]]),
|
|
np.array(
|
|
[
|
|
[False, True, False],
|
|
[True, True, False],
|
|
[False, False, False],
|
|
],
|
|
),
|
|
]
|
|
* 500,
|
|
dtype=object,
|
|
)
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def list_of_struct_array():
|
|
return pa.array([{"a": i}, {"a": -i}] for i in range(1000))
|
|
|
|
|
|
@pytest.fixture
|
|
def list_of_empty_struct_array():
|
|
return pa.array([{}, {}] for i in range(1000))
|
|
|
|
|
|
@pytest.fixture
|
|
def complex_nested_array():
|
|
return pa.UnionArray.from_sparse(
|
|
pa.array([0, 1] * 500, type=pa.int8()),
|
|
[
|
|
pa.array(
|
|
[
|
|
{
|
|
"a": i % 2 == 0,
|
|
"b": i,
|
|
"c": "bar",
|
|
}
|
|
for i in range(1000)
|
|
]
|
|
),
|
|
pa.array(
|
|
[list(zip("abcdefghij", range(10))) for _ in range(1000)],
|
|
type=pa.map_(pa.string(), pa.int64()),
|
|
),
|
|
],
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def pickled_objects_array():
|
|
elements = ["test", 20, False, {"some": "value"}, None, np.zeros((10, 10))]
|
|
elements *= 1 + 1000 // len(elements)
|
|
elements = elements[:1000]
|
|
|
|
arr = np.array(elements, dtype=object)
|
|
return ArrowPythonObjectArray.from_objects(arr)
|
|
|
|
|
|
pytest_custom_serialization_arrays = [
|
|
# Null array
|
|
(lazy_fixture("null_array"), 1.0),
|
|
# Int array
|
|
(lazy_fixture("int_array"), 0.1),
|
|
# Array with nulls
|
|
(lazy_fixture("int_array_with_nulls"), 0.1),
|
|
# Float array
|
|
(lazy_fixture("float_array"), 0.1),
|
|
# Boolean array
|
|
# Due to bit-packing, most of the pickle bytes are metadata.
|
|
(lazy_fixture("boolean_array"), 0.8),
|
|
# String array
|
|
(lazy_fixture("string_array"), 0.1),
|
|
# Large string array
|
|
(lazy_fixture("large_string_array"), 0.1),
|
|
# Binary array
|
|
(lazy_fixture("binary_array"), 0.1),
|
|
# Fixed size binary array
|
|
(lazy_fixture("fixed_size_binary_array"), 0.1),
|
|
# Large binary array
|
|
(lazy_fixture("large_binary_array"), 0.1),
|
|
# List array with nulls
|
|
(lazy_fixture("list_array"), 0.1),
|
|
# Large list array with nulls
|
|
(lazy_fixture("large_list_array"), 0.1),
|
|
# Fixed size list array
|
|
(lazy_fixture("fixed_size_list_array"), 0.1),
|
|
# Map array
|
|
(lazy_fixture("map_array"), 0.1),
|
|
# Struct array
|
|
(lazy_fixture("struct_array"), 0.1),
|
|
# Union array (sparse)
|
|
(lazy_fixture("sparse_union_array"), 0.1),
|
|
# Union array (dense)
|
|
(lazy_fixture("dense_union_array"), 0.1),
|
|
# Dictionary array
|
|
(lazy_fixture("dictionary_array"), 0.1),
|
|
# Tensor extension array
|
|
(lazy_fixture("tensor_array"), 0.1),
|
|
# Boolean tensor extension array
|
|
(lazy_fixture("boolean_tensor_array"), 0.25),
|
|
# Variable-shaped tensor extension array
|
|
(lazy_fixture("variable_shaped_tensor_array"), 0.1),
|
|
# Boolean variable-shaped tensor extension array
|
|
(lazy_fixture("boolean_variable_shaped_tensor_array"), 0.25),
|
|
# List of struct array
|
|
(lazy_fixture("list_of_struct_array"), 0.1),
|
|
# List of empty struct array
|
|
(lazy_fixture("list_of_empty_struct_array"), 0.1),
|
|
# Complex nested array
|
|
(lazy_fixture("complex_nested_array"), 0.1),
|
|
# Array of pickled objects
|
|
(lazy_fixture("pickled_objects_array"), 0.1),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("data,cap_mult", pytest_custom_serialization_arrays)
|
|
def test_custom_arrow_data_serializer(ray_start_regular_shared, data, cap_mult):
|
|
if len(data) == 0:
|
|
data = pa.table({"a": []})
|
|
else:
|
|
data = pa.Table.from_arrays(
|
|
[data, data, pa.array(range(1000), type=pa.int32())],
|
|
schema=pa.schema(
|
|
[
|
|
pa.field("arr1", data.type),
|
|
pa.field("arr2", data.type),
|
|
pa.field("arr3", pa.int32()),
|
|
],
|
|
metadata={b"foo": b"bar"},
|
|
),
|
|
)
|
|
ray._private.worker.global_worker.get_serialization_context()
|
|
data.validate()
|
|
pyarrow_version = get_pyarrow_version()
|
|
if pyarrow_version >= parse_version("7.0.0"):
|
|
# get_total_buffer_size API was added in Arrow 7.0.0.
|
|
buf_size = data.get_total_buffer_size()
|
|
# Create a zero-copy slice view of data.
|
|
view = data.slice(10, 10)
|
|
s_arr = pickle.dumps(data)
|
|
s_view = pickle.dumps(view)
|
|
post_slice = pickle.loads(s_view)
|
|
post_slice.validate()
|
|
# Check for round-trip equality.
|
|
assert view.equals(post_slice), post_slice
|
|
# Check that the slice view was truncated upon serialization.
|
|
assert len(s_view) <= cap_mult * len(s_arr)
|
|
for column, pre_column in zip(post_slice.columns, view.columns):
|
|
# Check that offset was reset on slice.
|
|
if column.num_chunks > 0:
|
|
assert column.chunk(0).offset == 0
|
|
# Check that null count was either properly cached or recomputed.
|
|
assert column.null_count == pre_column.null_count
|
|
if pyarrow_version >= parse_version("7.0.0"):
|
|
# Check that slice buffer only contains slice data.
|
|
slice_buf_size = post_slice.get_total_buffer_size()
|
|
if buf_size > 0:
|
|
assert buf_size / slice_buf_size - len(data) / len(post_slice) < 100
|
|
|
|
|
|
def test_custom_arrow_data_serializer_fallback(
|
|
ray_start_regular_shared, propagate_logs, caplog
|
|
):
|
|
# Reset serialization fallback set so warning is logged.
|
|
import ray._private.arrow_serialization as arrow_ser_module
|
|
|
|
arrow_ser_module._serialization_fallback_set = set()
|
|
|
|
data = pa.table(
|
|
{
|
|
"a": pa.UnionArray.from_dense(
|
|
pa.array([0, 1] * 500, type=pa.int8()),
|
|
pa.array(
|
|
[i if i % 2 == 0 else (i % 3) % 2 for i in range(1000)],
|
|
type=pa.int32(),
|
|
),
|
|
[pa.array(list(range(1000))), pa.array([True, False])],
|
|
)
|
|
}
|
|
)
|
|
cap_mult = 0.1
|
|
ray._private.worker.global_worker.get_serialization_context()
|
|
data.validate()
|
|
pyarrow_version = get_pyarrow_version()
|
|
if pyarrow_version >= parse_version("7.0.0"):
|
|
# get_total_buffer_size API was added in Arrow 7.0.0.
|
|
buf_size = data.get_total_buffer_size()
|
|
# Create a zero-copy slice view of data.
|
|
view = data.slice(10, 10)
|
|
# Confirm that (1) fallback works, and (2) warning is logged.
|
|
with caplog.at_level(
|
|
logging.WARNING,
|
|
logger="ray.data._internal.arrow_serialization",
|
|
):
|
|
s_arr = pickle.dumps(data)
|
|
assert "Failed to complete optimized serialization" in caplog.text
|
|
caplog.clear()
|
|
# Confirm that we only warn once per process.
|
|
with caplog.at_level(
|
|
logging.WARNING,
|
|
logger="ray.data._internal.arrow_serialization",
|
|
):
|
|
s_view = pickle.dumps(view)
|
|
assert "Failed to complete optimized serialization" not in caplog.text
|
|
post_slice = pickle.loads(s_view)
|
|
post_slice.validate()
|
|
# Check for round-trip equality.
|
|
assert view.equals(post_slice), post_slice
|
|
# Check that the slice view was truncated upon serialization.
|
|
assert len(s_view) <= cap_mult * len(s_arr)
|
|
for column, pre_column in zip(post_slice.columns, view.columns):
|
|
# Check that offset was reset on slice.
|
|
if column.num_chunks > 0:
|
|
assert column.chunk(0).offset == 0
|
|
# Check that null count was either properly cached or recomputed.
|
|
assert column.null_count == pre_column.null_count
|
|
if pyarrow_version >= parse_version("7.0.0"):
|
|
# Check that slice buffer only contains slice data.
|
|
slice_buf_size = post_slice.get_total_buffer_size()
|
|
if buf_size > 0:
|
|
assert buf_size / slice_buf_size - len(data) / len(post_slice) < 100
|
|
|
|
|
|
def test_arrow_scalar_conversion(ray_start_regular_shared):
|
|
ds = ray.data.from_items([1])
|
|
|
|
def fn(batch: list):
|
|
return {"id": np.array([1])}
|
|
|
|
ds = ds.map_batches(fn)
|
|
res = ds.take()
|
|
assert res == [{"id": 1}], res
|
|
|
|
|
|
def test_arrow_object_and_array_support(ray_start_regular_shared):
|
|
obj = types.SimpleNamespace(some_attribute="test")
|
|
|
|
def f(batch):
|
|
batch_size = len(batch["id"])
|
|
return {
|
|
"array": np.zeros((batch_size, 32, 32, 3)),
|
|
"unsupported": [obj] * batch_size,
|
|
}
|
|
|
|
res = ray.data.range(5).map_batches(f, batch_size=None).take(1)
|
|
assert res[0]["array"].shape == (32, 32, 3)
|
|
assert np.all(res[0]["array"] == 0)
|
|
assert res[0]["unsupported"] == obj
|
|
|
|
|
|
def test_custom_arrow_data_serializer_parquet_roundtrip(
|
|
ray_start_regular_shared, tmp_path
|
|
):
|
|
ray._private.worker.global_worker.get_serialization_context()
|
|
t = pa.table({"a": list(range(10000000))})
|
|
pq.write_table(t, f"{tmp_path}/test.parquet")
|
|
t2 = pq.read_table(f"{tmp_path}/test.parquet")
|
|
s_t = pickle.dumps(t)
|
|
s_t2 = pickle.dumps(t2)
|
|
# Check that the post-Parquet slice view chunks don't cause a serialization blow-up.
|
|
assert len(s_t2) < 1.1 * len(s_t)
|
|
# Check for round-trip equality.
|
|
assert t2.equals(pickle.loads(s_t2))
|
|
|
|
|
|
def test_arrow_schema_ipc_serialization(ray_start_regular_shared):
|
|
"""Test that Arrow Schema uses IPC serialization for performance."""
|
|
from ray._private.arrow_serialization import (
|
|
_arrow_schema_reduce,
|
|
_restore_schema_from_ipc,
|
|
)
|
|
|
|
# Verify the reducer is registered
|
|
ray._private.worker.global_worker.get_serialization_context()
|
|
assert pa.Schema in pickle.CloudPickler.dispatch
|
|
assert pickle.CloudPickler.dispatch[pa.Schema] == _arrow_schema_reduce
|
|
|
|
# Create a complex schema with various types
|
|
schema = pa.schema(
|
|
[
|
|
pa.field("id", pa.int64()),
|
|
pa.field("name", pa.string()),
|
|
pa.field("timestamp", pa.timestamp("us", tz="UTC")),
|
|
pa.field("tags", pa.list_(pa.string())),
|
|
pa.field("metadata", pa.map_(pa.string(), pa.string())),
|
|
pa.field(
|
|
"nested",
|
|
pa.struct(
|
|
[
|
|
pa.field("x", pa.float64()),
|
|
pa.field("y", pa.float64()),
|
|
]
|
|
),
|
|
),
|
|
pa.field("category", pa.dictionary(pa.int8(), pa.string())),
|
|
pa.field("decimal_val", pa.decimal128(18, 6)),
|
|
],
|
|
metadata={b"foo": b"bar"},
|
|
)
|
|
|
|
# Test roundtrip serialization
|
|
serialized = pickle.dumps(schema)
|
|
deserialized = pickle.loads(serialized)
|
|
assert schema.equals(deserialized)
|
|
assert schema.metadata == deserialized.metadata
|
|
|
|
# Verify the reducer uses IPC format (check via direct call)
|
|
restore_func, (ipc_bytes,) = _arrow_schema_reduce(schema)
|
|
assert restore_func == _restore_schema_from_ipc
|
|
# IPC bytes should match what schema.serialize() produces
|
|
assert ipc_bytes == schema.serialize().to_pybytes()
|
|
# Verify restore works
|
|
restored = restore_func(ipc_bytes)
|
|
assert schema.equals(restored)
|
|
|
|
|
|
def test_custom_arrow_data_serializer_disable(shutdown_only):
|
|
ray.shutdown()
|
|
ray.worker._post_init_hooks = []
|
|
context = ray.worker.global_worker.get_serialization_context()
|
|
context._unregister_cloudpickle_reducer(pa.Table)
|
|
# Disable custom Arrow array serialization.
|
|
os.environ["RAY_DISABLE_CUSTOM_ARROW_ARRAY_SERIALIZATION"] = "1"
|
|
ray.init()
|
|
# Create a zero-copy slice view of table.
|
|
t = pa.table({"a": list(range(10000000))})
|
|
view = t.slice(10, 10)
|
|
s_t = pickle.dumps(t)
|
|
s_view = pickle.dumps(view)
|
|
# Check that the slice view contains the full buffer of the underlying array.
|
|
d_view = pickle.loads(s_view)
|
|
assert d_view["a"].chunk(0).buffers()[1].size == t["a"].chunk(0).buffers()[1].size
|
|
# Check that the serialized slice view is large
|
|
assert len(s_view) > 0.8 * len(s_t)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|