chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,209 @@
|
||||
"""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__]))
|
||||
@@ -0,0 +1,147 @@
|
||||
"""Integration tests for boolean/logical expression operations.
|
||||
|
||||
These tests require Ray and test end-to-end boolean expression evaluation.
|
||||
"""
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
from packaging.version import parse as parse_version
|
||||
|
||||
import ray
|
||||
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 TestBooleanIntegration:
|
||||
"""Integration tests for boolean expressions with Ray Dataset."""
|
||||
|
||||
def test_boolean_filter_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test boolean expressions used for filtering with Ray Dataset."""
|
||||
ds = ray.data.from_items(
|
||||
[
|
||||
{"age": 17, "is_student": True, "name": "Alice"},
|
||||
{"age": 21, "is_student": True, "name": "Bob"},
|
||||
{"age": 25, "is_student": False, "name": "Charlie"},
|
||||
{"age": 30, "is_student": False, "name": "Diana"},
|
||||
]
|
||||
)
|
||||
|
||||
# Add boolean columns using expressions
|
||||
result = (
|
||||
ds.with_column("is_adult", col("age") >= 18)
|
||||
.with_column("adult_student", (col("age") >= 18) & col("is_student"))
|
||||
.with_column("minor_or_student", (col("age") < 18) | col("is_student"))
|
||||
.to_pandas()
|
||||
)
|
||||
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"age": [17, 21, 25, 30],
|
||||
"is_student": [True, True, False, False],
|
||||
"name": ["Alice", "Bob", "Charlie", "Diana"],
|
||||
"is_adult": [False, True, True, True],
|
||||
"adult_student": [False, True, False, False],
|
||||
"minor_or_student": [True, True, False, False],
|
||||
}
|
||||
)
|
||||
pd.testing.assert_frame_equal(result, expected, check_dtype=False)
|
||||
|
||||
def test_complex_boolean_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test complex boolean expressions with Ray Dataset."""
|
||||
ds = ray.data.from_items(
|
||||
[
|
||||
{"score": 85, "passed": True, "bonus": False},
|
||||
{"score": 70, "passed": True, "bonus": True},
|
||||
{"score": 45, "passed": False, "bonus": False},
|
||||
]
|
||||
)
|
||||
|
||||
# Complex: (score > 80) OR (passed AND bonus)
|
||||
result = ds.with_column(
|
||||
"eligible", (col("score") > 80) | (col("passed") & col("bonus"))
|
||||
).to_pandas()
|
||||
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"score": [85, 70, 45],
|
||||
"passed": [True, True, False],
|
||||
"bonus": [False, True, False],
|
||||
"eligible": [True, True, False],
|
||||
}
|
||||
)
|
||||
pd.testing.assert_frame_equal(result, expected, check_dtype=False)
|
||||
|
||||
def test_logical_not_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test logical NOT operation 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": [True, True, False, True, True]}
|
||||
)
|
||||
pd.testing.assert_frame_equal(result, expected, check_dtype=False)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"expression_factory,expected_results,test_id",
|
||||
[
|
||||
pytest.param(
|
||||
lambda: (col("age") > 18) & (col("country") == "USA"),
|
||||
[True, False, False],
|
||||
"complex_and",
|
||||
),
|
||||
pytest.param(
|
||||
lambda: (col("age") < 18) | (col("country") == "USA"),
|
||||
[True, True, False],
|
||||
"complex_or",
|
||||
),
|
||||
pytest.param(
|
||||
lambda: ~((col("age") < 25) & (col("country") != "USA")),
|
||||
[True, False, True],
|
||||
"complex_not",
|
||||
),
|
||||
pytest.param(
|
||||
lambda: (col("age") >= 21)
|
||||
& (col("score") >= 10)
|
||||
& col("active").is_not_null()
|
||||
& (col("active") == lit(True)),
|
||||
[True, False, False],
|
||||
"eligibility_flag",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_complex_boolean_expressions_with_dataset(
|
||||
self, ray_start_regular_shared, expression_factory, expected_results, test_id
|
||||
):
|
||||
"""Test complex boolean expressions with Ray Dataset."""
|
||||
test_data = [
|
||||
{"age": 25, "country": "USA", "active": True, "score": 20},
|
||||
{"age": 17, "country": "Canada", "active": False, "score": 10},
|
||||
{"age": 30, "country": "UK", "active": None, "score": 20},
|
||||
]
|
||||
|
||||
ds = ray.data.from_items(test_data)
|
||||
expression = expression_factory()
|
||||
result = ds.with_column("result", expression).to_pandas()
|
||||
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"age": [25, 17, 30],
|
||||
"country": ["USA", "Canada", "UK"],
|
||||
"active": [True, False, None],
|
||||
"score": [20, 10, 20],
|
||||
"result": expected_results,
|
||||
}
|
||||
)
|
||||
pd.testing.assert_frame_equal(result, expected, check_dtype=False)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,192 @@
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
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.datatype import DataType
|
||||
from ray.data.exceptions import UserCodeException
|
||||
from ray.data.expressions import col
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
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(
|
||||
"expr, target_type, expected_rows",
|
||||
[
|
||||
# Basic type conversions using Ray Data's DataType
|
||||
(col("id"), DataType.int64(), [{"id": i, "result": i} for i in range(5)]),
|
||||
(
|
||||
col("id"),
|
||||
DataType.float64(),
|
||||
[{"id": i, "result": float(i)} for i in range(5)],
|
||||
),
|
||||
(
|
||||
col("id"),
|
||||
DataType.string(),
|
||||
[{"id": i, "result": str(i)} for i in range(5)],
|
||||
),
|
||||
(
|
||||
col("id") / 2,
|
||||
DataType.int64(),
|
||||
[{"id": i, "result": i // 2} for i in range(5)],
|
||||
),
|
||||
# col("id")/2 uses integer division in expression layer, then cast to float64
|
||||
(
|
||||
col("id") / 2,
|
||||
DataType.float64(),
|
||||
[{"id": i, "result": float(i // 2)} for i in range(5)],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_cast_expression_basic(
|
||||
ray_start_regular_shared,
|
||||
expr,
|
||||
target_type,
|
||||
expected_rows,
|
||||
target_max_block_size_infinite_or_default,
|
||||
):
|
||||
"""Test basic type casting with cast() method."""
|
||||
ds = ray.data.range(5).with_column("result", expr.cast(target_type))
|
||||
actual = ds.take_all()
|
||||
assert rows_same(pd.DataFrame(actual), pd.DataFrame(expected_rows))
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
get_pyarrow_version() < parse_version("20.0.0"),
|
||||
reason="with_column requires PyArrow >= 20.0.0",
|
||||
)
|
||||
def test_cast_expression_usecase(
|
||||
ray_start_regular_shared, target_max_block_size_infinite_or_default
|
||||
):
|
||||
"""Test the user use case: converting float result from modulo to int64."""
|
||||
ds = ray.data.range(10)
|
||||
# The modulo operation returns float, cast it to int64
|
||||
ds = ds.with_column("part", (col("id") % 2).cast(DataType.int64()))
|
||||
actual = ds.take_all()
|
||||
expected_rows = [{"id": i, "part": i % 2} for i in range(10)]
|
||||
assert rows_same(pd.DataFrame(actual), pd.DataFrame(expected_rows))
|
||||
|
||||
# Verify the schema shows int64 type
|
||||
schema = ds.schema()
|
||||
assert "part" in schema.names
|
||||
part_type = schema.types[schema.names.index("part")]
|
||||
assert part_type == pa.int64()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
get_pyarrow_version() < parse_version("20.0.0"),
|
||||
reason="with_column requires PyArrow >= 20.0.0",
|
||||
)
|
||||
def test_cast_expression_chained(
|
||||
ray_start_regular_shared, target_max_block_size_infinite_or_default
|
||||
):
|
||||
"""Test that cast() can be chained with other expressions."""
|
||||
ds = ray.data.range(5)
|
||||
# Cast to float64 then multiply
|
||||
ds = ds.with_column("result", col("id").cast(DataType.float64()) * 2.5)
|
||||
actual = ds.take_all()
|
||||
expected_rows = [{"id": i, "result": i * 2.5} for i in range(5)]
|
||||
assert rows_same(pd.DataFrame(actual), pd.DataFrame(expected_rows))
|
||||
|
||||
# Cast result of arithmetic operation
|
||||
ds = ray.data.range(5)
|
||||
ds = ds.with_column("result", (col("id") + 1).cast(DataType.string()))
|
||||
actual = ds.take_all()
|
||||
expected_rows = [{"id": i, "result": str(i + 1)} for i in range(5)]
|
||||
assert rows_same(pd.DataFrame(actual), pd.DataFrame(expected_rows))
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
get_pyarrow_version() < parse_version("20.0.0"),
|
||||
reason="with_column requires PyArrow >= 20.0.0",
|
||||
)
|
||||
def test_cast_expression_safe_mode(
|
||||
ray_start_regular_shared, target_max_block_size_infinite_or_default
|
||||
):
|
||||
"""Test that safe=True (default) raises errors on invalid conversions."""
|
||||
ds = ray.data.from_items([{"value": "not_a_number"}])
|
||||
|
||||
# Attempting to cast non-numeric string to int should raise an error
|
||||
with pytest.raises((UserCodeException, ValueError, pa.ArrowInvalid)):
|
||||
ds.with_column("result", col("value").cast(DataType.int64())).materialize()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
get_pyarrow_version() < parse_version("20.0.0"),
|
||||
reason="with_column requires PyArrow >= 20.0.0",
|
||||
)
|
||||
def test_cast_expression_invalid_type(
|
||||
ray_start_regular_shared, target_max_block_size_infinite_or_default
|
||||
):
|
||||
"""Test that invalid type targets raise appropriate errors."""
|
||||
ds = ray.data.range(5)
|
||||
|
||||
# Passing a non-DataType target should raise TypeError
|
||||
with pytest.raises(
|
||||
TypeError, match="target_type must be a ray.data.datatype.DataType"
|
||||
):
|
||||
ds.with_column("result", col("id").cast("invalid_type")).materialize()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
get_pyarrow_version() < parse_version("20.0.0"),
|
||||
reason="with_column requires PyArrow >= 20.0.0",
|
||||
)
|
||||
def test_cast_expression_multiple_types(
|
||||
ray_start_regular_shared, target_max_block_size_infinite_or_default
|
||||
):
|
||||
"""Test casting with multiple different target types."""
|
||||
ds = ray.data.from_items([{"id": 42, "score": 3.14}])
|
||||
|
||||
# Cast id to different types
|
||||
ds = ds.with_column("id_int", col("id").cast(DataType.int64()))
|
||||
ds = ds.with_column("id_float", col("id").cast(DataType.float64()))
|
||||
ds = ds.with_column("id_str", col("id").cast(DataType.string()))
|
||||
|
||||
# Cast score to int (use safe=False to allow float truncation to int)
|
||||
ds = ds.with_column("score_int", col("score").cast(DataType.int64(), safe=False))
|
||||
|
||||
# Use rows_same to compare the full row content (expects DataFrames).
|
||||
results = ds.take_all()
|
||||
expected = [
|
||||
{
|
||||
"id": 42,
|
||||
"score": 3.14,
|
||||
"id_int": 42,
|
||||
"id_float": 42.0,
|
||||
"id_str": "42",
|
||||
"score_int": 3,
|
||||
}
|
||||
]
|
||||
assert rows_same(pd.DataFrame(results), pd.DataFrame(expected))
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
get_pyarrow_version() < parse_version("20.0.0"),
|
||||
reason="with_column requires PyArrow >= 20.0.0",
|
||||
)
|
||||
def test_cast_expression_python_type_datatype_error(
|
||||
ray_start_regular_shared, target_max_block_size_infinite_or_default
|
||||
):
|
||||
"""Test that using Python-type-backed DataType in cast() raises a clear error."""
|
||||
# Error is raised at expression build time when cast() is called (not at materialize).
|
||||
error_match = "Python-type-backed DataType.*requires.*values"
|
||||
with pytest.raises(TypeError, match=error_match):
|
||||
col("id").cast(DataType(int))
|
||||
with pytest.raises(TypeError, match=error_match):
|
||||
col("id").cast(DataType(str))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,78 @@
|
||||
"""Integration tests for comparison expression operations.
|
||||
|
||||
These tests require Ray and test end-to-end comparison expression evaluation.
|
||||
"""
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
from packaging.version import parse as parse_version
|
||||
|
||||
import ray
|
||||
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
|
||||
from ray.data.expressions import col
|
||||
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 TestComparisonIntegration:
|
||||
"""Integration tests for comparison expressions with Ray Dataset."""
|
||||
|
||||
def test_comparison_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test comparison expressions work correctly with Ray Dataset."""
|
||||
ds = ray.data.from_items(
|
||||
[
|
||||
{"age": 17, "name": "Alice"},
|
||||
{"age": 21, "name": "Bob"},
|
||||
{"age": 25, "name": "Charlie"},
|
||||
{"age": 18, "name": "Diana"},
|
||||
]
|
||||
)
|
||||
|
||||
result = ds.with_column("is_adult", col("age") >= 18).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"age": [17, 21, 25, 18],
|
||||
"name": ["Alice", "Bob", "Charlie", "Diana"],
|
||||
"is_adult": [False, True, True, True],
|
||||
}
|
||||
)
|
||||
pd.testing.assert_frame_equal(result, expected, check_dtype=False)
|
||||
|
||||
def test_multiple_comparisons_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test multiple comparison expressions with Ray Dataset."""
|
||||
ds = ray.data.from_items(
|
||||
[
|
||||
{"score": 45, "passing": 50},
|
||||
{"score": 75, "passing": 50},
|
||||
{"score": 50, "passing": 50},
|
||||
]
|
||||
)
|
||||
|
||||
result = (
|
||||
ds.with_column("passed", col("score") >= col("passing"))
|
||||
.with_column("failed", col("score") < col("passing"))
|
||||
.with_column("borderline", col("score") == col("passing"))
|
||||
.to_pandas()
|
||||
)
|
||||
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"score": [45, 75, 50],
|
||||
"passing": [50, 50, 50],
|
||||
"passed": [False, True, True],
|
||||
"failed": [True, False, False],
|
||||
"borderline": [False, False, True],
|
||||
}
|
||||
)
|
||||
pd.testing.assert_frame_equal(result, expected, check_dtype=False)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,63 @@
|
||||
"""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__]))
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Integration tests for datetime namespace expressions.
|
||||
|
||||
These tests require Ray and test end-to-end datetime namespace expression evaluation.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
|
||||
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",
|
||||
)
|
||||
|
||||
|
||||
class TestDatetimeNamespace:
|
||||
"""Tests for datetime namespace operations."""
|
||||
|
||||
def test_datetime_namespace_all_operations(self, ray_start_regular_shared):
|
||||
"""Test all datetime namespace operations on a datetime column."""
|
||||
ts = datetime.datetime(2024, 1, 2, 10, 30, 0)
|
||||
|
||||
ds = ray.data.from_items([{"ts": ts}])
|
||||
|
||||
result_ds = (
|
||||
ds.with_column("year", col("ts").dt.year())
|
||||
.with_column("month", col("ts").dt.month())
|
||||
.with_column("day", col("ts").dt.day())
|
||||
.with_column("hour", col("ts").dt.hour())
|
||||
.with_column("minute", col("ts").dt.minute())
|
||||
.with_column("second", col("ts").dt.second())
|
||||
.with_column("date_str", col("ts").dt.strftime("%Y-%m-%d"))
|
||||
.with_column("ts_floor", col("ts").dt.floor("day"))
|
||||
.with_column("ts_ceil", col("ts").dt.ceil("day"))
|
||||
.with_column("ts_round", col("ts").dt.round("day"))
|
||||
.drop_columns(["ts"])
|
||||
)
|
||||
|
||||
actual = result_ds.to_pandas()
|
||||
|
||||
expected = pd.DataFrame(
|
||||
[
|
||||
{
|
||||
"year": 2024,
|
||||
"month": 1,
|
||||
"day": 2,
|
||||
"hour": 10,
|
||||
"minute": 30,
|
||||
"second": 0,
|
||||
"date_str": "2024-01-02",
|
||||
"ts_floor": pd.Timestamp("2024-01-02"),
|
||||
"ts_ceil": pd.Timestamp("2024-01-03"),
|
||||
# round("day") rounds to nearest day; 10:30 < 12:00 so rounds down
|
||||
"ts_round": pd.Timestamp("2024-01-02"),
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
assert rows_same(actual, expected)
|
||||
|
||||
def test_dt_namespace_invalid_dtype_raises(self, ray_start_regular_shared):
|
||||
"""Test that dt namespace on non-datetime column raises an error."""
|
||||
ds = ray.data.from_items([{"value": 1}])
|
||||
|
||||
with pytest.raises(Exception):
|
||||
ds.with_column("year", col("value").dt.year()).to_pandas()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,183 @@
|
||||
"""Integration tests for list namespace expressions.
|
||||
|
||||
These tests require Ray and test end-to-end list 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.exceptions import RayTaskError
|
||||
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 _create_dataset(items_data, dataset_format, arrow_table=None):
|
||||
if dataset_format == "arrow":
|
||||
if arrow_table is not None:
|
||||
ds = ray.data.from_arrow(arrow_table)
|
||||
else:
|
||||
table = pa.Table.from_pylist(items_data)
|
||||
ds = ray.data.from_arrow(table)
|
||||
elif dataset_format == "pandas":
|
||||
if arrow_table is not None:
|
||||
df = arrow_table.to_pandas()
|
||||
else:
|
||||
df = pd.DataFrame(items_data)
|
||||
ds = ray.data.from_blocks([df])
|
||||
return ds
|
||||
|
||||
|
||||
DATASET_FORMATS = ["pandas", "arrow"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
||||
class TestListNamespace:
|
||||
"""Tests for list namespace operations."""
|
||||
|
||||
def test_list_len(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test list.len() returns length of each list."""
|
||||
data = [
|
||||
{"items": [1, 2, 3]},
|
||||
{"items": [4, 5]},
|
||||
{"items": []},
|
||||
]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column("len", col("items").list.len()).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"items": [[1, 2, 3], [4, 5], []],
|
||||
"len": [3, 2, 0],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_list_get(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test list.get() extracts element at index."""
|
||||
data = [
|
||||
{"items": [10, 20, 30]},
|
||||
{"items": [40, 50, 60]},
|
||||
]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column("first", col("items").list.get(0)).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"items": [[10, 20, 30], [40, 50, 60]],
|
||||
"first": [10, 40],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_list_bracket_index(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test list[i] bracket notation for element access."""
|
||||
data = [{"items": [10, 20, 30]}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column("elem", col("items").list[1]).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"items": [[10, 20, 30]],
|
||||
"elem": [20],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_list_with_arithmetic(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test list operations combined with arithmetic."""
|
||||
data = [{"items": [1, 2, 3]}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column("len_plus_one", col("items").list.len() + 1).to_pandas()
|
||||
expected = pd.DataFrame({"items": [[1, 2, 3]], "len_plus_one": [4]})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_list_sort(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test list.sort() sorts each list with custom options."""
|
||||
data = [
|
||||
{"items": [3, 1, 2]},
|
||||
{"items": [None, 4, 2]},
|
||||
]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
method = col("items").list.sort(order="descending", null_placement="at_start")
|
||||
result = ds.with_column("sorted", method).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"items": [[3, 1, 2], [None, 4, 2]],
|
||||
"sorted": [[3, 2, 1], [None, 4, 2]],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_list_flatten(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test list.flatten() removes one nesting level."""
|
||||
data = [
|
||||
{"items": [[1, 2], [3]]},
|
||||
{"items": [[], [4, 5]]},
|
||||
]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column("flattened", col("items").list.flatten()).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"items": [[[1, 2], [3]], [[], [4, 5]]],
|
||||
"flattened": [[1, 2, 3], [4, 5]],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_list_flatten_requires_nested_lists(
|
||||
self, ray_start_regular_shared, dataset_format
|
||||
):
|
||||
"""list.flatten() should raise if elements aren't lists."""
|
||||
data = [{"items": [1, 2]}, {"items": [3, 4]}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
with pytest.raises(RayTaskError):
|
||||
ds.with_column("flattened", col("items").list.flatten()).materialize()
|
||||
|
||||
def test_list_flatten_large_list_type(
|
||||
self, ray_start_regular_shared, dataset_format
|
||||
):
|
||||
"""Flatten should preserve LargeList type when present."""
|
||||
if dataset_format != "arrow":
|
||||
pytest.skip("LargeList type only available via Arrow tables.")
|
||||
|
||||
arrow_type = pa.large_list(pa.list_(pa.int64()))
|
||||
table = pa.Table.from_arrays(
|
||||
[
|
||||
pa.array(
|
||||
[
|
||||
[[1, 2], [3]],
|
||||
[[], [4, 5]],
|
||||
],
|
||||
type=arrow_type,
|
||||
)
|
||||
],
|
||||
names=["items"],
|
||||
)
|
||||
ds = _create_dataset(None, dataset_format, arrow_table=table)
|
||||
result = ds.with_column("flattened", col("items").list.flatten())
|
||||
arrow_refs = result.to_arrow_refs()
|
||||
tables = ray.get(arrow_refs)
|
||||
result_table = pa.concat_tables(tables) if len(tables) > 1 else tables[0]
|
||||
|
||||
flattened_type = result_table.schema.field("flattened").type
|
||||
assert flattened_type == pa.large_list(pa.int64())
|
||||
expected = pa.Table.from_arrays(
|
||||
[
|
||||
pa.array([[1, 2, 3], [4, 5]], type=pa.large_list(pa.int64())),
|
||||
],
|
||||
names=["flattened"],
|
||||
)
|
||||
assert result_table.select(["flattened"]).combine_chunks().equals(expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,150 @@
|
||||
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
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
version.parse(pa.__version__) < version.parse("19.0.0"),
|
||||
reason="Namespace expressions tests require PyArrow >= 19.0",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def map_dataset():
|
||||
"""Fixture that creates a dataset backed by an Arrow MapArray column."""
|
||||
map_items = [
|
||||
{"attrs": {"color": "red", "size": "M"}},
|
||||
{"attrs": {"brand": "Ray"}},
|
||||
]
|
||||
map_type = pa.map_(pa.string(), pa.string())
|
||||
arrow_table = pa.table(
|
||||
{"attrs": pa.array([row["attrs"] for row in map_items], type=map_type)}
|
||||
)
|
||||
return ray.data.from_arrow(arrow_table)
|
||||
|
||||
|
||||
def _assert_result(result_df: pd.DataFrame, expected_df: pd.DataFrame, drop_cols: list):
|
||||
"""Helper to drop columns and assert equality."""
|
||||
result_df = result_df.drop(columns=drop_cols)
|
||||
assert rows_same(result_df, expected_df)
|
||||
|
||||
|
||||
class TestMapNamespace:
|
||||
"""Tests for map namespace operations using the shared map_dataset fixture."""
|
||||
|
||||
def test_map_keys(self, map_dataset):
|
||||
result = map_dataset.with_column("keys", col("attrs").map.keys()).to_pandas()
|
||||
expected = pd.DataFrame({"keys": [["color", "size"], ["brand"]]})
|
||||
_assert_result(result, expected, drop_cols=["attrs"])
|
||||
|
||||
def test_map_values(self, map_dataset):
|
||||
result = map_dataset.with_column(
|
||||
"values", col("attrs").map.values()
|
||||
).to_pandas()
|
||||
expected = pd.DataFrame({"values": [["red", "M"], ["Ray"]]})
|
||||
_assert_result(result, expected, drop_cols=["attrs"])
|
||||
|
||||
def test_map_chaining(self, map_dataset):
|
||||
# map.keys() returns a list, so .list.len() should apply
|
||||
result = map_dataset.with_column(
|
||||
"num_keys", col("attrs").map.keys().list.len()
|
||||
).to_pandas()
|
||||
expected = pd.DataFrame({"num_keys": [2, 1]})
|
||||
_assert_result(result, expected, drop_cols=["attrs"])
|
||||
|
||||
|
||||
def test_physical_map_extraction():
|
||||
"""Test extraction works on List<Struct> (Physical Maps)."""
|
||||
# Construct List<Struct<k, v>>
|
||||
struct_type = pa.struct([pa.field("k", pa.string()), pa.field("v", pa.int64())])
|
||||
list_type = pa.list_(struct_type)
|
||||
|
||||
data_py = [[{"k": "a", "v": 1}], [{"k": "b", "v": 2}]]
|
||||
arrow_table = pa.Table.from_arrays(
|
||||
[pa.array(data_py, type=list_type)], names=["data"]
|
||||
)
|
||||
ds = ray.data.from_arrow(arrow_table)
|
||||
|
||||
result = (
|
||||
ds.with_column("keys", col("data").map.keys())
|
||||
.with_column("values", col("data").map.values())
|
||||
.to_pandas()
|
||||
)
|
||||
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"data": data_py,
|
||||
"keys": [["a"], ["b"]],
|
||||
"values": [[1], [2]],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
def test_map_sliced_offsets():
|
||||
"""Test extraction works correctly on sliced Arrow arrays (offset > 0)."""
|
||||
items = [{"m": {"id": i}} for i in range(10)]
|
||||
map_type = pa.map_(pa.string(), pa.int64())
|
||||
arrays = pa.array([row["m"] for row in items], type=map_type)
|
||||
table = pa.Table.from_arrays([arrays], names=["m"])
|
||||
|
||||
# Force offsets by slicing the table before ingestion
|
||||
sliced_table = table.slice(offset=7, length=3)
|
||||
ds = ray.data.from_arrow(sliced_table)
|
||||
|
||||
result = ds.with_column("vals", col("m").map.values()).to_pandas()
|
||||
expected = pd.DataFrame({"vals": [[7], [8], [9]]})
|
||||
_assert_result(result, expected, drop_cols=["m"])
|
||||
|
||||
|
||||
def test_map_nulls_and_empty():
|
||||
"""Test handling of null maps and empty maps."""
|
||||
items_data = [{"m": {"a": 1}}, {"m": {}}, {"m": None}]
|
||||
|
||||
map_type = pa.map_(pa.string(), pa.int64())
|
||||
arrays = pa.array([row["m"] for row in items_data], type=map_type)
|
||||
arrow_table = pa.Table.from_arrays([arrays], names=["m"])
|
||||
ds = ray.data.from_arrow(arrow_table)
|
||||
|
||||
rows = (
|
||||
ds.with_column("keys", col("m").map.keys())
|
||||
.with_column("values", col("m").map.values())
|
||||
.take_all()
|
||||
)
|
||||
|
||||
assert list(rows[0]["keys"]) == ["a"] and list(rows[0]["values"]) == [1]
|
||||
assert len(rows[1]["keys"]) == 0 and len(rows[1]["values"]) == 0
|
||||
assert rows[2]["keys"] is None and rows[2]["values"] is None
|
||||
|
||||
|
||||
def test_empty_chunked_array():
|
||||
"""Test extraction works on empty ChunkedArray (zero chunks)."""
|
||||
from ray.data.namespace_expressions.map_namespace import (
|
||||
MapComponent,
|
||||
_extract_map_component,
|
||||
)
|
||||
|
||||
# Create empty ChunkedArray with map type
|
||||
map_type = pa.map_(pa.string(), pa.int64())
|
||||
empty_chunked = pa.chunked_array([], type=map_type)
|
||||
assert empty_chunked.num_chunks == 0
|
||||
|
||||
# Extract keys - should return empty ChunkedArray with list<string> type
|
||||
keys_result = _extract_map_component(empty_chunked, MapComponent.KEYS)
|
||||
assert len(keys_result) == 0
|
||||
assert keys_result.type == pa.list_(pa.string())
|
||||
|
||||
# Extract values - should return empty ChunkedArray with list<int64> type
|
||||
values_result = _extract_map_component(empty_chunked, MapComponent.VALUES)
|
||||
assert len(values_result) == 0
|
||||
assert values_result.type == pa.list_(pa.int64())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,328 @@
|
||||
"""Integration tests for string namespace expressions.
|
||||
|
||||
These tests require Ray and test end-to-end string 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 _create_dataset(items_data, dataset_format, arrow_table=None):
|
||||
if dataset_format == "arrow":
|
||||
if arrow_table is not None:
|
||||
ds = ray.data.from_arrow(arrow_table)
|
||||
else:
|
||||
table = pa.Table.from_pylist(items_data)
|
||||
ds = ray.data.from_arrow(table)
|
||||
elif dataset_format == "pandas":
|
||||
if arrow_table is not None:
|
||||
df = arrow_table.to_pandas()
|
||||
else:
|
||||
df = pd.DataFrame(items_data)
|
||||
ds = ray.data.from_blocks([df])
|
||||
return ds
|
||||
|
||||
|
||||
DATASET_FORMATS = ["pandas", "arrow"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
||||
@pytest.mark.parametrize(
|
||||
"method_name,input_values,expected_results",
|
||||
[
|
||||
("len", ["Alice", "Bob"], [5, 3]),
|
||||
("byte_len", ["ABC"], [3]),
|
||||
],
|
||||
)
|
||||
class TestStringLength:
|
||||
"""Tests for string length operations."""
|
||||
|
||||
def test_string_length(
|
||||
self,
|
||||
ray_start_regular_shared,
|
||||
dataset_format,
|
||||
method_name,
|
||||
input_values,
|
||||
expected_results,
|
||||
):
|
||||
"""Test string length methods."""
|
||||
data = [{"name": v} for v in input_values]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
|
||||
method = getattr(col("name").str, method_name)
|
||||
result = ds.with_column("result", method()).to_pandas()
|
||||
|
||||
expected = pd.DataFrame({"name": input_values, "result": expected_results})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
||||
@pytest.mark.parametrize(
|
||||
"method_name,input_values,expected_values",
|
||||
[
|
||||
("upper", ["alice", "bob"], ["ALICE", "BOB"]),
|
||||
("lower", ["ALICE", "BOB"], ["alice", "bob"]),
|
||||
("capitalize", ["alice", "bob"], ["Alice", "Bob"]),
|
||||
("title", ["alice smith", "bob jones"], ["Alice Smith", "Bob Jones"]),
|
||||
("swapcase", ["AlIcE"], ["aLiCe"]),
|
||||
],
|
||||
)
|
||||
class TestStringCase:
|
||||
"""Tests for string case conversion."""
|
||||
|
||||
def test_string_case(
|
||||
self,
|
||||
ray_start_regular_shared,
|
||||
dataset_format,
|
||||
method_name,
|
||||
input_values,
|
||||
expected_values,
|
||||
):
|
||||
"""Test string case conversion methods."""
|
||||
data = [{"name": v} for v in input_values]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
|
||||
method = getattr(col("name").str, method_name)
|
||||
result = ds.with_column("result", method()).to_pandas()
|
||||
|
||||
expected = pd.DataFrame({"name": input_values, "result": expected_values})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
||||
@pytest.mark.parametrize(
|
||||
"method_name,input_values,expected_results",
|
||||
[
|
||||
("is_alpha", ["abc", "abc123", "123"], [True, False, False]),
|
||||
("is_alnum", ["abc123", "abc-123"], [True, False]),
|
||||
("is_digit", ["123", "12a"], [True, False]),
|
||||
("is_space", [" ", " a "], [True, False]),
|
||||
("is_lower", ["abc", "Abc"], [True, False]),
|
||||
("is_upper", ["ABC", "Abc"], [True, False]),
|
||||
("is_ascii", ["hello", "hello😊"], [True, False]),
|
||||
],
|
||||
)
|
||||
class TestStringPredicates:
|
||||
"""Tests for string predicate methods (is_*)."""
|
||||
|
||||
def test_string_predicate(
|
||||
self,
|
||||
ray_start_regular_shared,
|
||||
dataset_format,
|
||||
method_name,
|
||||
input_values,
|
||||
expected_results,
|
||||
):
|
||||
"""Test string predicate methods."""
|
||||
data = [{"val": v} for v in input_values]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
|
||||
method = getattr(col("val").str, method_name)
|
||||
result = ds.with_column("result", method()).to_pandas()
|
||||
|
||||
expected = pd.DataFrame({"val": input_values, "result": expected_results})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
||||
@pytest.mark.parametrize(
|
||||
"method_name,method_args,input_values,expected_values",
|
||||
[
|
||||
("strip", (), [" hello ", " world "], ["hello", "world"]),
|
||||
("strip", ("x",), ["xxxhelloxxx"], ["hello"]),
|
||||
("lstrip", (), [" hello "], ["hello "]),
|
||||
("rstrip", (), [" hello "], [" hello"]),
|
||||
],
|
||||
)
|
||||
class TestStringTrimming:
|
||||
"""Tests for string trimming operations."""
|
||||
|
||||
def test_string_trimming(
|
||||
self,
|
||||
ray_start_regular_shared,
|
||||
dataset_format,
|
||||
method_name,
|
||||
method_args,
|
||||
input_values,
|
||||
expected_values,
|
||||
):
|
||||
"""Test string trimming methods."""
|
||||
data = [{"val": v} for v in input_values]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
|
||||
method = getattr(col("val").str, method_name)
|
||||
result = ds.with_column("result", method(*method_args)).to_pandas()
|
||||
|
||||
expected = pd.DataFrame({"val": input_values, "result": expected_values})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
||||
@pytest.mark.parametrize(
|
||||
"method_name,method_kwargs,expected_value",
|
||||
[
|
||||
("pad", {"width": 5, "fillchar": "*", "side": "right"}, "hi***"),
|
||||
("pad", {"width": 5, "fillchar": "*", "side": "left"}, "***hi"),
|
||||
("pad", {"width": 6, "fillchar": "*", "side": "both"}, "**hi**"),
|
||||
("lpad", {"width": 5, "padding": "*"}, "***hi"),
|
||||
("rpad", {"width": 5, "padding": "*"}, "hi***"),
|
||||
("center", {"width": 6, "padding": "*"}, "**hi**"),
|
||||
],
|
||||
)
|
||||
class TestStringPadding:
|
||||
"""Tests for string padding operations."""
|
||||
|
||||
def test_string_padding(
|
||||
self,
|
||||
ray_start_regular_shared,
|
||||
dataset_format,
|
||||
method_name,
|
||||
method_kwargs,
|
||||
expected_value,
|
||||
):
|
||||
"""Test string padding methods."""
|
||||
data = [{"val": "hi"}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
|
||||
method = getattr(col("val").str, method_name)
|
||||
result = ds.with_column("result", method(**method_kwargs)).to_pandas()
|
||||
|
||||
expected = pd.DataFrame({"val": ["hi"], "result": [expected_value]})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
||||
@pytest.mark.parametrize(
|
||||
"method_name,method_args,method_kwargs,input_values,expected_results",
|
||||
[
|
||||
("starts_with", ("A",), {}, ["Alice", "Bob", "Alex"], [True, False, True]),
|
||||
("starts_with", ("A",), {"ignore_case": True}, ["alice", "bob"], [True, False]),
|
||||
("ends_with", ("e",), {}, ["Alice", "Bob"], [True, False]),
|
||||
("contains", ("li",), {}, ["Alice", "Bob", "Charlie"], [True, False, True]),
|
||||
("find", ("i",), {}, ["Alice", "Bob"], [2, -1]),
|
||||
("count", ("a",), {}, ["banana", "apple"], [3, 1]),
|
||||
("match", ("Al%",), {}, ["Alice", "Bob", "Alex"], [True, False, True]),
|
||||
],
|
||||
)
|
||||
class TestStringSearch:
|
||||
"""Tests for string searching operations."""
|
||||
|
||||
def test_string_search(
|
||||
self,
|
||||
ray_start_regular_shared,
|
||||
dataset_format,
|
||||
method_name,
|
||||
method_args,
|
||||
method_kwargs,
|
||||
input_values,
|
||||
expected_results,
|
||||
):
|
||||
"""Test string searching methods."""
|
||||
data = [{"val": v} for v in input_values]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
|
||||
method = getattr(col("val").str, method_name)
|
||||
result = ds.with_column(
|
||||
"result", method(*method_args, **method_kwargs)
|
||||
).to_pandas()
|
||||
|
||||
expected = pd.DataFrame({"val": input_values, "result": expected_results})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
||||
class TestStringTransform:
|
||||
"""Tests for string transformation operations."""
|
||||
|
||||
def test_reverse(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test str.reverse() reverses strings."""
|
||||
data = [{"val": "hello"}, {"val": "world"}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column("rev", col("val").str.reverse()).to_pandas()
|
||||
expected = pd.DataFrame({"val": ["hello", "world"], "rev": ["olleh", "dlrow"]})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_slice(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test str.slice() extracts substring."""
|
||||
data = [{"val": "hello"}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column("sliced", col("val").str.slice(1, 4)).to_pandas()
|
||||
expected = pd.DataFrame({"val": ["hello"], "sliced": ["ell"]})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_replace(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test str.replace() replaces substring."""
|
||||
data = [{"val": "hello world"}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column(
|
||||
"replaced", col("val").str.replace("world", "universe")
|
||||
).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{"val": ["hello world"], "replaced": ["hello universe"]}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_replace_with_max(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test str.replace() with max_replacements."""
|
||||
data = [{"val": "aaa"}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column(
|
||||
"replaced", col("val").str.replace("a", "X", max_replacements=2)
|
||||
).to_pandas()
|
||||
expected = pd.DataFrame({"val": ["aaa"], "replaced": ["XXa"]})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_repeat(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test str.repeat() repeats strings."""
|
||||
data = [{"val": "A"}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column("repeated", col("val").str.repeat(3)).to_pandas()
|
||||
expected = pd.DataFrame({"val": ["A"], "repeated": ["AAA"]})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_string_with_comparison(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test string operations combined with comparison."""
|
||||
data = [{"name": "Alice"}, {"name": "Bo"}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column("long_name", col("name").str.len() > 3).to_pandas()
|
||||
expected = pd.DataFrame({"name": ["Alice", "Bo"], "long_name": [True, False]})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_multiple_string_operations(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test multiple namespace operations in single pipeline."""
|
||||
data = [{"name": "alice"}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = (
|
||||
ds.with_column("upper", col("name").str.upper())
|
||||
.with_column("len", col("name").str.len())
|
||||
.with_column("starts_a", col("name").str.starts_with("a"))
|
||||
.to_pandas()
|
||||
)
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"name": ["alice"],
|
||||
"upper": ["ALICE"],
|
||||
"len": [5],
|
||||
"starts_a": [True],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,313 @@
|
||||
"""Integration tests for struct namespace expressions.
|
||||
|
||||
These tests require Ray and test end-to-end struct 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 _create_dataset(items_data, dataset_format, arrow_table=None):
|
||||
if dataset_format == "arrow":
|
||||
if arrow_table is not None:
|
||||
ds = ray.data.from_arrow(arrow_table)
|
||||
else:
|
||||
table = pa.Table.from_pylist(items_data)
|
||||
ds = ray.data.from_arrow(table)
|
||||
elif dataset_format == "pandas":
|
||||
if arrow_table is not None:
|
||||
df = arrow_table.to_pandas()
|
||||
else:
|
||||
df = pd.DataFrame(items_data)
|
||||
ds = ray.data.from_blocks([df])
|
||||
return ds
|
||||
|
||||
|
||||
DATASET_FORMATS = ["pandas", "arrow"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
||||
class TestStructNamespace:
|
||||
"""Tests for struct namespace operations."""
|
||||
|
||||
def test_struct_bracket_bool_index_raises(self, dataset_format):
|
||||
"""Test struct[bool] raises TypeError instead of being treated as int."""
|
||||
del dataset_format # Unused, required by class-level parametrization.
|
||||
|
||||
with pytest.raises(
|
||||
TypeError, match="Struct indices must be strings or integers"
|
||||
):
|
||||
col("user").struct[True]
|
||||
|
||||
@pytest.mark.parametrize("bad_index", ["1", True])
|
||||
def test_struct_field_by_index_non_integer_raises(self, dataset_format, bad_index):
|
||||
"""Test struct.field_by_index() rejects non-integer indices."""
|
||||
del dataset_format # Unused, required by class-level parametrization.
|
||||
|
||||
with pytest.raises(TypeError, match="Struct field index must be an integer"):
|
||||
col("user").struct.field_by_index(bad_index)
|
||||
|
||||
def test_struct_field_by_index_negative_raises(self, dataset_format):
|
||||
"""Test struct.field_by_index() rejects negative indices."""
|
||||
del dataset_format # Unused, required by class-level parametrization.
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="Struct field index must be non-negative, got -1"
|
||||
):
|
||||
col("user").struct.field_by_index(-1)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="Struct field index must be non-negative, got -1"
|
||||
):
|
||||
col("user").struct[-1]
|
||||
|
||||
def test_struct_field(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test struct.field() extracts field."""
|
||||
arrow_table = pa.table(
|
||||
{
|
||||
"user": pa.array(
|
||||
[
|
||||
{"name": "Alice", "age": 30},
|
||||
{"name": "Bob", "age": 25},
|
||||
],
|
||||
type=pa.struct(
|
||||
[
|
||||
pa.field("name", pa.string()),
|
||||
pa.field("age", pa.int32()),
|
||||
]
|
||||
),
|
||||
)
|
||||
}
|
||||
)
|
||||
items_data = [
|
||||
{"user": {"name": "Alice", "age": 30}},
|
||||
{"user": {"name": "Bob", "age": 25}},
|
||||
]
|
||||
ds = _create_dataset(items_data, dataset_format, arrow_table)
|
||||
|
||||
result = ds.with_column("age", col("user").struct.field("age")).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"user": [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}],
|
||||
"age": [30, 25],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_struct_bracket(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test struct['field'] bracket notation."""
|
||||
arrow_table = pa.table(
|
||||
{
|
||||
"user": pa.array(
|
||||
[
|
||||
{"name": "Alice", "age": 30},
|
||||
{"name": "Bob", "age": 25},
|
||||
],
|
||||
type=pa.struct(
|
||||
[
|
||||
pa.field("name", pa.string()),
|
||||
pa.field("age", pa.int32()),
|
||||
]
|
||||
),
|
||||
)
|
||||
}
|
||||
)
|
||||
items_data = [
|
||||
{"user": {"name": "Alice", "age": 30}},
|
||||
{"user": {"name": "Bob", "age": 25}},
|
||||
]
|
||||
ds = _create_dataset(items_data, dataset_format, arrow_table)
|
||||
|
||||
result = ds.with_column("name", col("user").struct["name"]).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"user": [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}],
|
||||
"name": ["Alice", "Bob"],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_struct_field_by_index(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test struct.field_by_index() extracts field by position."""
|
||||
if dataset_format == "pandas":
|
||||
pytest.skip(
|
||||
"Index-based struct access requires stable Arrow struct field ordering."
|
||||
)
|
||||
arrow_table = pa.table(
|
||||
{
|
||||
"user": pa.array(
|
||||
[
|
||||
{"name": "Alice", "age": 30},
|
||||
{"name": "Bob", "age": 25},
|
||||
],
|
||||
type=pa.struct(
|
||||
[
|
||||
pa.field("name", pa.string()),
|
||||
pa.field("age", pa.int32()),
|
||||
]
|
||||
),
|
||||
)
|
||||
}
|
||||
)
|
||||
items_data = [
|
||||
{"user": {"name": "Alice", "age": 30}},
|
||||
{"user": {"name": "Bob", "age": 25}},
|
||||
]
|
||||
ds = _create_dataset(items_data, dataset_format, arrow_table)
|
||||
|
||||
result = ds.with_column("age", col("user").struct.field_by_index(1)).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"user": [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}],
|
||||
"age": [30, 25],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_struct_bracket_with_index(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test struct[index] bracket notation."""
|
||||
if dataset_format == "pandas":
|
||||
pytest.skip(
|
||||
"Index-based struct access requires stable Arrow struct field ordering."
|
||||
)
|
||||
arrow_table = pa.table(
|
||||
{
|
||||
"user": pa.array(
|
||||
[
|
||||
{"name": "Alice", "age": 30},
|
||||
{"name": "Bob", "age": 25},
|
||||
],
|
||||
type=pa.struct(
|
||||
[
|
||||
pa.field("name", pa.string()),
|
||||
pa.field("age", pa.int32()),
|
||||
]
|
||||
),
|
||||
)
|
||||
}
|
||||
)
|
||||
items_data = [
|
||||
{"user": {"name": "Alice", "age": 30}},
|
||||
{"user": {"name": "Bob", "age": 25}},
|
||||
]
|
||||
ds = _create_dataset(items_data, dataset_format, arrow_table)
|
||||
|
||||
result = ds.with_column("name", col("user").struct[0]).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"user": [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}],
|
||||
"name": ["Alice", "Bob"],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_struct_nested_field(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test nested struct field access with .field()."""
|
||||
arrow_table = pa.table(
|
||||
{
|
||||
"user": pa.array(
|
||||
[
|
||||
{"name": "Alice", "address": {"city": "NYC", "zip": "10001"}},
|
||||
{"name": "Bob", "address": {"city": "LA", "zip": "90001"}},
|
||||
],
|
||||
type=pa.struct(
|
||||
[
|
||||
pa.field("name", pa.string()),
|
||||
pa.field(
|
||||
"address",
|
||||
pa.struct(
|
||||
[
|
||||
pa.field("city", pa.string()),
|
||||
pa.field("zip", pa.string()),
|
||||
]
|
||||
),
|
||||
),
|
||||
]
|
||||
),
|
||||
)
|
||||
}
|
||||
)
|
||||
items_data = [
|
||||
{"user": {"name": "Alice", "address": {"city": "NYC", "zip": "10001"}}},
|
||||
{"user": {"name": "Bob", "address": {"city": "LA", "zip": "90001"}}},
|
||||
]
|
||||
ds = _create_dataset(items_data, dataset_format, arrow_table)
|
||||
|
||||
result = ds.with_column(
|
||||
"city", col("user").struct.field("address").struct.field("city")
|
||||
).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"user": [
|
||||
{"name": "Alice", "address": {"city": "NYC", "zip": "10001"}},
|
||||
{"name": "Bob", "address": {"city": "LA", "zip": "90001"}},
|
||||
],
|
||||
"city": ["NYC", "LA"],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_struct_nested_bracket(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test nested struct field access with brackets."""
|
||||
arrow_table = pa.table(
|
||||
{
|
||||
"user": pa.array(
|
||||
[
|
||||
{"name": "Alice", "address": {"city": "NYC", "zip": "10001"}},
|
||||
{"name": "Bob", "address": {"city": "LA", "zip": "90001"}},
|
||||
],
|
||||
type=pa.struct(
|
||||
[
|
||||
pa.field("name", pa.string()),
|
||||
pa.field(
|
||||
"address",
|
||||
pa.struct(
|
||||
[
|
||||
pa.field("city", pa.string()),
|
||||
pa.field("zip", pa.string()),
|
||||
]
|
||||
),
|
||||
),
|
||||
]
|
||||
),
|
||||
)
|
||||
}
|
||||
)
|
||||
items_data = [
|
||||
{"user": {"name": "Alice", "address": {"city": "NYC", "zip": "10001"}}},
|
||||
{"user": {"name": "Bob", "address": {"city": "LA", "zip": "90001"}}},
|
||||
]
|
||||
ds = _create_dataset(items_data, dataset_format, arrow_table)
|
||||
|
||||
result = ds.with_column(
|
||||
"zip", col("user").struct["address"].struct["zip"]
|
||||
).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"user": [
|
||||
{"name": "Alice", "address": {"city": "NYC", "zip": "10001"}},
|
||||
{"name": "Bob", "address": {"city": "LA", "zip": "90001"}},
|
||||
],
|
||||
"zip": ["10001", "90001"],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,287 @@
|
||||
"""Integration tests for predicate expression operations.
|
||||
|
||||
These tests require Ray and test end-to-end predicate expression evaluation.
|
||||
"""
|
||||
|
||||
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
|
||||
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 TestPredicateIntegration:
|
||||
"""Integration tests for predicate expressions with Ray Dataset."""
|
||||
|
||||
def test_null_predicates_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test null predicate expressions with Ray Dataset."""
|
||||
ds = ray.data.from_items(
|
||||
[
|
||||
{"value": 10, "name": "Alice"},
|
||||
{"value": None, "name": "Bob"},
|
||||
{"value": 30, "name": None},
|
||||
{"value": None, "name": None},
|
||||
]
|
||||
)
|
||||
|
||||
result = (
|
||||
ds.with_column("value_is_null", col("value").is_null())
|
||||
.with_column("name_not_null", col("name").is_not_null())
|
||||
.with_column(
|
||||
"both_present", col("value").is_not_null() & col("name").is_not_null()
|
||||
)
|
||||
.to_pandas()
|
||||
)
|
||||
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"value": [10, None, 30, None],
|
||||
"name": ["Alice", "Bob", None, None],
|
||||
"value_is_null": [False, True, False, True],
|
||||
"name_not_null": [True, True, False, False],
|
||||
"both_present": [True, False, False, False],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_membership_predicates_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test membership predicate expressions with Ray Dataset."""
|
||||
ds = ray.data.from_items(
|
||||
[
|
||||
{"status": "active", "category": "A"},
|
||||
{"status": "inactive", "category": "B"},
|
||||
{"status": "pending", "category": "A"},
|
||||
{"status": "deleted", "category": "C"},
|
||||
]
|
||||
)
|
||||
|
||||
result = (
|
||||
ds.with_column(
|
||||
"is_valid_status", col("status").is_in(["active", "pending"])
|
||||
)
|
||||
.with_column("not_deleted", col("status").not_in(["deleted"]))
|
||||
.with_column("category_a", col("category").is_in(["A"]))
|
||||
.to_pandas()
|
||||
)
|
||||
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"status": ["active", "inactive", "pending", "deleted"],
|
||||
"category": ["A", "B", "A", "C"],
|
||||
"is_valid_status": [True, False, True, False],
|
||||
"not_deleted": [True, True, True, False],
|
||||
"category_a": [True, False, True, False],
|
||||
}
|
||||
)
|
||||
pd.testing.assert_frame_equal(result, expected, check_dtype=False)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"test_data,expression,expected_results,test_id",
|
||||
[
|
||||
pytest.param(
|
||||
[{"value": 1}, {"value": None}, {"value": 3}],
|
||||
col("value").is_null(),
|
||||
[False, True, False],
|
||||
"is_null_with_actual_nulls",
|
||||
),
|
||||
pytest.param(
|
||||
[{"value": 1}, {"value": None}, {"value": 3}],
|
||||
col("value").is_not_null(),
|
||||
[True, False, True],
|
||||
"is_not_null_with_actual_nulls",
|
||||
),
|
||||
pytest.param(
|
||||
[{"value": 1}, {"value": 2}, {"value": 3}],
|
||||
col("value").is_in([1, 3]),
|
||||
[True, False, True],
|
||||
"isin_operation",
|
||||
),
|
||||
pytest.param(
|
||||
[{"value": 1}, {"value": 2}, {"value": 3}],
|
||||
col("value").not_in([1, 3]),
|
||||
[False, True, False],
|
||||
"not_in_operation",
|
||||
),
|
||||
pytest.param(
|
||||
[{"name": "Alice"}, {"name": "Bob"}, {"name": "Charlie"}],
|
||||
col("name") == "Bob",
|
||||
[False, True, False],
|
||||
"string_equality",
|
||||
),
|
||||
pytest.param(
|
||||
[{"name": "Alice"}, {"name": "Bob"}, {"name": "Charlie"}],
|
||||
col("name") != "Bob",
|
||||
[True, False, True],
|
||||
"string_not_equal",
|
||||
),
|
||||
pytest.param(
|
||||
[{"name": "included"}, {"name": "excluded"}, {"name": None}],
|
||||
col("name").is_not_null() & (col("name") != "excluded"),
|
||||
[True, False, False],
|
||||
"string_filter",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_null_and_membership_with_dataset(
|
||||
self, ray_start_regular_shared, test_data, expression, expected_results, test_id
|
||||
):
|
||||
"""Test null checking and membership operations with Ray Dataset."""
|
||||
ds = ray.data.from_items(test_data)
|
||||
result = ds.with_column("result", expression).to_pandas()
|
||||
|
||||
expected_data = {}
|
||||
for key in test_data[0].keys():
|
||||
expected_data[key] = [row[key] for row in test_data]
|
||||
expected_data["result"] = expected_results
|
||||
expected = pd.DataFrame(expected_data)
|
||||
|
||||
assert rows_same(result, expected)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"filter_expr,test_data,expected_flags,test_id",
|
||||
[
|
||||
pytest.param(
|
||||
col("age") >= 21,
|
||||
[
|
||||
{"age": 20, "name": "Alice"},
|
||||
{"age": 21, "name": "Bob"},
|
||||
{"age": 25, "name": "Charlie"},
|
||||
],
|
||||
[False, True, True],
|
||||
"age_filter",
|
||||
),
|
||||
pytest.param(
|
||||
col("score") > 50,
|
||||
[
|
||||
{"score": 30, "status": "fail"},
|
||||
{"score": 50, "status": "pass"},
|
||||
{"score": 70, "status": "pass"},
|
||||
],
|
||||
[False, False, True],
|
||||
"score_filter",
|
||||
),
|
||||
pytest.param(
|
||||
(col("age") >= 18) & col("active"),
|
||||
[
|
||||
{"age": 17, "active": True},
|
||||
{"age": 18, "active": False},
|
||||
{"age": 25, "active": True},
|
||||
],
|
||||
[False, False, True],
|
||||
"complex_and_filter",
|
||||
),
|
||||
pytest.param(
|
||||
(col("status") == "approved") | (col("priority") == "high"),
|
||||
[
|
||||
{"status": "pending", "priority": "low"},
|
||||
{"status": "approved", "priority": "low"},
|
||||
{"status": "pending", "priority": "high"},
|
||||
],
|
||||
[False, True, True],
|
||||
"complex_or_filter",
|
||||
),
|
||||
pytest.param(
|
||||
col("value").is_not_null() & (col("value") > 0),
|
||||
[{"value": None}, {"value": -5}, {"value": 10}],
|
||||
[False, False, True],
|
||||
"null_aware_filter",
|
||||
),
|
||||
pytest.param(
|
||||
col("name").is_not_null() & (col("name") != "excluded"),
|
||||
[{"name": "included"}, {"name": "excluded"}, {"name": None}],
|
||||
[True, False, False],
|
||||
"string_filter",
|
||||
),
|
||||
pytest.param(
|
||||
col("category").is_in(["A", "B"]),
|
||||
[
|
||||
{"category": "A"},
|
||||
{"category": "B"},
|
||||
{"category": "C"},
|
||||
{"category": "D"},
|
||||
],
|
||||
[True, True, False, False],
|
||||
"membership_filter",
|
||||
),
|
||||
pytest.param(
|
||||
(col("score") >= 50) & (col("grade") != "F"),
|
||||
[
|
||||
{"score": 45, "grade": "F"},
|
||||
{"score": 55, "grade": "D"},
|
||||
{"score": 75, "grade": "B"},
|
||||
{"score": 30, "grade": "F"},
|
||||
],
|
||||
[False, True, True, False],
|
||||
"nested_filters",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_filter_expressions_with_dataset(
|
||||
self, ray_start_regular_shared, filter_expr, test_data, expected_flags, test_id
|
||||
):
|
||||
"""Test filter expressions with Ray Dataset."""
|
||||
ds = ray.data.from_items(test_data)
|
||||
result = ds.with_column("is_filtered", filter_expr).to_pandas()
|
||||
|
||||
expected = pd.DataFrame(test_data)
|
||||
expected["is_filtered"] = expected_flags
|
||||
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_filter_in_pipeline_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test filter expressions in a data processing pipeline."""
|
||||
test_data = [
|
||||
{"product": "A", "quantity": 10, "price": 100, "region": "North"},
|
||||
{"product": "B", "quantity": 5, "price": 200, "region": "South"},
|
||||
{"product": "C", "quantity": 20, "price": 50, "region": "North"},
|
||||
{"product": "D", "quantity": 15, "price": 75, "region": "East"},
|
||||
{"product": "E", "quantity": 3, "price": 300, "region": "West"},
|
||||
]
|
||||
|
||||
ds = ray.data.from_items(test_data)
|
||||
|
||||
result = (
|
||||
ds.with_column("revenue", col("quantity") * col("price"))
|
||||
.with_column("is_high_value", col("revenue") >= 1000)
|
||||
.with_column("is_bulk_order", col("quantity") >= 10)
|
||||
.with_column("is_premium", col("price") >= 100)
|
||||
.with_column(
|
||||
"needs_special_handling",
|
||||
(col("is_high_value")) | (col("is_bulk_order") & col("is_premium")),
|
||||
)
|
||||
.with_column("is_north_region", col("region") == "North")
|
||||
.to_pandas()
|
||||
)
|
||||
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"product": ["A", "B", "C", "D", "E"],
|
||||
"quantity": [10, 5, 20, 15, 3],
|
||||
"price": [100, 200, 50, 75, 300],
|
||||
"region": ["North", "South", "North", "East", "West"],
|
||||
"revenue": [1000, 1000, 1000, 1125, 900],
|
||||
"is_high_value": [True, True, True, True, False],
|
||||
"is_bulk_order": [True, False, True, True, False],
|
||||
"is_premium": [True, True, False, False, True],
|
||||
"needs_special_handling": [True, True, True, True, False],
|
||||
"is_north_region": [True, False, True, False, False],
|
||||
}
|
||||
)
|
||||
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
Reference in New Issue
Block a user