353 lines
13 KiB
Python
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__]))
|