wehub-resource-sync
210 lines
7.7 KiB
Python
210 lines
7.7 KiB
Python
"""Integration tests for arithmetic expression operations.
|
|
|
|
These tests require Ray and test end-to-end arithmetic expression evaluation.
|
|
"""
|
|
|
|
import math
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
from packaging.version import parse as 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.expressions import col, lit
|
|
from ray.data.tests.conftest import * # noqa
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
pytestmark = pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("20.0.0"),
|
|
reason="Expression integration tests require PyArrow >= 20.0.0",
|
|
)
|
|
|
|
|
|
class TestArithmeticIntegration:
|
|
"""Integration tests for arithmetic expressions with Ray Dataset."""
|
|
|
|
def test_arithmetic_with_dataset(self, ray_start_regular_shared):
|
|
"""Test arithmetic expressions work correctly with Ray Dataset."""
|
|
ds = ray.data.from_items(
|
|
[
|
|
{"price": 10.0, "quantity": 2},
|
|
{"price": 20.0, "quantity": 3},
|
|
{"price": 15.0, "quantity": 4},
|
|
]
|
|
)
|
|
|
|
result = ds.with_column("total", col("price") * col("quantity")).to_pandas()
|
|
expected = pd.DataFrame(
|
|
{
|
|
"price": [10.0, 20.0, 15.0],
|
|
"quantity": [2, 3, 4],
|
|
"total": [20.0, 60.0, 60.0],
|
|
}
|
|
)
|
|
assert rows_same(result, expected)
|
|
|
|
def test_chained_arithmetic_with_dataset(self, ray_start_regular_shared):
|
|
"""Test chained arithmetic expressions with Ray Dataset."""
|
|
ds = ray.data.from_items(
|
|
[
|
|
{"a": 10, "b": 5},
|
|
{"a": 20, "b": 3},
|
|
]
|
|
)
|
|
|
|
result = (
|
|
ds.with_column("sum", col("a") + col("b"))
|
|
.with_column("diff", col("a") - col("b"))
|
|
.with_column("product", col("a") * col("b"))
|
|
.to_pandas()
|
|
)
|
|
|
|
expected = pd.DataFrame(
|
|
{
|
|
"a": [10, 20],
|
|
"b": [5, 3],
|
|
"sum": [15, 23],
|
|
"diff": [5, 17],
|
|
"product": [50, 60],
|
|
}
|
|
)
|
|
assert rows_same(result, expected)
|
|
|
|
def test_floor_division_with_dataset(self, ray_start_regular_shared):
|
|
"""Test floor division operations with Ray Dataset."""
|
|
ds = ray.data.range(5)
|
|
result = ds.with_column("result", col("id") // 2).to_pandas()
|
|
expected = pd.DataFrame({"id": [0, 1, 2, 3, 4], "result": [0, 0, 1, 1, 2]})
|
|
assert rows_same(result, expected)
|
|
|
|
def test_literal_floor_division_with_dataset(self, ray_start_regular_shared):
|
|
"""Test literal floor division by expression with Ray Dataset."""
|
|
ds = ray.data.range(5)
|
|
result = ds.with_column("result", lit(10) // (col("id") + 2)).to_pandas()
|
|
expected = pd.DataFrame({"id": [0, 1, 2, 3, 4], "result": [5, 3, 2, 2, 1]})
|
|
assert rows_same(result, expected)
|
|
|
|
@pytest.mark.parametrize(
|
|
"expr_factory,expected_values",
|
|
[
|
|
pytest.param(lambda: col("value").ceil(), [-1, 0, 0, 1, 2], id="ceil"),
|
|
pytest.param(lambda: col("value").floor(), [-2, -1, 0, 0, 1], id="floor"),
|
|
pytest.param(lambda: col("value").round(), [-2, 0, 0, 0, 2], id="round"),
|
|
pytest.param(lambda: col("value").trunc(), [-1, 0, 0, 0, 1], id="trunc"),
|
|
],
|
|
)
|
|
def test_rounding_with_dataset(
|
|
self, ray_start_regular_shared, expr_factory, expected_values
|
|
):
|
|
"""Test rounding operations with Ray Dataset."""
|
|
values = [-1.75, -0.25, 0.0, 0.25, 1.75]
|
|
ds = ray.data.from_items([{"value": v} for v in values])
|
|
result = ds.with_column("result", expr_factory()).to_pandas()
|
|
expected = pd.DataFrame({"value": values, "result": expected_values})
|
|
assert rows_same(result, expected)
|
|
|
|
@pytest.mark.parametrize(
|
|
"expr_factory,expected_fn",
|
|
[
|
|
pytest.param(lambda: col("value").ln(), math.log, id="ln"),
|
|
pytest.param(lambda: col("value").log10(), math.log10, id="log10"),
|
|
pytest.param(lambda: col("value").log2(), math.log2, id="log2"),
|
|
pytest.param(lambda: col("value").exp(), math.exp, id="exp"),
|
|
],
|
|
)
|
|
def test_logarithmic_with_dataset(
|
|
self, ray_start_regular_shared, expr_factory, expected_fn
|
|
):
|
|
"""Test logarithmic operations with Ray Dataset."""
|
|
values = [1.0, math.e, 10.0, 4.0]
|
|
ds = ray.data.from_items([{"value": v} for v in values])
|
|
expected_values = [expected_fn(v) for v in values]
|
|
result = ds.with_column("result", expr_factory()).to_pandas()
|
|
expected = pd.DataFrame({"value": values, "result": expected_values})
|
|
assert rows_same(result, expected)
|
|
|
|
@pytest.mark.parametrize(
|
|
"expr_factory,expected_fn",
|
|
[
|
|
pytest.param(lambda: col("value").sin(), math.sin, id="sin"),
|
|
pytest.param(lambda: col("value").cos(), math.cos, id="cos"),
|
|
pytest.param(lambda: col("value").tan(), math.tan, id="tan"),
|
|
pytest.param(lambda: col("value").atan(), math.atan, id="atan"),
|
|
],
|
|
)
|
|
def test_trigonometric_with_dataset(
|
|
self, ray_start_regular_shared, expr_factory, expected_fn
|
|
):
|
|
"""Test trigonometric operations with Ray Dataset."""
|
|
values = [0.0, math.pi / 6, math.pi / 4, math.pi / 3]
|
|
ds = ray.data.from_items([{"value": v} for v in values])
|
|
expected_values = [expected_fn(v) for v in values]
|
|
result = ds.with_column("result", expr_factory()).to_pandas()
|
|
expected = pd.DataFrame({"value": values, "result": expected_values})
|
|
assert rows_same(result, expected)
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_data,expr_factory,expected_results",
|
|
[
|
|
pytest.param(
|
|
[{"x": 5}, {"x": -3}, {"x": 0}],
|
|
lambda: col("x").negate(),
|
|
[-5, 3, 0],
|
|
id="negate",
|
|
),
|
|
pytest.param(
|
|
[{"x": 5}, {"x": -3}, {"x": 0}],
|
|
lambda: col("x").sign(),
|
|
[1, -1, 0],
|
|
id="sign",
|
|
),
|
|
pytest.param(
|
|
[{"x": 5}, {"x": -3}, {"x": 0}],
|
|
lambda: col("x").abs(),
|
|
[5, 3, 0],
|
|
id="abs",
|
|
),
|
|
pytest.param(
|
|
[{"x": 2}, {"x": 3}, {"x": 4}],
|
|
lambda: col("x").power(2),
|
|
[4, 9, 16],
|
|
id="power_int",
|
|
),
|
|
pytest.param(
|
|
[{"x": 4}, {"x": 9}, {"x": 16}],
|
|
lambda: col("x").power(0.5),
|
|
[2.0, 3.0, 4.0],
|
|
id="power_sqrt",
|
|
),
|
|
],
|
|
)
|
|
def test_arithmetic_helpers_with_dataset(
|
|
self, ray_start_regular_shared, test_data, expr_factory, expected_results
|
|
):
|
|
"""Test arithmetic helper operations with Ray Dataset."""
|
|
ds = ray.data.from_items(test_data)
|
|
result = ds.with_column("result", expr_factory()).to_pandas()
|
|
expected = pd.DataFrame(test_data)
|
|
expected["result"] = expected_results
|
|
assert rows_same(result, expected)
|
|
|
|
def test_age_group_calculation_with_dataset(self, ray_start_regular_shared):
|
|
"""Test floor division for grouping values (e.g., age into decades)."""
|
|
test_data = [
|
|
{"age": 25},
|
|
{"age": 17},
|
|
{"age": 30},
|
|
]
|
|
ds = ray.data.from_items(test_data)
|
|
result = ds.with_column("age_group", col("age") // 10 * 10).to_pandas()
|
|
expected = pd.DataFrame({"age": [25, 17, 30], "age_group": [20, 10, 30]})
|
|
assert rows_same(result, expected)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|