wehub-resource-sync
1073 lines
36 KiB
Python
1073 lines
36 KiB
Python
import math
|
|
import time
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data._internal.tensor_extensions.arrow import (
|
|
MIN_PYARROW_VERSION_FIXED_SHAPE_TENSOR_SCALAR,
|
|
ArrowTensorArray,
|
|
FixedShapeTensorFormat,
|
|
create_arrow_fixed_shape_tensor_type,
|
|
)
|
|
from ray.data._internal.tensor_extensions.utils import _create_possibly_ragged_ndarray
|
|
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
|
|
from ray.data.block import BlockAccessor
|
|
from ray.data.context import DataContext
|
|
from ray.data.dataset import Schema
|
|
from ray.data.extensions.tensor_extension import (
|
|
ArrowTensorType,
|
|
ArrowTensorTypeV2,
|
|
ArrowVariableShapedTensorArray,
|
|
ArrowVariableShapedTensorType,
|
|
FixedShapeTensorType,
|
|
TensorArray,
|
|
TensorDtype,
|
|
)
|
|
from ray.data.tests.conftest import * # noqa
|
|
from ray.data.tests.util import extract_values
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
|
|
# https://github.com/ray-project/ray/issues/33695
|
|
def test_large_tensor_creation(ray_start_regular_shared, tensor_format_context):
|
|
"""Tests that large tensor read task creation can complete successfully without
|
|
hanging."""
|
|
start_time = time.time()
|
|
ray.data.range_tensor(1000, override_num_blocks=1000, shape=(80, 80, 100, 100))
|
|
end_time = time.time()
|
|
|
|
# Should not take more than 20 seconds.
|
|
assert end_time - start_time < 20
|
|
|
|
|
|
def test_tensors_basic(ray_start_regular_shared, tensor_format_context):
|
|
|
|
# Determine expected tensor type based on format
|
|
expected_type = create_arrow_fixed_shape_tensor_type(shape=(3, 5), dtype=pa.int64())
|
|
|
|
# Create directly.
|
|
tensor_shape = (3, 5)
|
|
ds = ray.data.range_tensor(6, shape=tensor_shape, override_num_blocks=6)
|
|
assert ds.count() == 6
|
|
assert ds.schema() == Schema(pa.schema([("data", expected_type)]))
|
|
# The actual size is slightly larger due to metadata.
|
|
# We add 6 (one per tensor) offset values of 8 bytes each to account for the
|
|
# in-memory representation of the PyArrow LargeList type
|
|
assert math.isclose(ds.size_bytes(), 5 * 3 * 6 * 8 + 6 * 8, rel_tol=0.1)
|
|
|
|
# Test row iterator yields tensors.
|
|
for tensor in ds.iter_rows():
|
|
tensor = tensor["data"]
|
|
assert isinstance(tensor, np.ndarray)
|
|
assert tensor.shape == tensor_shape
|
|
|
|
# Test batch iterator yields tensors.
|
|
for tensor in ds.iter_batches(batch_size=2):
|
|
tensor = tensor["data"]
|
|
assert isinstance(tensor, np.ndarray)
|
|
assert tensor.shape == (2,) + tensor_shape
|
|
|
|
# Native format.
|
|
def np_mapper(arr):
|
|
if "data" in arr:
|
|
arr = arr["data"]
|
|
else:
|
|
arr = arr["id"]
|
|
assert isinstance(arr, np.ndarray)
|
|
return {"data": arr + 1}
|
|
|
|
res = ray.data.range_tensor(2, shape=(2, 2)).map(np_mapper).take()
|
|
np.testing.assert_equal(
|
|
extract_values("data", res), [np.ones((2, 2)), 2 * np.ones((2, 2))]
|
|
)
|
|
|
|
# Explicit NumPy format.
|
|
res = (
|
|
ray.data.range_tensor(2, shape=(2, 2))
|
|
.map_batches(np_mapper, batch_format="numpy")
|
|
.take()
|
|
)
|
|
np.testing.assert_equal(
|
|
extract_values("data", res), [np.ones((2, 2)), 2 * np.ones((2, 2))]
|
|
)
|
|
|
|
# Pandas conversion.
|
|
def pd_mapper(df):
|
|
assert isinstance(df, pd.DataFrame)
|
|
return df + 2
|
|
|
|
res = ray.data.range_tensor(2).map_batches(pd_mapper, batch_format="pandas").take()
|
|
np.testing.assert_equal(extract_values("data", res), [np.array([2]), np.array([3])])
|
|
|
|
# Arrow columns in NumPy format.
|
|
def multi_mapper(col_arrs):
|
|
assert isinstance(col_arrs, dict)
|
|
assert list(col_arrs.keys()) == ["a", "b", "c"]
|
|
assert all(isinstance(col_arr, np.ndarray) for col_arr in col_arrs.values())
|
|
return {"a": col_arrs["a"] + 1, "b": col_arrs["b"] + 1, "c": col_arrs["c"] + 1}
|
|
|
|
# Multiple columns.
|
|
t = pa.table(
|
|
{
|
|
"a": [1, 2, 3],
|
|
"b": [4.0, 5.0, 6.0],
|
|
"c": ArrowTensorArray.from_numpy(np.array([[1, 2], [3, 4], [5, 6]])),
|
|
}
|
|
)
|
|
res = (
|
|
ray.data.from_arrow(t)
|
|
.map_batches(multi_mapper, batch_size=2, batch_format="numpy")
|
|
.take()
|
|
)
|
|
np.testing.assert_equal(
|
|
res,
|
|
[
|
|
{"a": 2, "b": 5.0, "c": np.array([2, 3])},
|
|
{"a": 3, "b": 6.0, "c": np.array([4, 5])},
|
|
{"a": 4, "b": 7.0, "c": np.array([6, 7])},
|
|
],
|
|
)
|
|
|
|
def single_mapper(col_arrs):
|
|
assert isinstance(col_arrs, dict)
|
|
assert list(col_arrs.keys()) == ["c"]
|
|
assert all(isinstance(col_arr, np.ndarray) for col_arr in col_arrs.values())
|
|
return {"c": col_arrs["c"] + 1}
|
|
|
|
# Single column (should still yield ndarray dict batches).
|
|
t = t.select(["c"])
|
|
res = (
|
|
ray.data.from_arrow(t)
|
|
.map_batches(single_mapper, batch_size=2, batch_format="numpy")
|
|
.take()
|
|
)
|
|
np.testing.assert_equal(
|
|
res,
|
|
[
|
|
{"c": np.array([2, 3])},
|
|
{"c": np.array([4, 5])},
|
|
{"c": np.array([6, 7])},
|
|
],
|
|
)
|
|
|
|
# Pandas columns in NumPy format.
|
|
def multi_mapper(col_arrs):
|
|
assert isinstance(col_arrs, dict)
|
|
assert list(col_arrs.keys()) == ["a", "b", "c"]
|
|
assert all(isinstance(col_arr, np.ndarray) for col_arr in col_arrs.values())
|
|
return pd.DataFrame(
|
|
{
|
|
"a": col_arrs["a"] + 1,
|
|
"b": col_arrs["b"] + 1,
|
|
"c": TensorArray(col_arrs["c"] + 1),
|
|
}
|
|
)
|
|
|
|
# Multiple columns.
|
|
df = pd.DataFrame(
|
|
{
|
|
"a": [1, 2, 3],
|
|
"b": [4.0, 5.0, 6.0],
|
|
"c": TensorArray(np.array([[1, 2], [3, 4], [5, 6]])),
|
|
}
|
|
)
|
|
res = (
|
|
ray.data.from_pandas(df)
|
|
.map_batches(multi_mapper, batch_size=2, batch_format="numpy")
|
|
.take()
|
|
)
|
|
np.testing.assert_equal(
|
|
res,
|
|
[
|
|
{"a": 2, "b": 5.0, "c": np.array([2, 3])},
|
|
{"a": 3, "b": 6.0, "c": np.array([4, 5])},
|
|
{"a": 4, "b": 7.0, "c": np.array([6, 7])},
|
|
],
|
|
)
|
|
|
|
# Single column (should still yield ndarray dict batches).
|
|
def single_mapper(col_arrs):
|
|
assert isinstance(col_arrs, dict)
|
|
assert list(col_arrs.keys()) == ["c"]
|
|
assert all(isinstance(col_arr, np.ndarray) for col_arr in col_arrs.values())
|
|
return pd.DataFrame({"c": TensorArray(col_arrs["c"] + 1)})
|
|
|
|
df = df[["c"]]
|
|
res = (
|
|
ray.data.from_pandas(df)
|
|
.map_batches(single_mapper, batch_size=2, batch_format="numpy")
|
|
.take()
|
|
)
|
|
np.testing.assert_equal(
|
|
res,
|
|
[
|
|
{"c": np.array([2, 3])},
|
|
{"c": np.array([4, 5])},
|
|
{"c": np.array([6, 7])},
|
|
],
|
|
)
|
|
|
|
# Simple dataset in NumPy format.
|
|
def mapper(arr):
|
|
arr = np_mapper(arr)
|
|
return arr
|
|
|
|
res = (
|
|
ray.data.range(10, override_num_blocks=2)
|
|
.map_batches(mapper, batch_format="numpy")
|
|
.take()
|
|
)
|
|
assert extract_values("data", res) == list(range(1, 11))
|
|
|
|
|
|
def test_batch_tensors(ray_start_regular_shared, tensor_format_context):
|
|
ds = ray.data.from_items(
|
|
[np.array([0, 0]) for _ in range(40)], override_num_blocks=40
|
|
)
|
|
|
|
batch = next(iter(ds.iter_batches()))
|
|
assert set(batch) == {"item"}
|
|
assert batch["item"].shape == (40, 2)
|
|
|
|
|
|
def test_tensors_shuffle(ray_start_regular_shared, tensor_format_context):
|
|
# Test Arrow table representation.
|
|
tensor_shape = (3, 5)
|
|
ds = ray.data.range_tensor(6, shape=tensor_shape)
|
|
shuffled_ds = ds.random_shuffle()
|
|
shuffled = extract_values("data", shuffled_ds.take())
|
|
base = extract_values("data", ds.take())
|
|
np.testing.assert_raises(
|
|
AssertionError,
|
|
np.testing.assert_equal,
|
|
shuffled,
|
|
base,
|
|
)
|
|
np.testing.assert_equal(
|
|
sorted(shuffled, key=lambda arr: arr.min()),
|
|
sorted(base, key=lambda arr: arr.min()),
|
|
)
|
|
|
|
# Test Pandas table representation.
|
|
tensor_shape = (3, 5)
|
|
ds = ray.data.range_tensor(6, shape=tensor_shape)
|
|
ds = ds.map_batches(lambda df: df, batch_format="pandas")
|
|
shuffled_ds = ds.random_shuffle()
|
|
shuffled = extract_values("data", shuffled_ds.take())
|
|
base = extract_values("data", ds.take())
|
|
np.testing.assert_raises(
|
|
AssertionError,
|
|
np.testing.assert_equal,
|
|
shuffled,
|
|
base,
|
|
)
|
|
np.testing.assert_equal(
|
|
sorted(shuffled, key=lambda arr: arr.min()),
|
|
sorted(base, key=lambda arr: arr.min()),
|
|
)
|
|
|
|
|
|
def test_tensors_sort(ray_start_regular_shared, tensor_format_context):
|
|
# Test Arrow table representation.
|
|
t = pa.table({"a": TensorArray(np.arange(32).reshape((2, 4, 4))), "b": [1, 2]})
|
|
ds = ray.data.from_arrow(t)
|
|
sorted_ds = ds.sort(key="b", descending=True)
|
|
sorted_arrs = [row["a"] for row in sorted_ds.take()]
|
|
base = [row["a"] for row in ds.take()]
|
|
np.testing.assert_raises(
|
|
AssertionError,
|
|
np.testing.assert_equal,
|
|
sorted_arrs,
|
|
base,
|
|
)
|
|
np.testing.assert_equal(
|
|
sorted_arrs,
|
|
sorted(base, key=lambda arr: -arr.min()),
|
|
)
|
|
|
|
# Test Pandas table representation.
|
|
df = pd.DataFrame({"a": TensorArray(np.arange(32).reshape((2, 4, 4))), "b": [1, 2]})
|
|
ds = ray.data.from_pandas(df)
|
|
sorted_ds = ds.sort(key="b", descending=True)
|
|
sorted_arrs = [np.asarray(row["a"]) for row in sorted_ds.take()]
|
|
base = [np.asarray(row["a"]) for row in ds.take()]
|
|
np.testing.assert_raises(
|
|
AssertionError,
|
|
np.testing.assert_equal,
|
|
sorted_arrs,
|
|
base,
|
|
)
|
|
np.testing.assert_equal(
|
|
sorted_arrs,
|
|
sorted(base, key=lambda arr: -arr.min()),
|
|
)
|
|
|
|
|
|
def test_tensors_inferred_from_map(ray_start_regular_shared, tensor_format_context):
|
|
tensor_format = tensor_format_context
|
|
# With tensor_format_context, ARROW_NATIVE only runs when supported,
|
|
# so to_type() is safe to use
|
|
expected_type = tensor_format.to_type()
|
|
|
|
# Test map.
|
|
ds = ray.data.range(10, override_num_blocks=10).map(
|
|
lambda _: {"data": np.ones((4, 4))}
|
|
)
|
|
ds = ds.materialize()
|
|
assert ds.count() == 10
|
|
schema = ds.schema()
|
|
assert schema.names == ["data"]
|
|
dtype = schema.types[0]
|
|
assert isinstance(dtype, expected_type)
|
|
assert tuple(dtype.shape) == (4, 4)
|
|
assert dtype.value_type == pa.float64()
|
|
|
|
# Test map_batches.
|
|
ds = ray.data.range(16, override_num_blocks=4).map_batches(
|
|
lambda _: {"data": np.ones((3, 4, 4))}, batch_size=2
|
|
)
|
|
ds = ds.materialize()
|
|
assert ds.count() == 24
|
|
schema = ds.schema()
|
|
assert schema.names == ["data"]
|
|
dtype = schema.types[0]
|
|
assert isinstance(dtype, expected_type)
|
|
assert tuple(dtype.shape) == (4, 4)
|
|
assert dtype.value_type == pa.float64()
|
|
|
|
# Test flat_map.
|
|
ds = ray.data.range(10, override_num_blocks=10).flat_map(
|
|
lambda _: [{"data": np.ones((4, 4))}, {"data": np.ones((4, 4))}]
|
|
)
|
|
ds = ds.materialize()
|
|
assert ds.count() == 20
|
|
schema = ds.schema()
|
|
assert schema.names == ["data"]
|
|
dtype = schema.types[0]
|
|
|
|
assert isinstance(dtype, expected_type)
|
|
assert tuple(dtype.shape) == (4, 4)
|
|
assert dtype.value_type == pa.float64()
|
|
|
|
# Test map_batches ndarray column.
|
|
ds = ray.data.range(16, override_num_blocks=4).map_batches(
|
|
lambda _: pd.DataFrame({"a": [np.ones((4, 4))] * 3}), batch_size=2
|
|
)
|
|
ds = ds.materialize()
|
|
assert ds.count() == 24
|
|
schema = ds.schema()
|
|
assert schema.names == ["a"]
|
|
dtype = schema.types[0]
|
|
assert isinstance(dtype, expected_type)
|
|
assert tuple(dtype.shape) == (4, 4)
|
|
assert dtype.value_type == pa.float64()
|
|
|
|
ds = ray.data.range(16, override_num_blocks=4).map_batches(
|
|
lambda _: pd.DataFrame({"a": [np.ones((2, 2)), np.ones((3, 3))]}),
|
|
batch_size=2,
|
|
)
|
|
ds = ds.materialize()
|
|
assert ds.count() == 16
|
|
schema = ds.schema()
|
|
assert schema.names == ["a"]
|
|
dtype = schema.types[0]
|
|
assert isinstance(dtype, ArrowVariableShapedTensorType)
|
|
assert tuple(dtype.shape) == (None, None)
|
|
assert dtype.value_type == pa.float64()
|
|
|
|
|
|
def test_tensor_array_block_slice(tensor_format_context):
|
|
tensor_format = tensor_format_context
|
|
|
|
# Test that ArrowBlock slicing works with tensor column extension type.
|
|
def check_for_copy(table1, table2, a, b, is_copy):
|
|
expected_slice = table1.slice(a, b - a)
|
|
assert table2.equals(expected_slice)
|
|
assert table2.schema == table1.schema
|
|
assert table1.num_columns == table2.num_columns
|
|
for col1, col2 in zip(table1.columns, table2.columns):
|
|
assert col1.num_chunks == col2.num_chunks
|
|
for chunk1, chunk2 in zip(col1.chunks, col2.chunks):
|
|
bufs1 = chunk1.buffers()
|
|
bufs2 = chunk2.buffers()
|
|
expected_offset = 0 if is_copy else a
|
|
assert chunk2.offset == expected_offset
|
|
assert len(chunk2) == b - a
|
|
index = 1
|
|
if (
|
|
tensor_format == FixedShapeTensorFormat.ARROW_NATIVE
|
|
and FixedShapeTensorType is not None
|
|
):
|
|
# The buffer for native fixed shaped tensors sits at index 2, not 1
|
|
index = 2
|
|
if is_copy:
|
|
assert bufs2[index].address != bufs1[index].address
|
|
else:
|
|
assert bufs2[index].address == bufs1[index].address
|
|
|
|
n = 20
|
|
one_arr = np.arange(4 * n).reshape(n, 2, 2)
|
|
df = pd.DataFrame({"one": TensorArray(one_arr), "two": ["a"] * n})
|
|
table = pa.Table.from_pandas(df)
|
|
a, b = 5, 10
|
|
block_accessor = BlockAccessor.for_block(table)
|
|
|
|
# Test with copy.
|
|
table2 = block_accessor.slice(a, b, True)
|
|
res = table2["one"].chunk(0).to_numpy_ndarray()
|
|
|
|
np.testing.assert_array_equal(res, one_arr[a:b, :, :])
|
|
check_for_copy(table, table2, a, b, is_copy=True)
|
|
|
|
# Test without copy. arrow_native requires a copy
|
|
table2 = block_accessor.slice(a, b, False)
|
|
res = table2["one"].chunk(0).to_numpy_ndarray()
|
|
np.testing.assert_array_equal(res, one_arr[a:b, :, :])
|
|
check_for_copy(table, table2, a, b, is_copy=False)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_data,a,b",
|
|
[
|
|
([[False, True], [True, False], [True, True], [False, False]], 1, 3),
|
|
([[False, True], [True, False], [True, True], [False, False]], 0, 1),
|
|
(
|
|
[
|
|
[False, True],
|
|
[True, False],
|
|
[True, True],
|
|
[False, False],
|
|
[True, False],
|
|
[False, False],
|
|
[False, True],
|
|
[True, True],
|
|
[False, False],
|
|
[True, True],
|
|
[False, True],
|
|
[True, False],
|
|
],
|
|
3,
|
|
6,
|
|
),
|
|
(
|
|
[
|
|
[False, True],
|
|
[True, False],
|
|
[True, True],
|
|
[False, False],
|
|
[True, False],
|
|
[False, False],
|
|
[False, True],
|
|
[True, True],
|
|
[False, False],
|
|
[True, True],
|
|
[False, True],
|
|
[True, False],
|
|
],
|
|
7,
|
|
11,
|
|
),
|
|
(
|
|
[
|
|
[False, True],
|
|
[True, False],
|
|
[True, True],
|
|
[False, False],
|
|
[True, False],
|
|
[False, False],
|
|
[False, True],
|
|
[True, True],
|
|
[False, False],
|
|
[True, True],
|
|
[False, True],
|
|
[True, False],
|
|
],
|
|
9,
|
|
12,
|
|
),
|
|
# Variable-shaped tensors.
|
|
(
|
|
[[False, True], [True, False, True], [False], [False, False, True, True]],
|
|
1,
|
|
3,
|
|
),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("init_with_pandas", [True, False])
|
|
def test_tensor_array_boolean_slice_pandas_roundtrip(init_with_pandas, test_data, a, b):
|
|
is_variable_shaped = len({len(elem) for elem in test_data}) > 1
|
|
n = len(test_data)
|
|
test_arr = _create_possibly_ragged_ndarray(test_data)
|
|
df = pd.DataFrame({"one": TensorArray(test_arr), "two": ["a"] * n})
|
|
if init_with_pandas:
|
|
table = pa.Table.from_pandas(df)
|
|
else:
|
|
if is_variable_shaped:
|
|
col = ArrowVariableShapedTensorArray.from_numpy(test_arr)
|
|
else:
|
|
col = ArrowTensorArray.from_numpy(test_arr)
|
|
table = pa.table({"one": col, "two": ["a"] * n})
|
|
block_accessor = BlockAccessor.for_block(table)
|
|
|
|
# Test without copy.
|
|
table2 = block_accessor.slice(a, b, False)
|
|
out = table2["one"].chunk(0).to_numpy()
|
|
expected = test_arr[a:b]
|
|
if is_variable_shaped:
|
|
for o, e in zip(out, expected):
|
|
np.testing.assert_array_equal(o, e)
|
|
else:
|
|
np.testing.assert_array_equal(out, expected)
|
|
pd.testing.assert_frame_equal(
|
|
table2.to_pandas().reset_index(drop=True), df[a:b].reset_index(drop=True)
|
|
)
|
|
|
|
# Test with copy.
|
|
table2 = block_accessor.slice(a, b, True)
|
|
out = table2["one"].chunk(0).to_numpy()
|
|
expected = test_arr[a:b]
|
|
if is_variable_shaped:
|
|
for o, e in zip(out, expected):
|
|
np.testing.assert_array_equal(o, e)
|
|
else:
|
|
np.testing.assert_array_equal(out, expected)
|
|
pd.testing.assert_frame_equal(
|
|
table2.to_pandas().reset_index(drop=True), df[a:b].reset_index(drop=True)
|
|
)
|
|
|
|
|
|
def test_tensors_in_tables_from_pandas(ray_start_regular_shared, tensor_format_context):
|
|
|
|
outer_dim = 3
|
|
inner_shape = (2, 2, 2)
|
|
shape = (outer_dim,) + inner_shape
|
|
num_items = np.prod(np.array(shape))
|
|
arr = np.arange(num_items).reshape(shape)
|
|
df = pd.DataFrame({"one": list(range(outer_dim)), "two": list(arr)})
|
|
# Cast column to tensor extension dtype.
|
|
df["two"] = df["two"].astype(TensorDtype(shape, np.int64))
|
|
ds = ray.data.from_pandas([df])
|
|
values = [[s["one"], s["two"]] for s in ds.take()]
|
|
expected = list(zip(list(range(outer_dim)), arr))
|
|
for v, e in zip(sorted(values), expected):
|
|
np.testing.assert_equal(v, e)
|
|
|
|
|
|
def test_tensors_in_tables_from_pandas_variable_shaped(
|
|
ray_start_regular_shared, tensor_format_context
|
|
):
|
|
|
|
shapes = [(2, 2), (3, 3), (4, 4)]
|
|
cumsum_sizes = np.cumsum([0] + [np.prod(shape) for shape in shapes[:-1]])
|
|
arrs = [
|
|
np.arange(offset, offset + np.prod(shape)).reshape(shape)
|
|
for offset, shape in zip(cumsum_sizes, shapes)
|
|
]
|
|
outer_dim = len(arrs)
|
|
df = pd.DataFrame({"one": list(range(outer_dim)), "two": arrs})
|
|
# Cast column to tensor extension dtype.
|
|
df["two"] = df["two"].astype(TensorDtype(None, np.int64))
|
|
ds = ray.data.from_pandas(df)
|
|
values = [[s["one"], s["two"]] for s in ds.take()]
|
|
expected = list(zip(range(outer_dim), arrs))
|
|
for v, e in zip(sorted(values), expected):
|
|
np.testing.assert_equal(v, e)
|
|
|
|
|
|
def test_tensors_in_tables_pandas_roundtrip(
|
|
ray_start_regular_shared,
|
|
enable_automatic_tensor_extension_cast,
|
|
tensor_format_context,
|
|
):
|
|
outer_dim = 3
|
|
inner_shape = (2, 2, 2)
|
|
shape = (outer_dim,) + inner_shape
|
|
num_items = np.prod(np.array(shape))
|
|
arr = np.arange(num_items).reshape(shape)
|
|
df = pd.DataFrame({"one": list(range(outer_dim)), "two": TensorArray(arr)})
|
|
ds = ray.data.from_pandas(df)
|
|
ds = ds.map_batches(lambda df: df + 1, batch_size=2, batch_format="pandas")
|
|
ds_df = ds.to_pandas()
|
|
expected_df = df + 1
|
|
if enable_automatic_tensor_extension_cast:
|
|
expected_df["two"] = list(expected_df["two"].to_numpy())
|
|
# Roundtrip may use Arrow-backed dtypes (e.g. int64[pyarrow]) for plain columns.
|
|
expected_df["one"] = expected_df["one"].astype(ds_df["one"].dtype)
|
|
pd.testing.assert_frame_equal(ds_df, expected_df)
|
|
|
|
|
|
def test_tensors_in_tables_pandas_roundtrip_variable_shaped(
|
|
ray_start_regular_shared,
|
|
enable_automatic_tensor_extension_cast,
|
|
tensor_format_context,
|
|
):
|
|
shapes = [(2, 2), (3, 3), (4, 4)]
|
|
cumsum_sizes = np.cumsum([0] + [np.prod(shape) for shape in shapes[:-1]])
|
|
arrs = [
|
|
np.arange(offset, offset + np.prod(shape)).reshape(shape)
|
|
for offset, shape in zip(cumsum_sizes, shapes)
|
|
]
|
|
outer_dim = len(arrs)
|
|
df = pd.DataFrame({"one": list(range(outer_dim)), "two": TensorArray(arrs)})
|
|
ds = ray.data.from_pandas(df)
|
|
ds = ds.map_batches(lambda df: df + 1, batch_size=2, batch_format="pandas")
|
|
ds_df = ds.to_pandas()
|
|
expected_df = df + 1
|
|
if enable_automatic_tensor_extension_cast:
|
|
expected_df["two"] = _create_possibly_ragged_ndarray(
|
|
expected_df["two"].to_numpy()
|
|
)
|
|
# Roundtrip may use Arrow-backed dtypes (e.g. int64[pyarrow]) for plain columns.
|
|
expected_df["one"] = expected_df["one"].astype(ds_df["one"].dtype)
|
|
pd.testing.assert_frame_equal(ds_df, expected_df)
|
|
|
|
|
|
def test_tensors_in_tables_parquet_roundtrip(
|
|
ray_start_regular_shared, tmp_path, tensor_format_context
|
|
):
|
|
outer_dim = 3
|
|
inner_shape = (2, 2, 2)
|
|
shape = (outer_dim,) + inner_shape
|
|
num_items = np.prod(np.array(shape))
|
|
arr = np.arange(num_items).reshape(shape)
|
|
df = pd.DataFrame({"one": list(range(outer_dim)), "two": TensorArray(arr)})
|
|
ds = ray.data.from_pandas(df)
|
|
ds = ds.map_batches(lambda df: df + 1, batch_size=2, batch_format="pandas")
|
|
ds.write_parquet(str(tmp_path))
|
|
ds = ray.data.read_parquet(str(tmp_path))
|
|
values = [[s["one"], s["two"]] for s in ds.take()]
|
|
expected = list(zip(list(range(1, outer_dim + 1)), arr + 1))
|
|
for v, e in zip(sorted(values), expected):
|
|
np.testing.assert_equal(v, e)
|
|
|
|
|
|
def test_tensors_in_tables_parquet_roundtrip_variable_shaped(
|
|
ray_start_regular_shared, tmp_path, tensor_format_context
|
|
):
|
|
shapes = [(2, 2), (3, 3), (4, 4)]
|
|
cumsum_sizes = np.cumsum([0] + [np.prod(shape) for shape in shapes[:-1]])
|
|
arrs = [
|
|
np.arange(offset, offset + np.prod(shape)).reshape(shape)
|
|
for offset, shape in zip(cumsum_sizes, shapes)
|
|
]
|
|
outer_dim = len(arrs)
|
|
df = pd.DataFrame({"one": list(range(outer_dim)), "two": TensorArray(arrs)})
|
|
ds = ray.data.from_pandas(df)
|
|
ds = ds.map_batches(lambda df: df + 1, batch_size=2, batch_format="pandas")
|
|
ds.write_parquet(str(tmp_path))
|
|
ds = ray.data.read_parquet(str(tmp_path))
|
|
values = [[s["one"], s["two"]] for s in ds.take()]
|
|
expected = list(zip(list(range(1, outer_dim + 1)), [arr + 1 for arr in arrs]))
|
|
for v, e in zip(sorted(values), expected):
|
|
np.testing.assert_equal(v, e)
|
|
|
|
|
|
def test_tensors_in_tables_parquet_with_schema(
|
|
ray_start_regular_shared,
|
|
tmp_path,
|
|
tensor_format_context,
|
|
):
|
|
outer_dim = 3
|
|
inner_shape = (2, 2, 2)
|
|
shape = (outer_dim,) + inner_shape
|
|
num_items = np.prod(np.array(shape))
|
|
arr = np.arange(num_items).reshape(shape)
|
|
df = pd.DataFrame({"one": list(range(outer_dim)), "two": TensorArray(arr)})
|
|
ds = ray.data.from_pandas([df])
|
|
ds.write_parquet(str(tmp_path))
|
|
|
|
tensor_type = create_arrow_fixed_shape_tensor_type(
|
|
shape=inner_shape,
|
|
dtype=pa.from_numpy_dtype(arr.dtype),
|
|
)
|
|
schema = pa.schema([("one", pa.int32()), ("two", tensor_type)])
|
|
ds = ray.data.read_parquet(str(tmp_path), schema=schema)
|
|
values = [[s["one"], s["two"]] for s in ds.take()]
|
|
expected = list(zip(list(range(outer_dim)), arr))
|
|
for v, e in zip(sorted(values), expected):
|
|
np.testing.assert_equal(v, e)
|
|
|
|
|
|
def test_tensors_in_tables_parquet_pickle_manual_serde(
|
|
ray_start_regular_shared, tmp_path, tensor_format_context
|
|
):
|
|
import pickle
|
|
|
|
outer_dim = 3
|
|
inner_shape = (2, 2, 2)
|
|
shape = (outer_dim,) + inner_shape
|
|
num_items = np.prod(np.array(shape))
|
|
arr = np.arange(num_items).reshape(shape)
|
|
df = pd.DataFrame(
|
|
{"one": list(range(outer_dim)), "two": [pickle.dumps(a) for a in arr]}
|
|
)
|
|
ds = ray.data.from_pandas([df])
|
|
ds.write_parquet(str(tmp_path))
|
|
ds = ray.data.read_parquet(str(tmp_path))
|
|
|
|
# Manually deserialize the tensor pickle bytes and cast to our tensor
|
|
# extension type.
|
|
def deser_mapper(batch: pd.DataFrame):
|
|
batch["two"] = [pickle.loads(a) for a in batch["two"]]
|
|
batch["two"] = batch["two"].astype(TensorDtype(shape, np.int64))
|
|
return batch
|
|
|
|
casted_ds = ds.map_batches(deser_mapper, batch_format="pandas")
|
|
|
|
values = [[s["one"], s["two"]] for s in casted_ds.take()]
|
|
expected = list(zip(list(range(outer_dim)), arr))
|
|
for v, e in zip(sorted(values), expected):
|
|
np.testing.assert_equal(v, e)
|
|
|
|
# Manually deserialize the pickle tensor bytes and directly cast it to a
|
|
# TensorArray.
|
|
def deser_mapper_direct(batch: pd.DataFrame):
|
|
batch["two"] = TensorArray([pickle.loads(a) for a in batch["two"]])
|
|
return batch
|
|
|
|
casted_ds = ds.map_batches(deser_mapper_direct, batch_format="pandas")
|
|
|
|
values = [[s["one"], s["two"]] for s in casted_ds.take()]
|
|
expected = list(zip(list(range(outer_dim)), arr))
|
|
for v, e in zip(sorted(values), expected):
|
|
np.testing.assert_equal(v, e)
|
|
|
|
|
|
def test_tensors_in_tables_parquet_bytes_manual_serde(
|
|
ray_start_regular_shared, tmp_path, tensor_format_context
|
|
):
|
|
outer_dim = 3
|
|
inner_shape = (2, 2, 2)
|
|
shape = (outer_dim,) + inner_shape
|
|
num_items = np.prod(np.array(shape))
|
|
arr = np.arange(num_items).reshape(shape)
|
|
df = pd.DataFrame(
|
|
{"one": list(range(outer_dim)), "two": [a.tobytes() for a in arr]}
|
|
)
|
|
ds = ray.data.from_pandas([df])
|
|
ds.write_parquet(str(tmp_path))
|
|
ds = ray.data.read_parquet(str(tmp_path))
|
|
|
|
tensor_col_name = "two"
|
|
|
|
# Manually deserialize the tensor bytes and cast to a TensorArray.
|
|
def np_deser_mapper(batch: pa.Table):
|
|
# NOTE(Clark): We use NumPy to consolidate these potentially
|
|
# non-contiguous buffers, and to do buffer bookkeeping in general.
|
|
np_col = np.array(
|
|
[
|
|
np.ndarray(inner_shape, buffer=buf.as_buffer(), dtype=arr.dtype)
|
|
for buf in batch.column(tensor_col_name)
|
|
]
|
|
)
|
|
|
|
return batch.set_column(
|
|
batch._ensure_integer_index(tensor_col_name),
|
|
tensor_col_name,
|
|
ArrowTensorArray.from_numpy(np_col),
|
|
)
|
|
|
|
ds = ds.map_batches(np_deser_mapper, batch_format="pyarrow")
|
|
|
|
values = [[s["one"], s["two"]] for s in ds.take()]
|
|
expected = list(zip(list(range(outer_dim)), arr))
|
|
for v, e in zip(sorted(values), expected):
|
|
np.testing.assert_equal(v, e)
|
|
|
|
|
|
def test_tensors_in_tables_parquet_bytes_manual_serde_udf(
|
|
ray_start_regular_shared, tmp_path, tensor_format_context
|
|
):
|
|
tensor_format = tensor_format_context
|
|
|
|
outer_dim = 3
|
|
inner_shape = (2, 2, 2)
|
|
shape = (outer_dim,) + inner_shape
|
|
num_items = np.prod(np.array(shape))
|
|
arr = np.arange(num_items).reshape(shape)
|
|
tensor_col_name = "two"
|
|
df = pd.DataFrame(
|
|
{"one": list(range(outer_dim)), tensor_col_name: [a.tobytes() for a in arr]}
|
|
)
|
|
ds = ray.data.from_pandas([df])
|
|
ds.write_parquet(str(tmp_path))
|
|
|
|
ds = ray.data.read_parquet(
|
|
str(tmp_path),
|
|
tensor_column_schema={tensor_col_name: (np.dtype(np.int64), inner_shape)},
|
|
)
|
|
|
|
# With tensor_format_context, ARROW_NATIVE only runs when supported,
|
|
# so to_type() is safe to use without fallback
|
|
assert isinstance(
|
|
ds.schema().base_schema.field_by_name(tensor_col_name).type,
|
|
tensor_format.to_type(),
|
|
)
|
|
|
|
values = [[s["one"], s["two"]] for s in ds.take()]
|
|
expected = list(zip(list(range(outer_dim)), arr))
|
|
for v, e in zip(sorted(values), expected):
|
|
np.testing.assert_equal(v, e)
|
|
|
|
|
|
def test_tensors_in_tables_parquet_bytes_manual_serde_col_schema(
|
|
ray_start_regular_shared, tmp_path, tensor_format_context
|
|
):
|
|
tensor_format = tensor_format_context
|
|
|
|
outer_dim = 3
|
|
inner_shape = (2, 2, 2)
|
|
shape = (outer_dim,) + inner_shape
|
|
num_items = np.prod(np.array(shape))
|
|
arr = np.arange(num_items).reshape(shape)
|
|
tensor_col_name = "two"
|
|
df = pd.DataFrame(
|
|
{"one": list(range(outer_dim)), tensor_col_name: [a.tobytes() for a in arr]}
|
|
)
|
|
ds = ray.data.from_pandas([df])
|
|
ds.write_parquet(str(tmp_path))
|
|
|
|
def _block_udf(block: pa.Table):
|
|
df = block.to_pandas()
|
|
df[tensor_col_name] += 1
|
|
return pa.Table.from_pandas(df, schema=block.schema)
|
|
|
|
ds = ray.data.read_parquet(
|
|
str(tmp_path),
|
|
tensor_column_schema={tensor_col_name: (arr.dtype, inner_shape)},
|
|
_block_udf=_block_udf,
|
|
)
|
|
|
|
# With tensor_format_context, ARROW_NATIVE only runs when supported,
|
|
# so to_type() is safe to use without fallback
|
|
assert isinstance(
|
|
ds.schema().base_schema.field_by_name(tensor_col_name).type,
|
|
tensor_format.to_type(),
|
|
)
|
|
|
|
values = [[s["one"], s["two"]] for s in ds.take()]
|
|
expected = list(zip(list(range(outer_dim)), arr + 1))
|
|
for v, e in zip(sorted(values), expected):
|
|
np.testing.assert_equal(v, e)
|
|
|
|
|
|
@pytest.mark.skip(
|
|
reason=(
|
|
"Waiting for Arrow to support registering custom ExtensionType "
|
|
"casting kernels. See "
|
|
"https://issues.apache.org/jira/browse/ARROW-5890#"
|
|
)
|
|
)
|
|
def test_tensors_in_tables_parquet_bytes_with_schema(
|
|
ray_start_regular_shared, tmp_path, tensor_format_context
|
|
):
|
|
|
|
outer_dim = 3
|
|
inner_shape = (2, 2, 2)
|
|
shape = (outer_dim,) + inner_shape
|
|
num_items = np.prod(np.array(shape))
|
|
arr = np.arange(num_items).reshape(shape)
|
|
df = pd.DataFrame(
|
|
{"one": list(range(outer_dim)), "two": [a.tobytes() for a in arr]}
|
|
)
|
|
ds = ray.data.from_pandas([df])
|
|
ds.write_parquet(str(tmp_path))
|
|
tensor_type = create_arrow_fixed_shape_tensor_type(
|
|
shape=inner_shape,
|
|
dtype=pa.from_numpy_dtype(arr.dtype),
|
|
)
|
|
schema = pa.schema([("one", pa.int32()), ("two", tensor_type)])
|
|
ds = ray.data.read_parquet(str(tmp_path), schema=schema)
|
|
values = [[s["one"], s["two"]] for s in ds.take()]
|
|
expected = list(zip(list(range(outer_dim)), arr))
|
|
for v, e in zip(sorted(values), expected):
|
|
np.testing.assert_equal(v, e)
|
|
|
|
|
|
def test_tensors_in_tables_iter_batches(
|
|
ray_start_regular_shared,
|
|
enable_automatic_tensor_extension_cast,
|
|
tensor_format_context,
|
|
):
|
|
outer_dim = 3
|
|
inner_shape = (2, 2, 2)
|
|
shape = (outer_dim,) + inner_shape
|
|
num_items = np.prod(np.array(shape))
|
|
arr = np.arange(num_items).reshape(shape)
|
|
df1 = pd.DataFrame(
|
|
{"one": TensorArray(arr), "two": TensorArray(arr + 1), "label": [1.0, 2.0, 3.0]}
|
|
)
|
|
arr2 = np.arange(num_items, 2 * num_items).reshape(shape)
|
|
df2 = pd.DataFrame(
|
|
{
|
|
"one": TensorArray(arr2),
|
|
"two": TensorArray(arr2 + 1),
|
|
"label": [4.0, 5.0, 6.0],
|
|
}
|
|
)
|
|
df = pd.concat([df1, df2], ignore_index=True)
|
|
if enable_automatic_tensor_extension_cast:
|
|
df["one"] = list(df["one"].to_numpy())
|
|
df["two"] = list(df["two"].to_numpy())
|
|
ds = ray.data.from_pandas([df1, df2])
|
|
batches = list(ds.iter_batches(batch_size=2, batch_format="pandas"))
|
|
assert len(batches) == 3
|
|
expected_batches = [df.iloc[:2], df.iloc[2:4], df.iloc[4:]]
|
|
for batch, expected_batch in zip(batches, expected_batches):
|
|
batch = batch.reset_index(drop=True)
|
|
expected_batch = expected_batch.reset_index(drop=True)
|
|
pd.testing.assert_frame_equal(batch, expected_batch)
|
|
|
|
|
|
def test_ragged_tensors(ray_start_regular_shared, tensor_format_context):
|
|
"""Test Arrow type promotion between ArrowTensorType and
|
|
ArrowVariableShapedTensorType when a column contains ragged tensors."""
|
|
import numpy as np
|
|
|
|
ds = ray.data.from_items(
|
|
[
|
|
{"spam": np.zeros((32, 32, 5))},
|
|
{"spam": np.zeros((64, 64, 5))},
|
|
]
|
|
)
|
|
new_type = ds.schema().types[0].value_type
|
|
assert ds.schema().types == [
|
|
ArrowVariableShapedTensorType(dtype=new_type, ndim=3),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"write_format",
|
|
[FixedShapeTensorFormat.V1, FixedShapeTensorFormat.V2],
|
|
)
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < MIN_PYARROW_VERSION_FIXED_SHAPE_TENSOR_SCALAR,
|
|
reason="Requires pyarrow>=16 for native FixedShapeTensorType, FixedShapeTensorScalar, FixedShapeTensorArray",
|
|
)
|
|
def test_tensor_format_conversion_v1_v2_to_native(
|
|
ray_start_regular_shared, tmp_path, restore_data_context, write_format
|
|
):
|
|
"""Test that data written in V1/V2 format can be read with native format
|
|
and written back while preserving types and data.
|
|
|
|
Steps:
|
|
1. Write tensor data using V1 or V2 format
|
|
2. Set context to use arrow_native format
|
|
3. Read the data back (should convert to native format)
|
|
4. Write to a different path
|
|
5. Read again and verify types/shapes/data are preserved in native format
|
|
"""
|
|
ctx = DataContext.get_current()
|
|
|
|
# Step 1: Write data using V1 or V2 format
|
|
ctx.arrow_fixed_shape_tensor_format = write_format
|
|
|
|
outer_dim = 4
|
|
inner_shape = (3, 2)
|
|
shape = (outer_dim,) + inner_shape
|
|
num_items = np.prod(np.array(shape))
|
|
arr = np.arange(num_items, dtype=np.float64).reshape(shape)
|
|
|
|
# Create dataset with tensor column
|
|
tensor_col_name = "tensor"
|
|
df = pd.DataFrame(
|
|
{
|
|
"id": list(range(outer_dim)),
|
|
tensor_col_name: list(arr),
|
|
}
|
|
)
|
|
ds = ray.data.from_pandas([df])
|
|
|
|
# Verify initial format matches write_format
|
|
schema = ds.schema()
|
|
col_index = schema.names.index(tensor_col_name)
|
|
initial_type = schema.types[col_index]
|
|
assert isinstance(initial_type, write_format.to_type())
|
|
|
|
# Write to first path
|
|
path1 = tmp_path / "v1_v2_data"
|
|
ds.write_parquet(str(path1))
|
|
|
|
# Step 2: Switch context to native format
|
|
ctx.arrow_fixed_shape_tensor_format = FixedShapeTensorFormat.ARROW_NATIVE
|
|
|
|
# Step 3: Read data back - note that reading preserves original format
|
|
ds_read = ray.data.read_parquet(str(path1))
|
|
|
|
# Reading parquet preserves the original extension type (V1 or V2)
|
|
# The context setting only affects NEW tensor arrays
|
|
schema_read = ds_read.schema()
|
|
col_index = schema_read.names.index(tensor_col_name)
|
|
read_type = schema_read.types[col_index]
|
|
if write_format == FixedShapeTensorFormat.V1:
|
|
assert isinstance(
|
|
read_type, ArrowTensorType
|
|
), f"Expected ArrowTensorType when reading V1 data, got {type(read_type).__name__}"
|
|
else:
|
|
assert isinstance(
|
|
read_type, ArrowTensorTypeV2
|
|
), f"Expected ArrowTensorTypeV2 when reading V2 data, got {type(read_type).__name__}"
|
|
|
|
# Step 4: Apply a transformation that recreates tensor arrays in native format
|
|
# map_batches with identity function will convert to native format
|
|
def convert_to_native(batch):
|
|
# This forces recreation of the tensor arrays using the current context format
|
|
return batch
|
|
|
|
ds_converted = ds_read.map_batches(convert_to_native, batch_format="numpy")
|
|
|
|
# Now verify the format is native
|
|
schema_converted = ds_converted.schema()
|
|
col_index = schema_converted.names.index(tensor_col_name)
|
|
converted_type = schema_converted.types[col_index]
|
|
assert isinstance(converted_type, FixedShapeTensorType), (
|
|
f"Expected FixedShapeTensorType after conversion, "
|
|
f"got {type(converted_type).__name__}"
|
|
)
|
|
|
|
# Verify shape is preserved
|
|
assert (
|
|
tuple(converted_type.shape) == inner_shape
|
|
), f"Shape mismatch: expected {inner_shape}, got {converted_type.shape}"
|
|
|
|
# Step 5: Write to different path
|
|
path2 = tmp_path / "native_data"
|
|
ds_converted.write_parquet(str(path2))
|
|
|
|
# Step 6: Read again and verify native format is preserved
|
|
ds_final = ray.data.read_parquet(str(path2))
|
|
|
|
schema_final = ds_final.schema()
|
|
col_index = schema_final.names.index(tensor_col_name)
|
|
final_type = schema_final.types[col_index]
|
|
assert isinstance(final_type, FixedShapeTensorType), (
|
|
f"Expected FixedShapeTensorType after round-trip, "
|
|
f"got {type(final_type).__name__}"
|
|
)
|
|
|
|
# Verify shape is still correct
|
|
assert tuple(final_type.shape) == inner_shape, (
|
|
f"Shape mismatch after round-trip: expected {inner_shape}, "
|
|
f"got {final_type.shape}"
|
|
)
|
|
|
|
# Verify data is preserved
|
|
final_values = ds_final.take()
|
|
for i, row in enumerate(sorted(final_values, key=lambda x: x["id"])):
|
|
assert row["id"] == i
|
|
np.testing.assert_array_equal(
|
|
row[tensor_col_name],
|
|
arr[i],
|
|
err_msg=f"Data mismatch at row {i}",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|