wehub-resource-sync
373 lines
14 KiB
Python
373 lines
14 KiB
Python
"""Tests for comparison expression operations.
|
|
|
|
This module tests:
|
|
- Comparison operators: GT (>), LT (<), GE (>=), LE (<=), EQ (==), NE (!=)
|
|
- Comparison with columns and literals
|
|
- Reverse comparisons (literal compared to column)
|
|
"""
|
|
|
|
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, col, lit
|
|
|
|
# ──────────────────────────────────────
|
|
# Basic Comparison Operations
|
|
# ──────────────────────────────────────
|
|
|
|
|
|
class TestComparisonOperators:
|
|
"""Tests for comparison operators (>, <, >=, <=, ==, !=)."""
|
|
|
|
@pytest.fixture
|
|
def sample_data(self):
|
|
"""Sample data for comparison tests."""
|
|
return pd.DataFrame(
|
|
{
|
|
"age": [18, 21, 25, 30, 16],
|
|
"score": [50, 75, 100, 60, 85],
|
|
"status": ["active", "inactive", "active", "pending", "active"],
|
|
}
|
|
)
|
|
|
|
# ── Greater Than ──
|
|
|
|
@pytest.mark.parametrize(
|
|
"expr,expected_values",
|
|
[
|
|
(col("age") > 21, [False, False, True, True, False]),
|
|
(col("age") > col("score") / 10, [True, True, True, True, True]),
|
|
],
|
|
ids=["col_gt_literal", "col_gt_col_expr"],
|
|
)
|
|
def test_greater_than(self, sample_data, expr, expected_values):
|
|
"""Test greater than (>) comparisons."""
|
|
assert isinstance(expr, BinaryExpr)
|
|
assert expr.op == Operation.GT
|
|
result = eval_expr(expr, sample_data)
|
|
pd.testing.assert_series_equal(
|
|
result.reset_index(drop=True),
|
|
pd.Series(expected_values),
|
|
check_names=False,
|
|
)
|
|
|
|
def test_greater_than_reverse(self, sample_data):
|
|
"""Test reverse greater than (literal > col)."""
|
|
expr = 22 > col("age")
|
|
assert isinstance(expr, BinaryExpr)
|
|
assert expr.op == Operation.LT # Reverse: 22 > age becomes age < 22
|
|
result = eval_expr(expr, sample_data)
|
|
expected = pd.Series([True, True, False, False, True])
|
|
pd.testing.assert_series_equal(
|
|
result.reset_index(drop=True), expected, check_names=False
|
|
)
|
|
|
|
# ── Less Than ──
|
|
|
|
@pytest.mark.parametrize(
|
|
"expr,expected_values",
|
|
[
|
|
(col("age") < 21, [True, False, False, False, True]),
|
|
(col("score") < 70, [True, False, False, True, False]),
|
|
],
|
|
ids=["col_lt_literal", "score_lt_70"],
|
|
)
|
|
def test_less_than(self, sample_data, expr, expected_values):
|
|
"""Test less than (<) comparisons."""
|
|
assert isinstance(expr, BinaryExpr)
|
|
assert expr.op == Operation.LT
|
|
result = eval_expr(expr, sample_data)
|
|
pd.testing.assert_series_equal(
|
|
result.reset_index(drop=True),
|
|
pd.Series(expected_values),
|
|
check_names=False,
|
|
)
|
|
|
|
def test_less_than_reverse(self, sample_data):
|
|
"""Test reverse less than (literal < col)."""
|
|
expr = 20 < col("age")
|
|
assert isinstance(expr, BinaryExpr)
|
|
assert expr.op == Operation.GT # Reverse: 20 < age becomes age > 20
|
|
result = eval_expr(expr, sample_data)
|
|
expected = pd.Series([False, True, True, True, False])
|
|
pd.testing.assert_series_equal(
|
|
result.reset_index(drop=True), expected, check_names=False
|
|
)
|
|
|
|
# ── Greater Than or Equal ──
|
|
|
|
@pytest.mark.parametrize(
|
|
"expr,expected_values",
|
|
[
|
|
(col("age") >= 21, [False, True, True, True, False]),
|
|
(col("score") >= 75, [False, True, True, False, True]),
|
|
],
|
|
ids=["col_ge_21", "score_ge_75"],
|
|
)
|
|
def test_greater_equal(self, sample_data, expr, expected_values):
|
|
"""Test greater than or equal (>=) comparisons."""
|
|
assert isinstance(expr, BinaryExpr)
|
|
assert expr.op == Operation.GE
|
|
result = eval_expr(expr, sample_data)
|
|
pd.testing.assert_series_equal(
|
|
result.reset_index(drop=True),
|
|
pd.Series(expected_values),
|
|
check_names=False,
|
|
)
|
|
|
|
def test_greater_equal_reverse(self, sample_data):
|
|
"""Test reverse greater equal (literal >= col)."""
|
|
expr = 21 >= col("age")
|
|
assert isinstance(expr, BinaryExpr)
|
|
assert expr.op == Operation.LE # Reverse
|
|
result = eval_expr(expr, sample_data)
|
|
expected = pd.Series([True, True, False, False, True])
|
|
pd.testing.assert_series_equal(
|
|
result.reset_index(drop=True), expected, check_names=False
|
|
)
|
|
|
|
# ── Less Than or Equal ──
|
|
|
|
@pytest.mark.parametrize(
|
|
"expr,expected_values",
|
|
[
|
|
(col("age") <= 21, [True, True, False, False, True]),
|
|
(col("score") <= 60, [True, False, False, True, False]),
|
|
],
|
|
ids=["col_le_21", "score_le_60"],
|
|
)
|
|
def test_less_equal(self, sample_data, expr, expected_values):
|
|
"""Test less than or equal (<=) comparisons."""
|
|
assert isinstance(expr, BinaryExpr)
|
|
assert expr.op == Operation.LE
|
|
result = eval_expr(expr, sample_data)
|
|
pd.testing.assert_series_equal(
|
|
result.reset_index(drop=True),
|
|
pd.Series(expected_values),
|
|
check_names=False,
|
|
)
|
|
|
|
def test_less_equal_reverse(self, sample_data):
|
|
"""Test reverse less equal (literal <= col)."""
|
|
expr = 25 <= col("age")
|
|
assert isinstance(expr, BinaryExpr)
|
|
assert expr.op == Operation.GE # Reverse
|
|
result = eval_expr(expr, sample_data)
|
|
expected = pd.Series([False, False, True, True, False])
|
|
pd.testing.assert_series_equal(
|
|
result.reset_index(drop=True), expected, check_names=False
|
|
)
|
|
|
|
# ── Equality ──
|
|
|
|
@pytest.mark.parametrize(
|
|
"expr,expected_values",
|
|
[
|
|
(col("age") == 21, [False, True, False, False, False]),
|
|
(col("status") == "active", [True, False, True, False, True]),
|
|
(col("score") == lit(100), [False, False, True, False, False]),
|
|
],
|
|
ids=["age_eq_21", "status_eq_active", "score_eq_100"],
|
|
)
|
|
def test_equality(self, sample_data, expr, expected_values):
|
|
"""Test equality (==) comparisons."""
|
|
assert isinstance(expr, BinaryExpr)
|
|
assert expr.op == Operation.EQ
|
|
result = eval_expr(expr, sample_data)
|
|
pd.testing.assert_series_equal(
|
|
result.reset_index(drop=True),
|
|
pd.Series(expected_values),
|
|
check_names=False,
|
|
)
|
|
|
|
# ── Not Equal ──
|
|
|
|
@pytest.mark.parametrize(
|
|
"expr,expected_values",
|
|
[
|
|
(col("age") != 21, [True, False, True, True, True]),
|
|
(col("status") != "active", [False, True, False, True, False]),
|
|
],
|
|
ids=["age_ne_21", "status_ne_active"],
|
|
)
|
|
def test_not_equal(self, sample_data, expr, expected_values):
|
|
"""Test not equal (!=) comparisons."""
|
|
assert isinstance(expr, BinaryExpr)
|
|
assert expr.op == Operation.NE
|
|
result = eval_expr(expr, sample_data)
|
|
pd.testing.assert_series_equal(
|
|
result.reset_index(drop=True),
|
|
pd.Series(expected_values),
|
|
check_names=False,
|
|
)
|
|
|
|
|
|
# ──────────────────────────────────────
|
|
# Column vs Column Comparisons
|
|
# ──────────────────────────────────────
|
|
|
|
|
|
class TestColumnToColumnComparison:
|
|
"""Tests for comparing columns against other columns."""
|
|
|
|
@pytest.fixture
|
|
def sample_data(self):
|
|
"""Sample data with comparable columns."""
|
|
return pd.DataFrame(
|
|
{
|
|
"value_a": [10, 20, 30, 40],
|
|
"value_b": [15, 20, 25, 45],
|
|
"threshold": [12, 18, 35, 35],
|
|
}
|
|
)
|
|
|
|
@pytest.mark.parametrize(
|
|
"expr_fn,expected_values",
|
|
[
|
|
(lambda: col("value_a") > col("value_b"), [False, False, True, False]),
|
|
(lambda: col("value_a") < col("threshold"), [True, False, True, False]),
|
|
(lambda: col("value_a") == col("value_b"), [False, True, False, False]),
|
|
(lambda: col("value_a") >= col("threshold"), [False, True, False, True]),
|
|
],
|
|
ids=["a_gt_b", "a_lt_threshold", "a_eq_b", "a_ge_threshold"],
|
|
)
|
|
def test_column_to_column_comparisons(self, sample_data, expr_fn, expected_values):
|
|
"""Test various column-to-column comparisons."""
|
|
expr = expr_fn()
|
|
result = eval_expr(expr, sample_data)
|
|
pd.testing.assert_series_equal(
|
|
result.reset_index(drop=True),
|
|
pd.Series(expected_values),
|
|
check_names=False,
|
|
)
|
|
|
|
|
|
# ──────────────────────────────────────
|
|
# Comparison with Expressions
|
|
# ──────────────────────────────────────
|
|
|
|
|
|
class TestComparisonWithExpressions:
|
|
"""Tests for comparing expressions against other expressions."""
|
|
|
|
@pytest.fixture
|
|
def sample_data(self):
|
|
"""Sample data for expression comparison tests."""
|
|
return pd.DataFrame(
|
|
{
|
|
"price": [100, 200, 150],
|
|
"discount": [10, 50, 30],
|
|
"min_price": [80, 160, 130],
|
|
}
|
|
)
|
|
|
|
def test_compare_computed_values(self, sample_data):
|
|
"""Test comparing computed expression results."""
|
|
# (price - discount) > min_price
|
|
expr = (col("price") - col("discount")) > col("min_price")
|
|
result = eval_expr(expr, sample_data)
|
|
expected = pd.Series([True, False, False]) # 90>80, 150>160, 120>130
|
|
pd.testing.assert_series_equal(
|
|
result.reset_index(drop=True), expected, check_names=False
|
|
)
|
|
|
|
def test_compare_scaled_values(self, sample_data):
|
|
"""Test comparing scaled column values."""
|
|
# price * 0.9 >= min_price (check if 10% discount still meets minimum)
|
|
expr = col("price") * 0.9 >= col("min_price")
|
|
result = eval_expr(expr, sample_data)
|
|
expected = pd.Series([True, True, True]) # 90>=80, 180>=160, 135>=130
|
|
pd.testing.assert_series_equal(
|
|
result.reset_index(drop=True), expected, check_names=False
|
|
)
|
|
|
|
|
|
# ──────────────────────────────────────
|
|
# String Comparisons
|
|
# ──────────────────────────────────────
|
|
|
|
|
|
class TestStringComparison:
|
|
"""Tests for string equality and inequality."""
|
|
|
|
@pytest.fixture
|
|
def sample_data(self):
|
|
"""Sample data with string columns."""
|
|
return pd.DataFrame(
|
|
{
|
|
"name": ["Alice", "Bob", "Charlie", "Alice"],
|
|
"city": ["NYC", "LA", "NYC", "SF"],
|
|
}
|
|
)
|
|
|
|
def test_string_equality(self, sample_data):
|
|
"""Test string equality comparison."""
|
|
expr = col("name") == "Alice"
|
|
result = eval_expr(expr, sample_data)
|
|
expected = pd.Series([True, False, False, True])
|
|
pd.testing.assert_series_equal(
|
|
result.reset_index(drop=True), expected, check_names=False
|
|
)
|
|
|
|
def test_string_inequality(self, sample_data):
|
|
"""Test string inequality comparison."""
|
|
expr = col("city") != "NYC"
|
|
result = eval_expr(expr, sample_data)
|
|
expected = pd.Series([False, True, False, True])
|
|
pd.testing.assert_series_equal(
|
|
result.reset_index(drop=True), expected, check_names=False
|
|
)
|
|
|
|
|
|
# ──────────────────────────────────────
|
|
# Boolean Comparisons
|
|
# ──────────────────────────────────────
|
|
|
|
|
|
class TestBooleanComparison:
|
|
"""Tests for boolean value comparisons."""
|
|
|
|
@pytest.fixture
|
|
def sample_data(self):
|
|
"""Sample data with boolean columns."""
|
|
return pd.DataFrame(
|
|
{
|
|
"is_active": [True, False, True, False],
|
|
"is_verified": [True, True, False, False],
|
|
}
|
|
)
|
|
|
|
def test_boolean_equality_true(self, sample_data):
|
|
"""Test boolean equality with True."""
|
|
expr = col("is_active") == lit(True)
|
|
result = eval_expr(expr, sample_data)
|
|
expected = pd.Series([True, False, True, False])
|
|
pd.testing.assert_series_equal(
|
|
result.reset_index(drop=True), expected, check_names=False
|
|
)
|
|
|
|
def test_boolean_equality_false(self, sample_data):
|
|
"""Test boolean equality with False."""
|
|
expr = col("is_active") == lit(False)
|
|
result = eval_expr(expr, sample_data)
|
|
expected = pd.Series([False, True, False, True])
|
|
pd.testing.assert_series_equal(
|
|
result.reset_index(drop=True), expected, check_names=False
|
|
)
|
|
|
|
def test_boolean_column_to_column(self, sample_data):
|
|
"""Test comparing two boolean columns."""
|
|
expr = col("is_active") == col("is_verified")
|
|
result = eval_expr(expr, sample_data)
|
|
expected = pd.Series([True, False, False, True])
|
|
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__]))
|