Files
ray-project--ray/python/ray/data/tests/unit/expressions/test_predicate.py
T
2026-07-13 13:17:40 +08:00

353 lines
13 KiB
Python

"""Tests for predicate expression operations.
This module tests:
- Null predicates: is_null(), is_not_null()
- Membership predicates: is_in(), not_in()
"""
import pandas as pd
import pytest
from ray.data._internal.planner.plan_expression.expression_evaluator import eval_expr
from ray.data.expressions import BinaryExpr, Operation, UnaryExpr, col, lit
# ──────────────────────────────────────
# Null Predicate Operations
# ──────────────────────────────────────
class TestIsNull:
"""Tests for is_null() predicate."""
@pytest.fixture
def sample_data(self):
"""Sample data with null values for null predicate tests."""
return pd.DataFrame(
{
"value": [1.0, None, 3.0, None, 5.0],
"name": ["Alice", None, "Charlie", "Diana", None],
}
)
def test_is_null_numeric(self, sample_data):
"""Test is_null on numeric column."""
expr = col("value").is_null()
assert isinstance(expr, UnaryExpr)
assert expr.op == Operation.IS_NULL
result = eval_expr(expr, sample_data)
expected = pd.Series([False, True, False, True, False])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
def test_is_null_string(self, sample_data):
"""Test is_null on string column."""
expr = col("name").is_null()
result = eval_expr(expr, sample_data)
expected = pd.Series([False, True, False, False, True])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
def test_is_null_structural_equality(self):
"""Test structural equality for is_null expressions."""
expr1 = col("value").is_null()
expr2 = col("value").is_null()
expr3 = col("other").is_null()
assert expr1.structurally_equals(expr2)
assert not expr1.structurally_equals(expr3)
class TestIsNotNull:
"""Tests for is_not_null() predicate."""
@pytest.fixture
def sample_data(self):
"""Sample data with null values."""
return pd.DataFrame(
{
"value": [1.0, None, 3.0, None, 5.0],
"name": ["Alice", None, "Charlie", "Diana", None],
}
)
def test_is_not_null_numeric(self, sample_data):
"""Test is_not_null on numeric column."""
expr = col("value").is_not_null()
assert isinstance(expr, UnaryExpr)
assert expr.op == Operation.IS_NOT_NULL
result = eval_expr(expr, sample_data)
expected = pd.Series([True, False, True, False, True])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
def test_is_not_null_string(self, sample_data):
"""Test is_not_null on string column."""
expr = col("name").is_not_null()
result = eval_expr(expr, sample_data)
expected = pd.Series([True, False, True, True, False])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
def test_is_not_null_structural_equality(self):
"""Test structural equality for is_not_null expressions."""
expr1 = col("value").is_not_null()
expr2 = col("value").is_not_null()
expr3 = col("other").is_not_null()
assert expr1.structurally_equals(expr2)
assert not expr1.structurally_equals(expr3)
class TestNullPredicateCombinations:
"""Tests for null predicates combined with other operations."""
@pytest.fixture
def sample_data(self):
"""Sample data with null values and other columns."""
return pd.DataFrame(
{
"value": [10.0, None, 30.0, None, 50.0],
"threshold": [5.0, 20.0, 25.0, 10.0, 40.0],
}
)
def test_null_aware_comparison(self, sample_data):
"""Test null-aware comparison (is_not_null AND comparison)."""
# Filter: value is not null AND value > threshold
expr = col("value").is_not_null() & (col("value") > col("threshold"))
result = eval_expr(expr, sample_data)
expected = pd.Series([True, False, True, False, True])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
def test_is_null_or_condition(self, sample_data):
"""Test is_null combined with OR."""
# value is null OR value > 40
expr = col("value").is_null() | (col("value") > 40)
result = eval_expr(expr, sample_data)
expected = pd.Series([False, True, False, True, True])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
# ──────────────────────────────────────
# Membership Predicate Operations
# ──────────────────────────────────────
class TestIsIn:
"""Tests for is_in() predicate."""
@pytest.fixture
def sample_data(self):
"""Sample data for membership tests."""
return pd.DataFrame(
{
"status": ["active", "inactive", "pending", "active", "deleted"],
"category": ["A", "B", "C", "A", "D"],
"value": [1, 2, 3, 4, 5],
}
)
def test_is_in_string_list(self, sample_data):
"""Test is_in with string list."""
expr = col("status").is_in(["active", "pending"])
assert isinstance(expr, BinaryExpr)
assert expr.op == Operation.IN
result = eval_expr(expr, sample_data)
expected = pd.Series([True, False, True, True, False])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
def test_is_in_single_value_list(self, sample_data):
"""Test is_in with single-value list."""
expr = col("status").is_in(["active"])
result = eval_expr(expr, sample_data)
expected = pd.Series([True, False, False, True, False])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
def test_is_in_numeric_list(self, sample_data):
"""Test is_in with numeric list."""
expr = col("value").is_in([1, 3, 5])
result = eval_expr(expr, sample_data)
expected = pd.Series([True, False, True, False, True])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
def test_is_in_empty_list(self, sample_data):
"""Test is_in with empty list (should return all False)."""
expr = col("status").is_in([])
result = eval_expr(expr, sample_data)
expected = pd.Series([False, False, False, False, False])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
def test_is_in_with_literal_expr(self, sample_data):
"""Test is_in with LiteralExpr containing list."""
values_expr = lit(["A", "C"])
expr = col("category").is_in(values_expr)
result = eval_expr(expr, sample_data)
expected = pd.Series([True, False, True, True, False])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
def test_is_in_structural_equality(self):
"""Test structural equality for is_in expressions."""
expr1 = col("status").is_in(["active", "pending"])
expr2 = col("status").is_in(["active", "pending"])
expr3 = col("status").is_in(["active"])
assert expr1.structurally_equals(expr2)
assert not expr1.structurally_equals(expr3)
class TestNotIn:
"""Tests for not_in() predicate."""
@pytest.fixture
def sample_data(self):
"""Sample data for membership tests."""
return pd.DataFrame(
{
"status": ["active", "inactive", "pending", "active", "deleted"],
"value": [1, 2, 3, 4, 5],
}
)
def test_not_in_string_list(self, sample_data):
"""Test not_in with string list."""
expr = col("status").not_in(["inactive", "deleted"])
assert isinstance(expr, BinaryExpr)
assert expr.op == Operation.NOT_IN
result = eval_expr(expr, sample_data)
expected = pd.Series([True, False, True, True, False])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
def test_not_in_numeric_list(self, sample_data):
"""Test not_in with numeric list."""
expr = col("value").not_in([2, 4])
result = eval_expr(expr, sample_data)
expected = pd.Series([True, False, True, False, True])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
def test_not_in_empty_list(self, sample_data):
"""Test not_in with empty list (should return all True)."""
expr = col("status").not_in([])
result = eval_expr(expr, sample_data)
expected = pd.Series([True, True, True, True, True])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
def test_not_in_structural_equality(self):
"""Test structural equality for not_in expressions."""
expr1 = col("status").not_in(["deleted"])
expr2 = col("status").not_in(["deleted"])
expr3 = col("status").not_in(["deleted", "inactive"])
assert expr1.structurally_equals(expr2)
assert not expr1.structurally_equals(expr3)
class TestMembershipWithNulls:
"""Tests for membership predicates with null values."""
@pytest.fixture
def sample_data(self):
"""Sample data with null values for membership tests."""
return pd.DataFrame(
{
"status": ["active", None, "pending", None, "deleted"],
"value": [1, None, 3, None, 5],
}
)
def test_is_in_with_nulls_in_data(self, sample_data):
"""Test is_in when data contains nulls."""
expr = col("status").is_in(["active", "pending"])
result = eval_expr(expr, sample_data)
# Nulls should return False (null is not in any list)
expected = pd.Series([True, False, True, False, False])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
def test_not_in_with_nulls_in_data(self, sample_data):
"""Test not_in when data contains nulls."""
expr = col("status").not_in(["active"])
result = eval_expr(expr, sample_data)
# Nulls should return True (null is not in the exclusion list)
expected = pd.Series([False, True, True, True, True])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
class TestMembershipCombinations:
"""Tests for membership predicates combined with other operations."""
@pytest.fixture
def sample_data(self):
"""Sample data for combination tests."""
return pd.DataFrame(
{
"status": ["active", "inactive", "pending", "active", "deleted"],
"priority": ["high", "low", "high", "medium", "low"],
"value": [100, 50, 75, 200, 25],
}
)
def test_is_in_and_comparison(self, sample_data):
"""Test is_in combined with comparison."""
# status in ["active", "pending"] AND value > 50
expr = col("status").is_in(["active", "pending"]) & (col("value") > 50)
result = eval_expr(expr, sample_data)
expected = pd.Series([True, False, True, True, False])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
def test_multiple_is_in(self, sample_data):
"""Test multiple is_in predicates."""
# status in ["active"] AND priority in ["high", "medium"]
expr = col("status").is_in(["active"]) & col("priority").is_in(
["high", "medium"]
)
result = eval_expr(expr, sample_data)
expected = pd.Series([True, False, False, True, False])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
def test_is_in_or_not_in(self, sample_data):
"""Test is_in combined with not_in."""
# status in ["active"] OR priority not_in ["low"]
expr = col("status").is_in(["active"]) | col("priority").not_in(["low"])
result = eval_expr(expr, sample_data)
expected = pd.Series([True, False, True, True, False])
pd.testing.assert_series_equal(
result.reset_index(drop=True), expected, check_names=False
)
if __name__ == "__main__":
import sys
sys.exit(pytest.main(["-v", __file__]))