wehub-resource-sync
362 lines
13 KiB
Python
362 lines
13 KiB
Python
import warnings
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
from ray.data._internal.tensor_extensions.arrow import ArrowTensorArray
|
|
from ray.data._internal.tensor_extensions.pandas import TensorArray, TensorDtype
|
|
from ray.data.constants import TENSOR_COLUMN_NAME
|
|
from ray.data.util.data_batch_conversion import (
|
|
BatchFormat,
|
|
_cast_ndarray_columns_to_tensor_extension,
|
|
_cast_tensor_columns_to_ndarrays,
|
|
_convert_batch_type_to_numpy,
|
|
_convert_batch_type_to_pandas,
|
|
_convert_pandas_to_batch_type,
|
|
)
|
|
from ray.data.util.torch_utils import convert_ndarray_to_torch_tensor
|
|
|
|
|
|
def test_pandas_pandas():
|
|
input_data = pd.DataFrame({"x": [1, 2, 3]})
|
|
expected_output = input_data
|
|
actual_output = _convert_batch_type_to_pandas(input_data)
|
|
pd.testing.assert_frame_equal(expected_output, actual_output)
|
|
|
|
actual_output = _convert_pandas_to_batch_type(
|
|
actual_output, type=BatchFormat.PANDAS
|
|
)
|
|
pd.testing.assert_frame_equal(actual_output, input_data)
|
|
|
|
|
|
def test_numpy_to_numpy():
|
|
input_data = {"x": np.arange(12).reshape(3, 4)}
|
|
expected_output = input_data
|
|
actual_output = _convert_batch_type_to_numpy(input_data)
|
|
assert expected_output == actual_output
|
|
|
|
input_data = {
|
|
"column_1": np.arange(12).reshape(3, 4),
|
|
"column_2": np.arange(12).reshape(3, 4),
|
|
}
|
|
expected_output = {
|
|
"column_1": np.arange(12).reshape(3, 4),
|
|
"column_2": np.arange(12).reshape(3, 4),
|
|
}
|
|
actual_output = _convert_batch_type_to_numpy(input_data)
|
|
assert input_data.keys() == expected_output.keys()
|
|
np.testing.assert_array_equal(input_data["column_1"], expected_output["column_1"])
|
|
np.testing.assert_array_equal(input_data["column_2"], expected_output["column_2"])
|
|
|
|
input_data = np.arange(12).reshape(3, 4)
|
|
expected_output = input_data
|
|
actual_output = _convert_batch_type_to_numpy(input_data)
|
|
np.testing.assert_array_equal(expected_output, actual_output)
|
|
|
|
|
|
def test_arrow_to_numpy():
|
|
input_data = pa.table({"column_1": [1, 2, 3, 4]})
|
|
expected_output = {"column_1": np.array([1, 2, 3, 4])}
|
|
actual_output = _convert_batch_type_to_numpy(input_data)
|
|
assert expected_output.keys() == actual_output.keys()
|
|
np.testing.assert_array_equal(
|
|
expected_output["column_1"], actual_output["column_1"]
|
|
)
|
|
|
|
input_data = pa.table(
|
|
{
|
|
TENSOR_COLUMN_NAME: ArrowTensorArray.from_numpy(
|
|
np.arange(12).reshape(3, 2, 2)
|
|
)
|
|
}
|
|
)
|
|
expected_output = np.arange(12).reshape(3, 2, 2)
|
|
actual_output = _convert_batch_type_to_numpy(input_data)
|
|
np.testing.assert_array_equal(expected_output, actual_output)
|
|
|
|
input_data = pa.table(
|
|
{
|
|
"column_1": [1, 2, 3, 4],
|
|
"column_2": [1, -1, 1, -1],
|
|
}
|
|
)
|
|
expected_output = {
|
|
"column_1": np.array([1, 2, 3, 4]),
|
|
"column_2": np.array([1, -1, 1, -1]),
|
|
}
|
|
|
|
actual_output = _convert_batch_type_to_numpy(input_data)
|
|
assert expected_output.keys() == actual_output.keys()
|
|
np.testing.assert_array_equal(
|
|
expected_output["column_1"], actual_output["column_1"]
|
|
)
|
|
np.testing.assert_array_equal(
|
|
expected_output["column_2"], actual_output["column_2"]
|
|
)
|
|
|
|
|
|
def test_pd_dataframe_to_numpy():
|
|
input_data = pd.DataFrame({"column_1": [1, 2, 3, 4]})
|
|
expected_output = np.array([1, 2, 3, 4])
|
|
actual_output = _convert_batch_type_to_numpy(input_data)
|
|
np.testing.assert_array_equal(expected_output, actual_output)
|
|
|
|
input_data = pd.DataFrame(
|
|
{TENSOR_COLUMN_NAME: TensorArray(np.arange(12).reshape(3, 4))}
|
|
)
|
|
expected_output = np.arange(12).reshape(3, 4)
|
|
actual_output = _convert_batch_type_to_numpy(input_data)
|
|
np.testing.assert_array_equal(expected_output, actual_output)
|
|
|
|
input_data = pd.DataFrame({"column_1": [1, 2, 3, 4], "column_2": [1, -1, 1, -1]})
|
|
expected_output = {
|
|
"column_1": np.array([1, 2, 3, 4]),
|
|
"column_2": np.array([1, -1, 1, -1]),
|
|
}
|
|
actual_output = _convert_batch_type_to_numpy(input_data)
|
|
assert expected_output.keys() == actual_output.keys()
|
|
np.testing.assert_array_equal(
|
|
expected_output["column_1"], actual_output["column_1"]
|
|
)
|
|
np.testing.assert_array_equal(
|
|
expected_output["column_2"], actual_output["column_2"]
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("use_tensor_extension_for_input", [True, False])
|
|
@pytest.mark.parametrize("cast_tensor_columns", [True, False])
|
|
def test_pandas_multi_dim_pandas(cast_tensor_columns, use_tensor_extension_for_input):
|
|
input_tensor = np.arange(12).reshape((3, 2, 2))
|
|
input_data = pd.DataFrame(
|
|
{
|
|
"x": TensorArray(input_tensor)
|
|
if use_tensor_extension_for_input
|
|
else list(input_tensor)
|
|
}
|
|
)
|
|
expected_output = pd.DataFrame(
|
|
{
|
|
"x": (
|
|
list(input_tensor)
|
|
if cast_tensor_columns or not use_tensor_extension_for_input
|
|
else TensorArray(input_tensor)
|
|
)
|
|
}
|
|
)
|
|
actual_output = _convert_batch_type_to_pandas(input_data, cast_tensor_columns)
|
|
pd.testing.assert_frame_equal(expected_output, actual_output)
|
|
|
|
actual_output = _convert_pandas_to_batch_type(
|
|
actual_output, type=BatchFormat.PANDAS, cast_tensor_columns=cast_tensor_columns
|
|
)
|
|
pd.testing.assert_frame_equal(actual_output, input_data)
|
|
|
|
|
|
def test_no_pandas_future_warning():
|
|
"""Tests that Pandas in-place FutureWarning is
|
|
suppressed during tensor extension casting."""
|
|
|
|
input_tensor = np.arange(12).reshape((3, 2, 2))
|
|
input_data = pd.DataFrame({"x": TensorArray(input_tensor)})
|
|
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("error", category=FutureWarning)
|
|
data_no_tensor_array = _cast_tensor_columns_to_ndarrays(input_data)
|
|
_cast_ndarray_columns_to_tensor_extension(data_no_tensor_array)
|
|
|
|
|
|
@pytest.mark.parametrize("cast_tensor_columns", [True, False])
|
|
def test_numpy_pandas(cast_tensor_columns):
|
|
input_data = np.array([1, 2, 3])
|
|
expected_output = pd.DataFrame({TENSOR_COLUMN_NAME: input_data})
|
|
actual_output = _convert_batch_type_to_pandas(input_data, cast_tensor_columns)
|
|
pd.testing.assert_frame_equal(expected_output, actual_output)
|
|
|
|
output_array = _convert_pandas_to_batch_type(
|
|
actual_output, type=BatchFormat.NUMPY, cast_tensor_columns=cast_tensor_columns
|
|
)
|
|
|
|
np.testing.assert_equal(output_array, input_data)
|
|
|
|
|
|
@pytest.mark.parametrize("cast_tensor_columns", [True, False])
|
|
def test_numpy_multi_dim_pandas(cast_tensor_columns):
|
|
input_data = np.arange(12).reshape((3, 2, 2))
|
|
expected_output = pd.DataFrame({TENSOR_COLUMN_NAME: list(input_data)})
|
|
actual_output = _convert_batch_type_to_pandas(input_data, cast_tensor_columns)
|
|
pd.testing.assert_frame_equal(expected_output, actual_output)
|
|
|
|
output_array = _convert_pandas_to_batch_type(
|
|
actual_output, type=BatchFormat.NUMPY, cast_tensor_columns=cast_tensor_columns
|
|
)
|
|
np.testing.assert_array_equal(np.array(list(output_array)), input_data)
|
|
|
|
|
|
def test_numpy_object_pandas():
|
|
input_data = np.array([[1, 2, 3], [1]], dtype=object)
|
|
expected_output = pd.DataFrame({TENSOR_COLUMN_NAME: input_data})
|
|
actual_output = _convert_batch_type_to_pandas(input_data)
|
|
pd.testing.assert_frame_equal(expected_output, actual_output)
|
|
|
|
np.testing.assert_array_equal(
|
|
_convert_pandas_to_batch_type(actual_output, type=BatchFormat.NUMPY), input_data
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("writable", [False, True])
|
|
def test_numpy_to_tensor_warning(writable):
|
|
input_data = np.array([[1, 2, 3]], dtype=int)
|
|
input_data.setflags(write=writable)
|
|
|
|
with pytest.warns(None) as record:
|
|
tensor = convert_ndarray_to_torch_tensor(input_data)
|
|
assert not record.list, [w.message for w in record.list]
|
|
assert tensor is not None
|
|
|
|
|
|
def test_dict_fail():
|
|
input_data = {"x": "y"}
|
|
with pytest.raises(ValueError):
|
|
_convert_batch_type_to_pandas(input_data)
|
|
|
|
|
|
@pytest.mark.parametrize("cast_tensor_columns", [True, False])
|
|
def test_dict_pandas(cast_tensor_columns):
|
|
input_data = {"x": np.array([1, 2, 3])}
|
|
expected_output = pd.DataFrame({"x": input_data["x"]})
|
|
actual_output = _convert_batch_type_to_pandas(input_data, cast_tensor_columns)
|
|
pd.testing.assert_frame_equal(expected_output, actual_output)
|
|
|
|
output_array = _convert_pandas_to_batch_type(
|
|
actual_output, type=BatchFormat.NUMPY, cast_tensor_columns=cast_tensor_columns
|
|
)
|
|
|
|
np.testing.assert_array_equal(output_array, input_data["x"])
|
|
|
|
|
|
@pytest.mark.parametrize("cast_tensor_columns", [True, False])
|
|
def test_dict_multi_dim_to_pandas(cast_tensor_columns):
|
|
tensor = np.arange(12).reshape((3, 2, 2))
|
|
input_data = {"x": tensor}
|
|
expected_output = pd.DataFrame({"x": list(tensor)})
|
|
actual_output = _convert_batch_type_to_pandas(input_data, cast_tensor_columns)
|
|
pd.testing.assert_frame_equal(expected_output, actual_output)
|
|
|
|
output_array = _convert_pandas_to_batch_type(
|
|
actual_output, type=BatchFormat.NUMPY, cast_tensor_columns=cast_tensor_columns
|
|
)
|
|
np.testing.assert_array_equal(np.array(list(output_array)), input_data["x"])
|
|
|
|
|
|
@pytest.mark.parametrize("cast_tensor_columns", [True, False])
|
|
def test_dict_pandas_multi_column(cast_tensor_columns):
|
|
array_dict = {"x": np.array([1, 2, 3]), "y": np.array([4, 5, 6])}
|
|
expected_output = pd.DataFrame(array_dict)
|
|
actual_output = _convert_batch_type_to_pandas(array_dict, cast_tensor_columns)
|
|
pd.testing.assert_frame_equal(expected_output, actual_output)
|
|
|
|
output_dict = _convert_pandas_to_batch_type(
|
|
actual_output, type=BatchFormat.NUMPY, cast_tensor_columns=cast_tensor_columns
|
|
)
|
|
|
|
for k, v in output_dict.items():
|
|
np.testing.assert_array_equal(v, array_dict[k])
|
|
|
|
|
|
def test_arrow_pandas():
|
|
df = pd.DataFrame({"x": [1, 2, 3]})
|
|
input_data = pa.Table.from_pandas(df)
|
|
expected_output = df
|
|
actual_output = _convert_batch_type_to_pandas(input_data)
|
|
pd.testing.assert_frame_equal(expected_output, actual_output)
|
|
|
|
assert _convert_pandas_to_batch_type(actual_output, type=BatchFormat.ARROW).equals(
|
|
input_data
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("cast_tensor_columns", [True, False])
|
|
def test_arrow_tensor_pandas(cast_tensor_columns):
|
|
np_array = np.arange(12).reshape((3, 2, 2))
|
|
input_data = pa.Table.from_arrays(
|
|
[ArrowTensorArray.from_numpy(np_array)], names=["x"]
|
|
)
|
|
actual_output = _convert_batch_type_to_pandas(input_data, cast_tensor_columns)
|
|
expected_output = pd.DataFrame({"x": list(np_array)})
|
|
expected_output = pd.DataFrame(
|
|
{"x": (list(np_array) if cast_tensor_columns else TensorArray(np_array))}
|
|
)
|
|
pd.testing.assert_frame_equal(expected_output, actual_output)
|
|
|
|
arrow_output = _convert_pandas_to_batch_type(
|
|
actual_output, type=BatchFormat.ARROW, cast_tensor_columns=cast_tensor_columns
|
|
)
|
|
assert arrow_output.equals(input_data)
|
|
|
|
|
|
def _make_object_column(arrays):
|
|
"""Build a 1-D object-dtype ndarray whose elements are the given ndarrays.
|
|
|
|
``np.array([...], dtype=object)`` would build a 2-D array, so we fill an
|
|
empty object array element-by-element to keep each ndarray as a cell value.
|
|
"""
|
|
out = np.empty(len(arrays), dtype=object)
|
|
for i, arr in enumerate(arrays):
|
|
out[i] = arr
|
|
return out
|
|
|
|
|
|
def test_cast_ndarray_columns_duplicate_names():
|
|
"""
|
|
Casting ndarray to tensor columns must handle duplicate column names,
|
|
keeping each column's data intact.
|
|
"""
|
|
col_a = [np.array([1, 2, 3]), np.array([4, 5, 6]), np.array([7, 8, 9])]
|
|
col_b = [np.array([10, 20]), np.array([30, 40]), np.array([50, 60])]
|
|
df = pd.DataFrame(
|
|
{"_a": _make_object_column(col_a), "_b": _make_object_column(col_b)}
|
|
)
|
|
df.columns = ["x", "x"]
|
|
|
|
actual = _cast_ndarray_columns_to_tensor_extension(df)
|
|
|
|
# Both physical columns must be cast to the tensor extension type.
|
|
assert [isinstance(dt, TensorDtype) for _, dt in actual.dtypes.items()] == [
|
|
True,
|
|
True,
|
|
]
|
|
# Values must be preserved per physical column (no write-back broadcasting
|
|
# one column's data across both duplicate labels).
|
|
np.testing.assert_array_equal(np.array(list(actual.iloc[:, 0])), col_a)
|
|
np.testing.assert_array_equal(np.array(list(actual.iloc[:, 1])), col_b)
|
|
|
|
|
|
def test_cast_tensor_columns_duplicate_names():
|
|
"""
|
|
Casting tensor columns back to ndarrays must handle duplicate names,
|
|
keeping each column's data intact.
|
|
"""
|
|
col_a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
|
|
col_b = np.array([[10, 20], [30, 40], [50, 60]])
|
|
df = pd.DataFrame({"_a": TensorArray(col_a), "_b": TensorArray(col_b)})
|
|
df.columns = ["x", "x"]
|
|
|
|
actual = _cast_tensor_columns_to_ndarrays(df)
|
|
|
|
# Neither physical column should remain a tensor extension column.
|
|
assert [isinstance(dt, TensorDtype) for _, dt in actual.dtypes.items()] == [
|
|
False,
|
|
False,
|
|
]
|
|
# Each physical column must retain its own original data.
|
|
np.testing.assert_array_equal(np.array(list(actual.iloc[:, 0])), col_a)
|
|
np.testing.assert_array_equal(np.array(list(actual.iloc[:, 1])), col_b)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-sv", __file__]))
|