wehub-resource-sync
64 lines
1.7 KiB
Python
64 lines
1.7 KiB
Python
"""Integration tests for array namespace expressions.
|
|
|
|
These tests require Ray and test end-to-end array namespace expression evaluation.
|
|
"""
|
|
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pytest
|
|
from packaging import version
|
|
|
|
import ray
|
|
from ray.data._internal.util import rows_same
|
|
from ray.data.expressions import col
|
|
from ray.data.tests.conftest import * # noqa
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
pytestmark = pytest.mark.skipif(
|
|
version.parse(pa.__version__) < version.parse("19.0.0"),
|
|
reason="Namespace expressions tests require PyArrow >= 19.0",
|
|
)
|
|
|
|
|
|
def _make_fixed_size_list_table() -> pa.Table:
|
|
values = pa.array([1, 2, 3, 4, 5, 6], type=pa.int64())
|
|
fixed = pa.FixedSizeListArray.from_arrays(values, list_size=2)
|
|
return pa.Table.from_arrays([fixed], names=["features"])
|
|
|
|
|
|
def test_arr_to_list_fixed_size(ray_start_regular_shared):
|
|
table = _make_fixed_size_list_table()
|
|
ds = ray.data.from_arrow(table)
|
|
|
|
result = (
|
|
ds.with_column("features", col("features").arr.to_list())
|
|
.select_columns(["features"])
|
|
.to_pandas()
|
|
)
|
|
expected = pd.DataFrame(
|
|
[
|
|
{"features": [1, 2]},
|
|
{"features": [3, 4]},
|
|
{"features": [5, 6]},
|
|
]
|
|
)
|
|
|
|
assert rows_same(result, expected)
|
|
|
|
|
|
def test_arr_to_list_invalid_dtype_raises(ray_start_regular_shared):
|
|
ds = ray.data.from_items([{"value": 1}, {"value": 2}])
|
|
|
|
with pytest.raises(
|
|
(ray.exceptions.RayTaskError, ray.exceptions.UserCodeException)
|
|
) as exc_info:
|
|
ds.with_column("value_list", col("value").arr.to_list()).to_pandas()
|
|
|
|
assert "to_list() can only be called on list-like columns" in str(exc_info.value)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|