wehub-resource-sync
1024 lines
33 KiB
Python
1024 lines
33 KiB
Python
import pandas as pd
|
|
import pyarrow as pa
|
|
import pyarrow.compute as pc
|
|
import pytest
|
|
from pkg_resources import parse_version
|
|
|
|
import ray
|
|
from ray.data._internal.util import rows_same
|
|
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
|
|
from ray.data.datatype import DataType
|
|
from ray.data.exceptions import UserCodeException
|
|
from ray.data.expressions import col, lit, udf
|
|
from ray.data.tests.conftest import * # noqa
|
|
from ray.exceptions import RayTaskError
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"column_name, expr, expected_value",
|
|
[
|
|
# Arithmetic operations
|
|
("result", col("id") + 1, 1), # 0 + 1 = 1
|
|
("result", col("id") + 5, 5), # 0 + 5 = 5
|
|
("result", col("id") - 1, -1), # 0 - 1 = -1
|
|
("result", col("id") * 2, 0), # 0 * 2 = 0
|
|
("result", col("id") * 3, 0), # 0 * 3 = 0
|
|
("result", col("id") / 2, 0.0), # 0 / 2 = 0.0
|
|
# More complex arithmetic
|
|
("result", (col("id") + 1) * 2, 2), # (0 + 1) * 2 = 2
|
|
("result", (col("id") * 2) + 3, 3), # 0 * 2 + 3 = 3
|
|
# Comparison operations
|
|
("result", col("id") > 0, False), # 0 > 0 = False
|
|
("result", col("id") >= 0, True), # 0 >= 0 = True
|
|
("result", col("id") < 1, True), # 0 < 1 = True
|
|
("result", col("id") <= 0, True), # 0 <= 0 = True
|
|
("result", col("id") == 0, True), # 0 == 0 = True
|
|
# Operations with literals
|
|
("result", col("id") + lit(10), 10), # 0 + 10 = 10
|
|
("result", col("id") * lit(5), 0), # 0 * 5 = 0
|
|
("result", lit(2) + col("id"), 2), # 2 + 0 = 2
|
|
("result", lit(10) / (col("id") + 1), 10.0), # 10 / (0 + 1) = 10.0
|
|
],
|
|
)
|
|
def test_with_column(
|
|
ray_start_regular_shared,
|
|
column_name,
|
|
expr,
|
|
expected_value,
|
|
target_max_block_size_infinite_or_default,
|
|
):
|
|
"""Verify that `with_column` works with various operations."""
|
|
ds = ray.data.range(5).with_column(column_name, expr)
|
|
result = ds.take(1)[0]
|
|
assert result["id"] == 0
|
|
assert result[column_name] == expected_value
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
def test_with_column_nonexistent_column(
|
|
ray_start_regular_shared, target_max_block_size_infinite_or_default
|
|
):
|
|
"""Verify that referencing a non-existent column with col() raises an exception."""
|
|
# Create a dataset with known column "id"
|
|
ds = ray.data.range(5)
|
|
|
|
# Try to reference a non-existent column - this should raise an exception
|
|
with pytest.raises(UserCodeException):
|
|
ds.with_column("result", col("nonexistent_column") + 1).materialize()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
def test_with_column_multiple_expressions(
|
|
ray_start_regular_shared, target_max_block_size_infinite_or_default
|
|
):
|
|
"""Verify that `with_column` correctly handles multiple expressions at once."""
|
|
ds = ray.data.range(5)
|
|
|
|
ds = ds.with_column("plus_one", col("id") + 1)
|
|
ds = ds.with_column("times_two", col("id") * 2)
|
|
ds = ds.with_column("ten_minus_id", 10 - col("id"))
|
|
|
|
first_row = ds.take(1)[0]
|
|
assert first_row["id"] == 0
|
|
assert first_row["plus_one"] == 1
|
|
assert first_row["times_two"] == 0
|
|
assert first_row["ten_minus_id"] == 10
|
|
|
|
# Ensure all new columns exist in the schema.
|
|
assert set(ds.schema().names) == {"id", "plus_one", "times_two", "ten_minus_id"}
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"udf_function, column_name, expected_result",
|
|
[
|
|
# Single column UDF - add one to each value
|
|
pytest.param(
|
|
lambda: udf(DataType.int64())(lambda x: pc.add(x, 1)),
|
|
"add_one",
|
|
1, # 0 + 1 = 1
|
|
id="single_column_add_one",
|
|
),
|
|
# Single column UDF - multiply by 2
|
|
pytest.param(
|
|
lambda: udf(DataType.int64())(lambda x: pc.multiply(x, 2)),
|
|
"times_two",
|
|
0, # 0 * 2 = 0
|
|
id="single_column_multiply",
|
|
),
|
|
# Single column UDF - square the value
|
|
pytest.param(
|
|
lambda: udf(DataType.int64())(lambda x: pc.multiply(x, x)),
|
|
"squared",
|
|
0, # 0 * 0 = 0
|
|
id="single_column_square",
|
|
),
|
|
# Single column UDF with string return type
|
|
pytest.param(
|
|
lambda: udf(DataType.string())(lambda x: pc.cast(x, pa.string())),
|
|
"id_str",
|
|
"0", # Convert 0 to "0"
|
|
id="single_column_to_string",
|
|
),
|
|
# Single column UDF with float return type
|
|
pytest.param(
|
|
lambda: udf(DataType.float64())(lambda x: pc.divide(x, 2.0)),
|
|
"half",
|
|
0.0, # 0 / 2.0 = 0.0
|
|
id="single_column_divide_float",
|
|
),
|
|
],
|
|
)
|
|
def test_with_column_udf_single_column(
|
|
ray_start_regular_shared,
|
|
udf_function,
|
|
column_name,
|
|
expected_result,
|
|
target_max_block_size_infinite_or_default,
|
|
):
|
|
"""Test UDFExpr functionality with single column operations in with_column."""
|
|
ds = ray.data.range(5)
|
|
udf_fn = udf_function()
|
|
|
|
# Apply the UDF to the "id" column
|
|
ds_with_udf = ds.with_column(column_name, udf_fn(col("id")))
|
|
|
|
result = ds_with_udf.take(1)[0]
|
|
assert result["id"] == 0
|
|
assert result[column_name] == expected_result
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"test_scenario",
|
|
[
|
|
# Multi-column UDF - add two columns
|
|
pytest.param(
|
|
{
|
|
"data": [{"a": 1, "b": 2}, {"a": 3, "b": 4}],
|
|
"udf": lambda: udf(DataType.int64())(lambda x, y: pc.add(x, y)),
|
|
"column_name": "sum_ab",
|
|
"expected_first": 3, # 1 + 2 = 3
|
|
"expected_second": 7, # 3 + 4 = 7
|
|
},
|
|
id="multi_column_add",
|
|
),
|
|
# Multi-column UDF - multiply two columns
|
|
pytest.param(
|
|
{
|
|
"data": [{"x": 2, "y": 3}, {"x": 4, "y": 5}],
|
|
"udf": lambda: udf(DataType.int64())(lambda x, y: pc.multiply(x, y)),
|
|
"column_name": "product_xy",
|
|
"expected_first": 6, # 2 * 3 = 6
|
|
"expected_second": 20, # 4 * 5 = 20
|
|
},
|
|
id="multi_column_multiply",
|
|
),
|
|
# Multi-column UDF - string concatenation
|
|
pytest.param(
|
|
{
|
|
"data": [
|
|
{"first": "John", "last": "Doe"},
|
|
{"first": "Jane", "last": "Smith"},
|
|
],
|
|
"udf": lambda: udf(DataType.string())(
|
|
lambda first, last: pc.binary_join_element_wise(first, last, " ")
|
|
),
|
|
"column_name": "full_name",
|
|
"expected_first": "John Doe",
|
|
"expected_second": "Jane Smith",
|
|
},
|
|
id="multi_column_string_concat",
|
|
),
|
|
],
|
|
)
|
|
def test_with_column_udf_multi_column(
|
|
ray_start_regular_shared,
|
|
test_scenario,
|
|
target_max_block_size_infinite_or_default,
|
|
):
|
|
"""Test UDFExpr functionality with multi-column operations in with_column."""
|
|
data = test_scenario["data"]
|
|
udf_fn = test_scenario["udf"]()
|
|
column_name = test_scenario["column_name"]
|
|
expected_first = test_scenario["expected_first"]
|
|
expected_second = test_scenario["expected_second"]
|
|
|
|
ds = ray.data.from_items(data)
|
|
|
|
# Apply UDF to multiple columns based on the scenario
|
|
if "a" in data[0] and "b" in data[0]:
|
|
ds_with_udf = ds.with_column(column_name, udf_fn(col("a"), col("b")))
|
|
elif "x" in data[0] and "y" in data[0]:
|
|
ds_with_udf = ds.with_column(column_name, udf_fn(col("x"), col("y")))
|
|
else: # first/last name scenario
|
|
ds_with_udf = ds.with_column(column_name, udf_fn(col("first"), col("last")))
|
|
|
|
results = ds_with_udf.take(2)
|
|
assert results[0][column_name] == expected_first
|
|
assert results[1][column_name] == expected_second
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"expression_scenario",
|
|
[
|
|
# UDF in arithmetic expression
|
|
pytest.param(
|
|
{
|
|
"expression_factory": lambda add_one_udf: add_one_udf(col("id")) * 2,
|
|
"expected": 2, # (0 + 1) * 2 = 2
|
|
"column_name": "udf_times_two",
|
|
},
|
|
id="udf_in_arithmetic",
|
|
),
|
|
# UDF with literal addition
|
|
pytest.param(
|
|
{
|
|
"expression_factory": lambda add_one_udf: add_one_udf(col("id"))
|
|
+ lit(10),
|
|
"expected": 11, # (0 + 1) + 10 = 11
|
|
"column_name": "udf_plus_literal",
|
|
},
|
|
id="udf_plus_literal",
|
|
),
|
|
# UDF in comparison
|
|
pytest.param(
|
|
{
|
|
"expression_factory": lambda add_one_udf: add_one_udf(col("id")) > 0,
|
|
"expected": True, # (0 + 1) > 0 = True
|
|
"column_name": "udf_comparison",
|
|
},
|
|
id="udf_in_comparison",
|
|
),
|
|
# Nested UDF operations (UDF + regular expression)
|
|
pytest.param(
|
|
{
|
|
"expression_factory": lambda add_one_udf: add_one_udf(col("id") + 5),
|
|
"expected": 6, # add_one(0 + 5) = add_one(5) = 6
|
|
"column_name": "nested_udf",
|
|
},
|
|
id="nested_udf_expression",
|
|
),
|
|
],
|
|
)
|
|
def test_with_column_udf_in_complex_expressions(
|
|
ray_start_regular_shared,
|
|
expression_scenario,
|
|
target_max_block_size_infinite_or_default,
|
|
):
|
|
"""Test UDFExpr functionality in complex expressions with with_column."""
|
|
ds = ray.data.range(5)
|
|
|
|
# Create a simple add_one UDF for use in expressions
|
|
@udf(DataType.int64())
|
|
def add_one(x: pa.Array) -> pa.Array:
|
|
return pc.add(x, 1)
|
|
|
|
expression = expression_scenario["expression_factory"](add_one)
|
|
expected = expression_scenario["expected"]
|
|
column_name = expression_scenario["column_name"]
|
|
|
|
ds_with_expr = ds.with_column(column_name, expression)
|
|
|
|
result = ds_with_expr.take(1)[0]
|
|
assert result["id"] == 0
|
|
assert result[column_name] == expected
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
def test_with_column_udf_multiple_udfs(
|
|
ray_start_regular_shared, target_max_block_size_infinite_or_default
|
|
):
|
|
"""Test applying multiple UDFs in sequence with with_column."""
|
|
ds = ray.data.range(5)
|
|
|
|
# Define multiple UDFs
|
|
@udf(DataType.int64())
|
|
def add_one(x: pa.Array) -> pa.Array:
|
|
return pc.add(x, 1)
|
|
|
|
@udf(DataType.int64())
|
|
def multiply_by_two(x: pa.Array) -> pa.Array:
|
|
return pc.multiply(x, 2)
|
|
|
|
@udf(DataType.float64())
|
|
def divide_by_three(x: pa.Array) -> pa.Array:
|
|
return pc.divide(x, 3.0)
|
|
|
|
# Apply UDFs in sequence
|
|
ds = ds.with_column("plus_one", add_one(col("id")))
|
|
ds = ds.with_column("times_two", multiply_by_two(col("plus_one")))
|
|
ds = ds.with_column("div_three", divide_by_three(col("times_two")))
|
|
|
|
# Convert to pandas and compare with expected result
|
|
result_df = ds.to_pandas()
|
|
|
|
expected_df = pd.DataFrame(
|
|
{
|
|
"id": [0, 1, 2, 3, 4],
|
|
"plus_one": [1, 2, 3, 4, 5], # id + 1
|
|
"times_two": [2, 4, 6, 8, 10], # (id + 1) * 2
|
|
"div_three": [
|
|
2.0 / 3.0,
|
|
4.0 / 3.0,
|
|
2.0,
|
|
8.0 / 3.0,
|
|
10.0 / 3.0,
|
|
], # ((id + 1) * 2) / 3
|
|
}
|
|
)
|
|
expected_df = expected_df.astype(result_df.dtypes.to_dict())
|
|
|
|
pd.testing.assert_frame_equal(result_df, expected_df)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
def test_with_column_mixed_udf_and_regular_expressions(
|
|
ray_start_regular_shared, target_max_block_size_infinite_or_default
|
|
):
|
|
"""Test mixing UDF expressions and regular expressions in with_column operations."""
|
|
ds = ray.data.range(5)
|
|
|
|
# Define a UDF for testing
|
|
@udf(DataType.int64())
|
|
def multiply_by_three(x: pa.Array) -> pa.Array:
|
|
return pc.multiply(x, 3)
|
|
|
|
# Mix regular expressions and UDF expressions
|
|
ds = ds.with_column("plus_ten", col("id") + 10) # Regular expression
|
|
ds = ds.with_column("times_three", multiply_by_three(col("id"))) # UDF expression
|
|
ds = ds.with_column("minus_five", col("id") - 5) # Regular expression
|
|
ds = ds.with_column(
|
|
"udf_plus_regular", multiply_by_three(col("id")) + col("plus_ten")
|
|
) # Mixed: UDF + regular
|
|
ds = ds.with_column(
|
|
"comparison", col("times_three") > col("plus_ten")
|
|
) # Regular expression using UDF result
|
|
|
|
# Convert to pandas and compare with expected result
|
|
result_df = ds.to_pandas()
|
|
|
|
expected_df = pd.DataFrame(
|
|
{
|
|
"id": [0, 1, 2, 3, 4],
|
|
"plus_ten": [10, 11, 12, 13, 14], # id + 10
|
|
"times_three": [0, 3, 6, 9, 12], # id * 3
|
|
"minus_five": [-5, -4, -3, -2, -1], # id - 5
|
|
"udf_plus_regular": [10, 14, 18, 22, 26], # (id * 3) + (id + 10)
|
|
"comparison": [False, False, False, False, False], # times_three > plus_ten
|
|
}
|
|
)
|
|
expected_df = expected_df.astype(result_df.dtypes.to_dict())
|
|
|
|
pd.testing.assert_frame_equal(result_df, expected_df)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
def test_with_column_udf_invalid_return_type_validation(
|
|
ray_start_regular_shared, target_max_block_size_infinite_or_default
|
|
):
|
|
"""Test that UDFs returning invalid types raise TypeError with clear message."""
|
|
ds = ray.data.range(3)
|
|
|
|
# Test UDF returning invalid type (dict) - expecting string but returning dict
|
|
@udf(DataType.string())
|
|
def invalid_dict_return(x: pa.Array) -> dict:
|
|
return {"invalid": "return_type"}
|
|
|
|
# Test UDF returning invalid type (str) - expecting string but returning plain str
|
|
@udf(DataType.string())
|
|
def invalid_str_return(x: pa.Array) -> str:
|
|
return "invalid_string"
|
|
|
|
# Test UDF returning invalid type (int) - expecting int64 but returning plain int
|
|
@udf(DataType.int64())
|
|
def invalid_int_return(x: pa.Array) -> int:
|
|
return 42
|
|
|
|
# Test each invalid return type
|
|
test_cases = [
|
|
(invalid_dict_return, "dict"),
|
|
(invalid_str_return, "str"),
|
|
(invalid_int_return, "int"),
|
|
]
|
|
|
|
for invalid_udf, expected_type_name in test_cases:
|
|
with pytest.raises((RayTaskError, UserCodeException)) as exc_info:
|
|
ds.with_column("invalid_col", invalid_udf(col("id"))).take(1)
|
|
|
|
# The actual TypeError gets wrapped, so we need to check the exception chain
|
|
error_message = str(exc_info.value)
|
|
assert f"returned invalid type {expected_type_name}" in error_message
|
|
assert "Expected type" in error_message
|
|
assert "pandas.Series" in error_message and "numpy.ndarray" in error_message
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"scenario",
|
|
[
|
|
pytest.param(
|
|
{
|
|
"data": [
|
|
{"name": "Alice"},
|
|
{"name": "Bob"},
|
|
{"name": "Charlie"},
|
|
],
|
|
"expr_factory": lambda: col("name") + "_X",
|
|
"column_name": "name_with_suffix",
|
|
"expected": ["Alice_X", "Bob_X", "Charlie_X"],
|
|
},
|
|
id="string_col_plus_python_literal_rhs",
|
|
),
|
|
pytest.param(
|
|
{
|
|
"data": [
|
|
{"name": "Alice"},
|
|
{"name": "Bob"},
|
|
{"name": "Charlie"},
|
|
],
|
|
"expr_factory": lambda: "_X" + col("name"),
|
|
"column_name": "name_with_prefix",
|
|
"expected": ["_XAlice", "_XBob", "_XCharlie"],
|
|
},
|
|
id="python_literal_lhs_plus_string_col",
|
|
),
|
|
pytest.param(
|
|
{
|
|
"data": [
|
|
{"first": "John", "last": "Doe"},
|
|
{"first": "Jane", "last": "Smith"},
|
|
],
|
|
"expr_factory": lambda: col("first") + col("last"),
|
|
"column_name": "full_name",
|
|
"expected": ["JohnDoe", "JaneSmith"],
|
|
},
|
|
id="string_col_plus_string_col",
|
|
),
|
|
pytest.param(
|
|
{
|
|
"arrow_table": pa.table(
|
|
{"name": pa.array(["Alice", "Bob"]).dictionary_encode()}
|
|
),
|
|
"expr_factory": lambda: col("name") + "_X",
|
|
"column_name": "name_with_suffix",
|
|
"expected": ["Alice_X", "Bob_X"],
|
|
},
|
|
id="dict_encoded_string_col_plus_literal_rhs",
|
|
),
|
|
pytest.param(
|
|
{
|
|
"data": [
|
|
{"name": "Alice"},
|
|
{"name": "Bob"},
|
|
],
|
|
"expr_factory": lambda: col("name") + lit("_X"),
|
|
"column_name": "name_with_suffix",
|
|
"expected": ["Alice_X", "Bob_X"],
|
|
},
|
|
id="string_col_plus_lit_literal_rhs",
|
|
),
|
|
],
|
|
)
|
|
def test_with_column_string_concat_combinations(
|
|
ray_start_regular_shared,
|
|
scenario,
|
|
):
|
|
if "arrow_table" in scenario:
|
|
ds = ray.data.from_arrow(scenario["arrow_table"])
|
|
else:
|
|
ds = ray.data.from_items(scenario["data"])
|
|
|
|
expr = scenario["expr_factory"]()
|
|
column_name = scenario["column_name"]
|
|
|
|
ds2 = ds.with_column(column_name, expr)
|
|
out = ds2.to_pandas()
|
|
assert out[column_name].tolist() == scenario["expected"]
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
def test_with_column_string_concat_type_mismatch_raises(
|
|
ray_start_regular_shared,
|
|
):
|
|
# int + string should raise a user-facing error
|
|
ds = ray.data.range(3)
|
|
with pytest.raises((RayTaskError, UserCodeException)):
|
|
ds.with_column("bad", col("id") + "_X").materialize()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"expr_factory, expected_columns, alias_name, expected_values",
|
|
[
|
|
(
|
|
lambda: col("id").alias("new_id"),
|
|
["id", "new_id"],
|
|
"new_id",
|
|
[0, 1, 2, 3, 4], # Copy of id column
|
|
),
|
|
(
|
|
lambda: (col("id") + 1).alias("id_plus_one"),
|
|
["id", "id_plus_one"],
|
|
"id_plus_one",
|
|
[1, 2, 3, 4, 5], # id + 1
|
|
),
|
|
(
|
|
lambda: (col("id") * 2 + 5).alias("transformed"),
|
|
["id", "transformed"],
|
|
"transformed",
|
|
[5, 7, 9, 11, 13], # id * 2 + 5
|
|
),
|
|
(
|
|
lambda: lit(42).alias("constant"),
|
|
["id", "constant"],
|
|
"constant",
|
|
[42, 42, 42, 42, 42], # lit(42)
|
|
),
|
|
(
|
|
lambda: (col("id") >= 0).alias("is_non_negative"),
|
|
["id", "is_non_negative"],
|
|
"is_non_negative",
|
|
[True, True, True, True, True], # id >= 0
|
|
),
|
|
(
|
|
lambda: (col("id") + 1).alias("id"),
|
|
["id"], # Only one column since we're overwriting id
|
|
"id",
|
|
[1, 2, 3, 4, 5], # id + 1 replaces original id
|
|
),
|
|
],
|
|
ids=[
|
|
"col_alias",
|
|
"arithmetic_alias",
|
|
"complex_alias",
|
|
"literal_alias",
|
|
"comparison_alias",
|
|
"overwrite_existing_column",
|
|
],
|
|
)
|
|
def test_with_column_alias_expressions(
|
|
ray_start_regular_shared,
|
|
expr_factory,
|
|
expected_columns,
|
|
alias_name,
|
|
expected_values,
|
|
):
|
|
"""Test that alias expressions work correctly with with_column."""
|
|
expr = expr_factory()
|
|
|
|
# Verify the alias name matches what we expect
|
|
assert expr.name == alias_name
|
|
|
|
# Apply the aliased expression
|
|
ds = ray.data.range(5).with_column(alias_name, expr)
|
|
|
|
# Convert to pandas for comprehensive comparison
|
|
result_df = ds.to_pandas()
|
|
|
|
# Create expected DataFrame
|
|
expected_df = pd.DataFrame({"id": [0, 1, 2, 3, 4], alias_name: expected_values})
|
|
|
|
# Ensure column order matches expected_columns
|
|
expected_df = expected_df[expected_columns]
|
|
expected_df = expected_df.astype(result_df.dtypes.to_dict())
|
|
|
|
# Assert the entire DataFrame is equal
|
|
pd.testing.assert_frame_equal(result_df, expected_df)
|
|
# Verify the alias expression evaluates the same as the non-aliased version
|
|
non_aliased_expr = expr
|
|
ds_non_aliased = ray.data.range(5).with_column(alias_name, non_aliased_expr)
|
|
|
|
non_aliased_df = ds_non_aliased.to_pandas()
|
|
|
|
pd.testing.assert_frame_equal(result_df, non_aliased_df)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
def test_with_column_callable_class_udf_actor_semantics(ray_start_regular_shared):
|
|
"""Test that callable class UDFs maintain state across batches using actor semantics."""
|
|
import pyarrow.compute as pc
|
|
|
|
# Create a callable class UDF that tracks the number of times it's called
|
|
@udf(return_dtype=DataType.int32())
|
|
class InvocationCounter:
|
|
def __init__(self, offset=0):
|
|
self.offset = offset
|
|
self.call_count = 0
|
|
|
|
def __call__(self, x):
|
|
# Increment call count each time the UDF is invoked
|
|
self.call_count += 1
|
|
# Add the offset plus the call count to show state is maintained
|
|
return pc.add(pc.add(x, self.offset), self.call_count)
|
|
|
|
# Create a dataset with multiple blocks to ensure multiple invocations
|
|
ds = ray.data.range(20, override_num_blocks=4)
|
|
|
|
# Use the callable class UDF
|
|
counter = InvocationCounter(offset=100)
|
|
result_ds = ds.with_column("result", counter(col("id")))
|
|
|
|
# Convert to list to trigger execution
|
|
results = result_ds.take_all()
|
|
|
|
# The results should show that the call_count incremented across batches
|
|
# Since we have 4 blocks, the UDF should be called 4 times on the same actor
|
|
# The exact values will depend on which batch each row came from
|
|
# But we can verify that the offset (100) was applied
|
|
for result in results:
|
|
# Each result should have the base id + offset (100) + at least 1 (first call)
|
|
assert result["result"] >= result["id"] + 100 + 1
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
def test_with_column_callable_class_udf_with_constructor_args(
|
|
ray_start_regular_shared,
|
|
):
|
|
"""Test that callable class UDFs correctly use constructor arguments."""
|
|
import pyarrow.compute as pc
|
|
|
|
@udf(return_dtype=DataType.int32())
|
|
class AddOffset:
|
|
def __init__(self, offset):
|
|
self.offset = offset
|
|
|
|
def __call__(self, x):
|
|
return pc.add(x, self.offset)
|
|
|
|
# Create dataset
|
|
ds = ray.data.range(10)
|
|
|
|
# Test with different offsets
|
|
add_five = AddOffset(5)
|
|
add_ten = AddOffset(10)
|
|
|
|
result_5 = ds.with_column("plus_five", add_five(col("id"))).to_pandas()
|
|
result_10 = ds.with_column("plus_ten", add_ten(col("id"))).to_pandas()
|
|
|
|
# Verify the offsets were applied correctly
|
|
expected_5 = pd.DataFrame({"id": list(range(10)), "plus_five": list(range(5, 15))})
|
|
expected_10 = pd.DataFrame({"id": list(range(10)), "plus_ten": list(range(10, 20))})
|
|
|
|
assert rows_same(result_5, expected_5)
|
|
assert rows_same(result_10, expected_10)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
def test_with_column_multiple_callable_class_udfs(ray_start_regular_shared):
|
|
"""Test that multiple callable class UDFs can be used in the same projection."""
|
|
import pyarrow.compute as pc
|
|
|
|
@udf(return_dtype=DataType.int32())
|
|
class Multiplier:
|
|
def __init__(self, factor):
|
|
self.factor = factor
|
|
|
|
def __call__(self, x):
|
|
return pc.multiply(x, self.factor)
|
|
|
|
@udf(return_dtype=DataType.int32())
|
|
class Adder:
|
|
def __init__(self, addend):
|
|
self.addend = addend
|
|
|
|
def __call__(self, x):
|
|
return pc.add(x, self.addend)
|
|
|
|
# Create dataset
|
|
ds = ray.data.range(5)
|
|
|
|
# Use multiple callable class UDFs
|
|
times_two = Multiplier(2)
|
|
plus_ten = Adder(10)
|
|
|
|
result = ds.with_column("doubled", times_two(col("id"))).with_column(
|
|
"plus_ten", plus_ten(col("id"))
|
|
)
|
|
|
|
result_df = result.to_pandas()
|
|
expected_df = pd.DataFrame(
|
|
{
|
|
"id": [0, 1, 2, 3, 4],
|
|
"doubled": [0, 2, 4, 6, 8],
|
|
"plus_ten": [10, 11, 12, 13, 14],
|
|
}
|
|
)
|
|
|
|
assert rows_same(result_df, expected_df)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
def test_with_column_same_callable_class_different_constructor_args(
|
|
ray_start_regular_shared,
|
|
):
|
|
"""Test that the same callable class with different constructor args works correctly.
|
|
|
|
This test ensures that when the same callable class is instantiated with different
|
|
constructor arguments, each instance maintains its own state. This is important for
|
|
future-proofing in case Actor->Actor fusion becomes enabled.
|
|
"""
|
|
import pyarrow.compute as pc
|
|
|
|
@udf(return_dtype=DataType.int32())
|
|
class Multiplier:
|
|
def __init__(self, factor):
|
|
self.factor = factor
|
|
|
|
def __call__(self, x):
|
|
return pc.multiply(x, self.factor)
|
|
|
|
# Create dataset
|
|
ds = ray.data.range(5)
|
|
|
|
# Use the SAME class with DIFFERENT constructor arguments
|
|
times_two = Multiplier(2)
|
|
times_three = Multiplier(3)
|
|
|
|
result = ds.with_column("times_two", times_two(col("id"))).with_column(
|
|
"times_three", times_three(col("id"))
|
|
)
|
|
|
|
print(result.explain())
|
|
|
|
result_df = result.to_pandas()
|
|
expected_df = pd.DataFrame(
|
|
{
|
|
"id": [0, 1, 2, 3, 4],
|
|
"times_two": [0, 2, 4, 6, 8], # id * 2
|
|
"times_three": [0, 3, 6, 9, 12], # id * 3
|
|
}
|
|
)
|
|
|
|
from ray.data._internal.util import rows_same
|
|
|
|
assert rows_same(result_df, expected_df)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
def test_with_column_callable_class_udf_with_compute_strategy(
|
|
ray_start_regular_shared,
|
|
):
|
|
"""Test that compute strategy can be specified for callable class UDFs."""
|
|
import pyarrow.compute as pc
|
|
|
|
@udf(return_dtype=DataType.int32())
|
|
class AddOffset:
|
|
def __init__(self, offset):
|
|
self.offset = offset
|
|
|
|
def __call__(self, x):
|
|
return pc.add(x, self.offset)
|
|
|
|
# Create dataset
|
|
ds = ray.data.range(10)
|
|
|
|
# Use a specific compute strategy
|
|
add_five = AddOffset(5)
|
|
result = ds.with_column(
|
|
"result",
|
|
add_five(col("id")),
|
|
compute=ray.data.ActorPoolStrategy(size=2),
|
|
)
|
|
|
|
result_df = result.to_pandas()
|
|
expected_df = pd.DataFrame({"id": list(range(10)), "result": list(range(5, 15))})
|
|
|
|
assert rows_same(result_df, expected_df)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
def test_with_column_async_callable_class_udf(ray_start_regular_shared):
|
|
"""Test that async callable class UDFs work correctly with actor semantics."""
|
|
import asyncio
|
|
|
|
import pyarrow.compute as pc
|
|
|
|
@udf(return_dtype=DataType.int32())
|
|
class AsyncAddOffset:
|
|
def __init__(self, offset):
|
|
self.offset = offset
|
|
self.call_count = 0
|
|
|
|
async def __call__(self, x):
|
|
# Simulate async work
|
|
await asyncio.sleep(0.001)
|
|
self.call_count += 1
|
|
# Add offset to show the UDF was called
|
|
return pc.add(x, self.offset)
|
|
|
|
# Create dataset
|
|
ds = ray.data.range(10, override_num_blocks=2)
|
|
|
|
# Use async callable class UDF
|
|
add_five = AsyncAddOffset(5)
|
|
result = ds.with_column("result", add_five(col("id")))
|
|
|
|
result_df = result.to_pandas()
|
|
expected_df = pd.DataFrame({"id": list(range(10)), "result": list(range(5, 15))})
|
|
|
|
assert rows_same(result_df, expected_df)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
def test_with_column_async_callable_class_udf_with_state(ray_start_regular_shared):
|
|
"""Test that async callable class UDFs maintain state across batches."""
|
|
import asyncio
|
|
|
|
import pyarrow.compute as pc
|
|
|
|
@udf(return_dtype=DataType.int32())
|
|
class AsyncCounter:
|
|
def __init__(self):
|
|
self.total_processed = 0
|
|
|
|
async def __call__(self, x):
|
|
# Simulate async work
|
|
await asyncio.sleep(0.001)
|
|
# Track how many items we've processed
|
|
batch_size = len(x)
|
|
self.total_processed += batch_size
|
|
# Return the running count
|
|
return pc.add(x, self.total_processed - batch_size)
|
|
|
|
# Create dataset with multiple blocks
|
|
ds = ray.data.range(20, override_num_blocks=4)
|
|
|
|
# Use async callable class UDF with state
|
|
counter = AsyncCounter()
|
|
result = ds.with_column("running_total", counter(col("id")))
|
|
|
|
# Just verify we got results without errors
|
|
# The exact values will depend on execution order
|
|
results = result.take_all()
|
|
assert len(results) == 20
|
|
# All values should be at least the original id
|
|
for i, result in enumerate(results):
|
|
assert result["running_total"] >= result["id"]
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
def test_with_column_multiple_async_callable_class_udfs(ray_start_regular_shared):
|
|
"""Test that multiple async callable class UDFs can work together."""
|
|
import asyncio
|
|
|
|
import pyarrow.compute as pc
|
|
|
|
@udf(return_dtype=DataType.int32())
|
|
class AsyncMultiplier:
|
|
def __init__(self, factor):
|
|
self.factor = factor
|
|
|
|
async def __call__(self, x):
|
|
await asyncio.sleep(0.001)
|
|
return pc.multiply(x, self.factor)
|
|
|
|
@udf(return_dtype=DataType.int32())
|
|
class AsyncAdder:
|
|
def __init__(self, addend):
|
|
self.addend = addend
|
|
|
|
async def __call__(self, x):
|
|
await asyncio.sleep(0.001)
|
|
return pc.add(x, self.addend)
|
|
|
|
# Create dataset
|
|
ds = ray.data.range(5)
|
|
|
|
# Use multiple async callable class UDFs
|
|
times_two = AsyncMultiplier(2)
|
|
plus_ten = AsyncAdder(10)
|
|
|
|
result = ds.with_column("doubled", times_two(col("id"))).with_column(
|
|
"plus_ten", plus_ten(col("id"))
|
|
)
|
|
|
|
result_df = result.to_pandas()
|
|
expected_df = pd.DataFrame(
|
|
{
|
|
"id": [0, 1, 2, 3, 4],
|
|
"doubled": [0, 2, 4, 6, 8],
|
|
"plus_ten": [10, 11, 12, 13, 14],
|
|
}
|
|
)
|
|
|
|
assert rows_same(result_df, expected_df)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="with_column requires PyArrow >= 20.0.0",
|
|
)
|
|
def test_with_column_async_generator_udf_multiple_yields(ray_start_regular_shared):
|
|
"""Test that async generator UDFs correctly handle multiple yields.
|
|
|
|
When an async generator UDF yields multiple values, the last (most recent)
|
|
value is returned. This matches map_batches behavior of collecting all yields,
|
|
while adapting to expression context where a single value per row is required.
|
|
"""
|
|
|
|
import pyarrow.compute as pc
|
|
|
|
@udf(return_dtype=DataType.int32())
|
|
class AsyncGeneratorMultiYield:
|
|
"""UDF that yields multiple values - last yield is returned."""
|
|
|
|
def __init__(self, offset):
|
|
self.offset = offset
|
|
|
|
async def __call__(self, x):
|
|
# Yield multiple values for the same input
|
|
# Fix: Last yield is returned (most recent/final result)
|
|
yield pc.add(x, self.offset) # First yield: x + offset
|
|
yield pc.multiply(x, self.offset + 10) # Second yield: x * (offset + 10)
|
|
yield pc.add(x, self.offset * 2) # Third yield: x + (offset * 2) - RETURNED
|
|
|
|
# Create dataset
|
|
ds = ray.data.range(5, override_num_blocks=1)
|
|
|
|
# Use async generator UDF
|
|
udf_instance = AsyncGeneratorMultiYield(5)
|
|
result = ds.with_column("result", udf_instance(col("id")))
|
|
|
|
result_df = result.to_pandas()
|
|
|
|
# Fixed behavior: last yield is returned
|
|
# Input: [0, 1, 2, 3, 4]
|
|
# First yield: [0+5, 1+5, 2+5, 3+5, 4+5] = [5, 6, 7, 8, 9]
|
|
# Second yield: [0*15, 1*15, 2*15, 3*15, 4*15] = [0, 15, 30, 45, 60]
|
|
# Third yield (RETURNED): [0+10, 1+10, 2+10, 3+10, 4+10] = [10, 11, 12, 13, 14]
|
|
|
|
expected_after_fix = pd.DataFrame(
|
|
{
|
|
"id": [0, 1, 2, 3, 4],
|
|
"result": [10, 11, 12, 13, 14], # Last yield returned: id + (5*2) = id + 10
|
|
}
|
|
)
|
|
|
|
assert rows_same(result_df, expected_after_fix)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|