wehub-resource-sync
378 lines
14 KiB
Python
378 lines
14 KiB
Python
from datetime import datetime
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
import torch
|
|
|
|
import ray
|
|
from ray.data._internal.tensor_extensions.utils import create_ragged_ndarray
|
|
from ray.data.context import DataContext
|
|
from ray.data.tests.conftest import * # noqa
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
|
|
class UserObj:
|
|
def __eq__(self, other):
|
|
return isinstance(other, UserObj)
|
|
|
|
|
|
def do_map_batches(data):
|
|
ds = ray.data.range(1)
|
|
ds = ds.map_batches(lambda x: {"output": data})
|
|
return ds.take_batch()["output"]
|
|
|
|
|
|
def assert_structure_equals(a, b):
|
|
assert type(a) is type(b), (type(a), type(b))
|
|
assert type(a[0]) == type(b[0]), (type(a[0]), type(b[0])) # noqa: E721
|
|
assert a.dtype == b.dtype
|
|
assert a.shape == b.shape
|
|
for i in range(len(a)):
|
|
assert np.array_equal(a[i], b[i]), (i, a[i], b[i])
|
|
|
|
|
|
def test_list_of_scalars(ray_start_regular_shared, restore_data_context):
|
|
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
data = [1, 2, 3]
|
|
output = do_map_batches(data)
|
|
assert_structure_equals(output, np.array([1, 2, 3], dtype=np.int64))
|
|
|
|
|
|
def test_list_of_numpy_scalars(ray_start_regular_shared, restore_data_context):
|
|
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
data = [np.int64(1), np.int64(2), np.int64(3)]
|
|
output = do_map_batches(data)
|
|
assert_structure_equals(output, np.array([1, 2, 3], dtype=np.int64))
|
|
|
|
|
|
def test_list_of_objects(ray_start_regular_shared, restore_data_context):
|
|
# NOTE: Fallback is enabled by default, this is purely for notational purposes
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = True
|
|
|
|
data = [1, 2, 3, UserObj()]
|
|
output = do_map_batches(data)
|
|
assert_structure_equals(output, np.array([1, 2, 3, UserObj()]))
|
|
|
|
|
|
# Define datetime values with different precisions
|
|
DATETIME_DAY_PRECISION = datetime(year=2024, month=1, day=1)
|
|
DATETIME_HOUR_PRECISION = datetime(year=2024, month=1, day=1, hour=1)
|
|
DATETIME_MIN_PRECISION = datetime(year=2024, month=1, day=1, minute=1)
|
|
DATETIME_SEC_PRECISION = datetime(year=2024, month=1, day=1, second=1)
|
|
DATETIME_MILLISEC_PRECISION = datetime(year=2024, month=1, day=1, microsecond=1000)
|
|
DATETIME_MICROSEC_PRECISION = datetime(year=2024, month=1, day=1, microsecond=1)
|
|
|
|
# Define pandas values for different precisions
|
|
PANDAS_DAY_PRECISION = pd.Timestamp(year=2024, month=1, day=1)
|
|
PANDAS_HOUR_PRECISION = pd.Timestamp(year=2024, month=1, day=1, hour=1)
|
|
PANDAS_MIN_PRECISION = pd.Timestamp(year=2024, month=1, day=1, minute=1)
|
|
PANDAS_SEC_PRECISION = pd.Timestamp(year=2024, month=1, day=1, second=1)
|
|
PANDAS_MILLISEC_PRECISION = pd.Timestamp(year=2024, month=1, day=1, microsecond=1000)
|
|
PANDAS_MICROSEC_PRECISION = pd.Timestamp(year=2024, month=1, day=1, microsecond=1)
|
|
PANDAS_NANOSEC_PRECISION = pd.Timestamp(
|
|
year=2024, month=1, day=1, hour=1, minute=1, second=1
|
|
) + pd.Timedelta(nanoseconds=1)
|
|
|
|
# Define numpy.datetime64 values for comparison
|
|
DATETIME64_DAY_PRECISION = np.datetime64("2024-01-01")
|
|
DATETIME64_HOUR_PRECISION = np.datetime64("2024-01-01T01:00", "s")
|
|
DATETIME64_MIN_PRECISION = np.datetime64("2024-01-01T00:01", "s")
|
|
DATETIME64_SEC_PRECISION = np.datetime64("2024-01-01T00:00:01")
|
|
DATETIME64_MILLISEC_PRECISION = np.datetime64("2024-01-01T00:00:00.001")
|
|
DATETIME64_MICROSEC_PRECISION = np.datetime64("2024-01-01T00:00:00.000001")
|
|
DATETIME64_NANOSEC_PRECISION = np.datetime64("2024-01-01T01:01:01.000000001", "ns")
|
|
|
|
|
|
# Parametrized test to validate datetime values and expected numpy.datetime64 results
|
|
@pytest.mark.parametrize(
|
|
"data,expected_output",
|
|
[
|
|
(
|
|
[DATETIME_DAY_PRECISION],
|
|
np.array([DATETIME64_DAY_PRECISION], dtype="datetime64[s]"),
|
|
),
|
|
([DATETIME_HOUR_PRECISION], np.array([DATETIME64_HOUR_PRECISION])),
|
|
([DATETIME_MIN_PRECISION], np.array([DATETIME64_MIN_PRECISION])),
|
|
([DATETIME_SEC_PRECISION], np.array([DATETIME64_SEC_PRECISION])),
|
|
([DATETIME_MILLISEC_PRECISION], np.array([DATETIME64_MILLISEC_PRECISION])),
|
|
([DATETIME_MICROSEC_PRECISION], np.array([DATETIME64_MICROSEC_PRECISION])),
|
|
(
|
|
[DATETIME_MICROSEC_PRECISION, DATETIME_MILLISEC_PRECISION],
|
|
np.array(
|
|
[DATETIME64_MICROSEC_PRECISION, DATETIME64_MILLISEC_PRECISION],
|
|
dtype="datetime64[us]",
|
|
),
|
|
),
|
|
(
|
|
[DATETIME_SEC_PRECISION, DATETIME_MILLISEC_PRECISION],
|
|
np.array(
|
|
[DATETIME64_SEC_PRECISION, DATETIME64_MILLISEC_PRECISION],
|
|
dtype="datetime64[ms]",
|
|
),
|
|
),
|
|
(
|
|
[DATETIME_DAY_PRECISION, DATETIME_SEC_PRECISION],
|
|
np.array(
|
|
[DATETIME64_DAY_PRECISION, DATETIME64_SEC_PRECISION],
|
|
dtype="datetime64[s]",
|
|
),
|
|
),
|
|
(
|
|
[PANDAS_DAY_PRECISION],
|
|
np.array([DATETIME64_DAY_PRECISION], dtype="datetime64[s]"),
|
|
),
|
|
([PANDAS_HOUR_PRECISION], np.array([DATETIME64_HOUR_PRECISION])),
|
|
([PANDAS_MIN_PRECISION], np.array([DATETIME64_MIN_PRECISION])),
|
|
([PANDAS_SEC_PRECISION], np.array([DATETIME64_SEC_PRECISION])),
|
|
([PANDAS_MILLISEC_PRECISION], np.array([DATETIME64_MILLISEC_PRECISION])),
|
|
([PANDAS_MICROSEC_PRECISION], np.array([DATETIME64_MICROSEC_PRECISION])),
|
|
([PANDAS_NANOSEC_PRECISION], np.array([DATETIME64_NANOSEC_PRECISION])),
|
|
(
|
|
[PANDAS_NANOSEC_PRECISION, PANDAS_MICROSEC_PRECISION],
|
|
np.array(
|
|
[DATETIME64_NANOSEC_PRECISION, DATETIME64_MICROSEC_PRECISION],
|
|
dtype="datetime64[ns]",
|
|
),
|
|
),
|
|
(
|
|
[PANDAS_MICROSEC_PRECISION, PANDAS_MILLISEC_PRECISION],
|
|
np.array(
|
|
[DATETIME64_MICROSEC_PRECISION, DATETIME64_MILLISEC_PRECISION],
|
|
dtype="datetime64[us]",
|
|
),
|
|
),
|
|
(
|
|
[PANDAS_SEC_PRECISION, PANDAS_MILLISEC_PRECISION],
|
|
np.array(
|
|
[DATETIME64_SEC_PRECISION, DATETIME64_MILLISEC_PRECISION],
|
|
dtype="datetime64[ms]",
|
|
),
|
|
),
|
|
(
|
|
[PANDAS_DAY_PRECISION, PANDAS_SEC_PRECISION],
|
|
np.array(
|
|
[DATETIME64_DAY_PRECISION, DATETIME64_SEC_PRECISION],
|
|
dtype="datetime64[s]",
|
|
),
|
|
),
|
|
],
|
|
)
|
|
def test_list_of_datetimes(
|
|
data, expected_output, ray_start_regular_shared, restore_data_context
|
|
):
|
|
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
output = do_map_batches(data)
|
|
assert_structure_equals(output, expected_output)
|
|
|
|
|
|
def test_array_like(ray_start_regular_shared, restore_data_context):
|
|
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
data = torch.Tensor([1, 2, 3])
|
|
output = do_map_batches(data)
|
|
assert_structure_equals(output, np.array([1.0, 2.0, 3.0], dtype=np.float32))
|
|
|
|
|
|
def test_list_of_arrays(ray_start_regular_shared, restore_data_context):
|
|
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
data = [np.array([1, 2, 3]), np.array([4, 5, 6])]
|
|
output = do_map_batches(data)
|
|
assert_structure_equals(output, np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64))
|
|
|
|
|
|
def test_list_of_array_like(ray_start_regular_shared, restore_data_context):
|
|
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
data = [torch.Tensor([1, 2, 3]), torch.Tensor([4, 5, 6])]
|
|
output = do_map_batches(data)
|
|
assert_structure_equals(output, np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32))
|
|
|
|
|
|
def test_ragged_tensors_map_batches(ray_start_regular_shared, restore_data_context):
|
|
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
data = [torch.Tensor([1, 2, 3]), torch.Tensor([1, 2])]
|
|
output = do_map_batches(data)
|
|
assert_structure_equals(
|
|
output, create_ragged_ndarray([np.array([1, 2, 3]), np.array([1, 2])])
|
|
)
|
|
|
|
data = [torch.zeros((3, 5, 10)), torch.zeros((3, 8, 8))]
|
|
output = do_map_batches(data)
|
|
assert_structure_equals(
|
|
output, create_ragged_ndarray([np.zeros((3, 5, 10)), np.zeros((3, 8, 8))])
|
|
)
|
|
|
|
|
|
def test_scalar_nested_arrays(ray_start_regular_shared, restore_data_context):
|
|
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
data = [[[1]], [[2]]]
|
|
output = do_map_batches(data)
|
|
|
|
assert_structure_equals(
|
|
output,
|
|
create_ragged_ndarray(
|
|
[np.array([1], dtype=np.object_), np.array([2], dtype=np.object_)]
|
|
),
|
|
)
|
|
|
|
|
|
def test_scalar_lists_not_converted(ray_start_regular_shared, restore_data_context):
|
|
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
data = [[1, 2], [1, 2]]
|
|
output = do_map_batches(data)
|
|
assert_structure_equals(
|
|
output, create_ragged_ndarray([np.array([1, 2]), np.array([1, 2])])
|
|
)
|
|
|
|
data = [[1, 2, 3], [1, 2]]
|
|
output = do_map_batches(data)
|
|
assert_structure_equals(
|
|
output, create_ragged_ndarray([np.array([1, 2, 3]), np.array([1, 2])])
|
|
)
|
|
|
|
|
|
def test_scalar_numpy(ray_start_regular_shared, restore_data_context):
|
|
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
data = np.int64(1)
|
|
ds = ray.data.range(2, override_num_blocks=1)
|
|
ds = ds.map(lambda x: {"output": data})
|
|
output = ds.take_batch()["output"]
|
|
assert_structure_equals(output, np.array([1, 1], dtype=np.int64))
|
|
|
|
|
|
def test_scalar_arrays(ray_start_regular_shared, restore_data_context):
|
|
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
data = np.array([1, 2, 3])
|
|
ds = ray.data.range(2, override_num_blocks=1)
|
|
ds = ds.map(lambda x: {"output": data})
|
|
output = ds.take_batch()["output"]
|
|
assert_structure_equals(output, np.array([[1, 2, 3], [1, 2, 3]], dtype=np.int64))
|
|
|
|
|
|
def test_bytes(ray_start_regular_shared, restore_data_context):
|
|
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
"""Tests that bytes are converted to object dtype instead of zero-terminated."""
|
|
data = b"\x1a\n\x00\n\x1a"
|
|
ds = ray.data.range(1, override_num_blocks=1)
|
|
ds = ds.map(lambda x: {"output": data})
|
|
output = ds.take_batch()["output"]
|
|
assert_structure_equals(output, np.array([b"\x1a\n\x00\n\x1a"], dtype=object))
|
|
|
|
|
|
def test_uniform_tensors(ray_start_regular_shared, restore_data_context):
|
|
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
data = torch.Tensor([1, 2, 3])
|
|
ds = ray.data.range(2, override_num_blocks=1)
|
|
ds = ds.map(lambda x: {"output": data})
|
|
output = ds.take_batch()["output"]
|
|
assert_structure_equals(output, np.array([[1, 2, 3], [1, 2, 3]], dtype=np.float32))
|
|
|
|
|
|
def test_scalar_ragged_arrays(ray_start_regular_shared, restore_data_context):
|
|
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
data = [np.array([1, 2, 3]), np.array([1, 2])]
|
|
ds = ray.data.range(2, override_num_blocks=1)
|
|
ds = ds.map(lambda x: {"output": data[x["id"]]})
|
|
output = ds.take_batch()["output"]
|
|
|
|
assert_structure_equals(
|
|
output, np.array([np.array([1, 2, 3]), np.array([1, 2])], dtype=object)
|
|
)
|
|
|
|
|
|
def test_ragged_tensors(ray_start_regular_shared, restore_data_context):
|
|
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
data = [torch.Tensor([1, 2, 3]), torch.Tensor([1, 2])]
|
|
ds = ray.data.range(2, override_num_blocks=1)
|
|
ds = ds.map(lambda x: {"output": data[x["id"]]})
|
|
output = ds.take_batch()["output"]
|
|
assert_structure_equals(
|
|
output, np.array([np.array([1, 2, 3]), np.array([1, 2])], dtype=object)
|
|
)
|
|
|
|
data = [torch.zeros((3, 5, 10)), torch.zeros((3, 8, 8))]
|
|
ds = ray.data.range(2, override_num_blocks=1)
|
|
ds = ds.map(lambda x: {"output": data[x["id"]]})
|
|
output = ds.take_batch()["output"]
|
|
assert_structure_equals(
|
|
output, create_ragged_ndarray([np.zeros((3, 5, 10)), np.zeros((3, 8, 8))])
|
|
)
|
|
|
|
|
|
def test_nested_ragged_arrays(ray_start_regular_shared, restore_data_context):
|
|
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
data = [
|
|
{"a": [[1], [2, 3]]},
|
|
{"a": [[4, 5], [6]]},
|
|
]
|
|
|
|
def f(row):
|
|
return data[row["id"]]
|
|
|
|
output = ray.data.range(2).map(f).take_all()
|
|
assert output == data
|
|
|
|
|
|
# https://github.com/ray-project/ray/issues/35340
|
|
def test_complex_ragged_arrays(ray_start_regular_shared, restore_data_context):
|
|
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
data = [[{"a": 1}, {"a": 2}, {"a": 3}], [{"b": 1}]]
|
|
output = do_map_batches(data)
|
|
|
|
# Assert resulting objects are coerced to appropriate shape, following
|
|
# table's schema
|
|
assert_structure_equals(
|
|
output,
|
|
create_ragged_ndarray(
|
|
[
|
|
np.array(
|
|
[{"a": 1, "b": None}, {"a": 2, "b": None}, {"a": 3, "b": None}]
|
|
),
|
|
np.array([{"a": None, "b": 1}]),
|
|
]
|
|
),
|
|
)
|
|
|
|
data = ["hi", 1, None, [[[[]]]], {"a": [[{"b": 2, "c": UserObj()}]]}, UserObj()]
|
|
output = do_map_batches(data)
|
|
assert_structure_equals(output, create_ragged_ndarray(data))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|