wehub-resource-sync
629 lines
21 KiB
Python
629 lines
21 KiB
Python
"""Tests for expression conversion to PyArrow and Iceberg.
|
|
|
|
This module tests:
|
|
- Conversion to PyArrow compute expressions (to_pyarrow)
|
|
- Conversion to PyIceberg expressions (IcebergExpressionVisitor)
|
|
- Unsupported expression handling
|
|
"""
|
|
|
|
import pyarrow as pa
|
|
import pyarrow.compute as pc
|
|
import pyarrow.dataset as ds
|
|
import pytest
|
|
from packaging.version import parse as version_parse
|
|
from pyiceberg.expressions import (
|
|
And,
|
|
EqualTo,
|
|
GreaterThan,
|
|
GreaterThanOrEqual,
|
|
In,
|
|
IsNull,
|
|
LessThan,
|
|
LessThanOrEqual,
|
|
Not,
|
|
NotEqualTo,
|
|
NotIn,
|
|
NotNull,
|
|
Or,
|
|
Reference,
|
|
literal,
|
|
)
|
|
|
|
from ray.data._internal.datasource.iceberg_datasource import _IcebergExpressionVisitor
|
|
from ray.data.datatype import DataType
|
|
from ray.data.expressions import (
|
|
BinaryExpr,
|
|
Operation,
|
|
UDFExpr,
|
|
col,
|
|
download,
|
|
lit,
|
|
star,
|
|
)
|
|
|
|
# ──────────────────────────────────────
|
|
# PyArrow Conversion Tests
|
|
# ──────────────────────────────────────
|
|
|
|
|
|
class TestToPyArrow:
|
|
"""Test conversion of Ray Data expressions to PyArrow compute expressions."""
|
|
|
|
@pytest.fixture
|
|
def test_table(self):
|
|
"""Sample PyArrow table for testing expressions."""
|
|
return pa.table(
|
|
{
|
|
"age": [15, 25, 45, 70],
|
|
"x": [1, 2, 3, 4],
|
|
"price": [10.0, 20.0, 30.0, 40.0],
|
|
"quantity": [2, 3, 1, 5],
|
|
"tax": [1.0, 2.0, 3.0, 4.0],
|
|
"status": ["active", "pending", "inactive", "active"],
|
|
"value": [1, None, 3, None],
|
|
"active": [True, False, True, False],
|
|
}
|
|
)
|
|
|
|
# ── Basic Expressions ──
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_expr,equivalent_pyarrow_expr",
|
|
[
|
|
(col("age"), lambda: pc.field("age")),
|
|
(lit(42), lambda: pc.scalar(42)),
|
|
(lit("hello"), lambda: pc.scalar("hello")),
|
|
],
|
|
ids=["col", "int_lit", "str_lit"],
|
|
)
|
|
def test_basic_expressions(self, test_table, ray_expr, equivalent_pyarrow_expr):
|
|
"""Test conversion of basic expressions."""
|
|
converted = ray_expr.to_pyarrow()
|
|
expected = equivalent_pyarrow_expr()
|
|
assert converted.equals(expected)
|
|
|
|
# ── Arithmetic Expressions ──
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_expr,equivalent_pyarrow_expr",
|
|
[
|
|
(
|
|
col("x") + 5,
|
|
lambda: pc.add(pc.field("x"), pc.scalar(5)),
|
|
),
|
|
(
|
|
col("x") - 3,
|
|
lambda: pc.subtract(pc.field("x"), pc.scalar(3)),
|
|
),
|
|
(
|
|
col("x") * 2,
|
|
lambda: pc.multiply(pc.field("x"), pc.scalar(2)),
|
|
),
|
|
(
|
|
col("x") / 2,
|
|
lambda: pc.divide(pc.field("x"), pc.scalar(2)),
|
|
),
|
|
],
|
|
ids=["add", "sub", "mul", "div"],
|
|
)
|
|
def test_arithmetic_expressions(
|
|
self, test_table, ray_expr, equivalent_pyarrow_expr
|
|
):
|
|
"""Test conversion of arithmetic expressions."""
|
|
converted = ray_expr.to_pyarrow()
|
|
expected = equivalent_pyarrow_expr()
|
|
assert converted.equals(expected)
|
|
|
|
# ── Comparison Expressions ──
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_expr,equivalent_pyarrow_expr,description",
|
|
[
|
|
(
|
|
col("age") > 18,
|
|
lambda: pc.greater(pc.field("age"), pc.scalar(18)),
|
|
"greater than",
|
|
),
|
|
(
|
|
col("age") < 65,
|
|
lambda: pc.less(pc.field("age"), pc.scalar(65)),
|
|
"less than",
|
|
),
|
|
(
|
|
col("age") >= 21,
|
|
lambda: pc.greater_equal(pc.field("age"), pc.scalar(21)),
|
|
"greater equal",
|
|
),
|
|
(
|
|
col("age") <= 30,
|
|
lambda: pc.less_equal(pc.field("age"), pc.scalar(30)),
|
|
"less equal",
|
|
),
|
|
(
|
|
col("status") == "active",
|
|
lambda: pc.equal(pc.field("status"), pc.scalar("active")),
|
|
"equality",
|
|
),
|
|
(
|
|
col("status") != "deleted",
|
|
lambda: pc.not_equal(pc.field("status"), pc.scalar("deleted")),
|
|
"not equal",
|
|
),
|
|
],
|
|
ids=["gt", "lt", "ge", "le", "eq", "ne"],
|
|
)
|
|
def test_comparison_expressions(
|
|
self, test_table, ray_expr, equivalent_pyarrow_expr, description
|
|
):
|
|
"""Test conversion of comparison expressions."""
|
|
converted = ray_expr.to_pyarrow()
|
|
expected = equivalent_pyarrow_expr()
|
|
|
|
# Verify they produce the same results on sample data
|
|
dataset = ds.dataset(test_table)
|
|
try:
|
|
result_converted = dataset.scanner(filter=converted).to_table()
|
|
result_expected = dataset.scanner(filter=expected).to_table()
|
|
assert result_converted.equals(result_expected)
|
|
except (TypeError, pa.lib.ArrowInvalid, pa.lib.ArrowNotImplementedError):
|
|
pass
|
|
|
|
# ── Boolean Expressions ──
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_expr,equivalent_pyarrow_expr,description",
|
|
[
|
|
(
|
|
(col("age") > 18) & (col("age") < 65),
|
|
lambda: pc.and_kleene(
|
|
pc.greater(pc.field("age"), pc.scalar(18)),
|
|
pc.less(pc.field("age"), pc.scalar(65)),
|
|
),
|
|
"logical AND",
|
|
),
|
|
(
|
|
(col("status") == "active") | (col("status") == "pending"),
|
|
lambda: pc.or_kleene(
|
|
pc.equal(pc.field("status"), pc.scalar("active")),
|
|
pc.equal(pc.field("status"), pc.scalar("pending")),
|
|
),
|
|
"logical OR",
|
|
),
|
|
(
|
|
~col("active"),
|
|
lambda: pc.invert(pc.field("active")),
|
|
"logical NOT",
|
|
),
|
|
],
|
|
ids=["and", "or", "not"],
|
|
)
|
|
def test_boolean_expressions(
|
|
self, test_table, ray_expr, equivalent_pyarrow_expr, description
|
|
):
|
|
"""Test conversion of boolean expressions."""
|
|
converted = ray_expr.to_pyarrow()
|
|
expected = equivalent_pyarrow_expr()
|
|
|
|
dataset = ds.dataset(test_table)
|
|
try:
|
|
result_converted = dataset.scanner(filter=converted).to_table()
|
|
result_expected = dataset.scanner(filter=expected).to_table()
|
|
assert result_converted.equals(result_expected)
|
|
except (TypeError, pa.lib.ArrowInvalid, pa.lib.ArrowNotImplementedError):
|
|
pass
|
|
|
|
# ── Predicate Expressions ──
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_expr,equivalent_pyarrow_expr,description",
|
|
[
|
|
(
|
|
col("value").is_null(),
|
|
lambda: pc.is_null(pc.field("value")),
|
|
"is_null check",
|
|
),
|
|
(
|
|
col("value").is_not_null(),
|
|
lambda: pc.is_valid(pc.field("value")),
|
|
"is_not_null check",
|
|
),
|
|
(
|
|
col("status").is_in(["active", "pending"]),
|
|
lambda: pc.is_in(pc.field("status"), pa.array(["active", "pending"])),
|
|
"is_in with list",
|
|
),
|
|
],
|
|
ids=["is_null", "is_not_null", "is_in"],
|
|
)
|
|
def test_predicate_expressions(
|
|
self, test_table, ray_expr, equivalent_pyarrow_expr, description
|
|
):
|
|
"""Test conversion of predicate expressions."""
|
|
converted = ray_expr.to_pyarrow()
|
|
expected = equivalent_pyarrow_expr()
|
|
|
|
dataset = ds.dataset(test_table)
|
|
try:
|
|
result_converted = dataset.scanner(filter=converted).to_table()
|
|
result_expected = dataset.scanner(filter=expected).to_table()
|
|
assert result_converted.equals(result_expected)
|
|
except (TypeError, pa.lib.ArrowInvalid, pa.lib.ArrowNotImplementedError):
|
|
pass
|
|
|
|
# ── Nested Expressions ──
|
|
|
|
def test_nested_arithmetic(self, test_table):
|
|
"""Test nested arithmetic expressions."""
|
|
ray_expr = (col("price") * col("quantity")) + col("tax")
|
|
converted = ray_expr.to_pyarrow()
|
|
assert isinstance(converted, pc.Expression)
|
|
|
|
# ── Alias Expressions ──
|
|
|
|
def test_alias_expressions(self, test_table):
|
|
"""Test that alias expressions unwrap to inner expression."""
|
|
ray_expr = (col("x") + 5).alias("result")
|
|
converted = ray_expr.to_pyarrow()
|
|
assert isinstance(converted, pc.Expression)
|
|
|
|
# ── PyArrow Compute UDF Expressions ──
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_expr,equivalent_pyarrow_expr",
|
|
[
|
|
pytest.param(
|
|
col("name").str.match_regex("foo.*bar"),
|
|
lambda: pc.match_substring_regex(pc.field("name"), "foo.*bar"),
|
|
id="match_regex",
|
|
),
|
|
pytest.param(
|
|
col("name").str.starts_with("foo"),
|
|
lambda: pc.starts_with(pc.field("name"), "foo"),
|
|
id="starts_with",
|
|
),
|
|
pytest.param(
|
|
col("name").str.ends_with("bar"),
|
|
lambda: pc.ends_with(pc.field("name"), "bar"),
|
|
id="ends_with",
|
|
),
|
|
pytest.param(
|
|
col("name").str.contains("baz"),
|
|
lambda: pc.match_substring(pc.field("name"), "baz"),
|
|
id="contains",
|
|
),
|
|
pytest.param(
|
|
col("name").str.upper(),
|
|
lambda: pc.utf8_upper(pc.field("name")),
|
|
id="upper",
|
|
),
|
|
pytest.param(
|
|
col("x").ceil(),
|
|
lambda: pc.ceil(pc.field("x")),
|
|
id="ceil",
|
|
),
|
|
pytest.param(
|
|
col("x").abs(),
|
|
lambda: pc.abs_checked(pc.field("x")),
|
|
id="abs",
|
|
),
|
|
],
|
|
)
|
|
def test_pyarrow_compute_udf_expressions(
|
|
self, test_table, ray_expr, equivalent_pyarrow_expr
|
|
):
|
|
"""Test that PyArrow-compute-backed UDFs convert to PyArrow expressions."""
|
|
converted = ray_expr.to_pyarrow()
|
|
expected = equivalent_pyarrow_expr()
|
|
assert converted.equals(expected)
|
|
|
|
@pytest.mark.skipif(
|
|
version_parse(pa.__version__) < version_parse("19.0.0"),
|
|
reason="Requires PyArrow >= 19 for string compute UDF pushdown",
|
|
)
|
|
def test_negated_pyarrow_compute_udf(self, test_table):
|
|
"""Test that negated PyArrow compute UDF expressions convert correctly."""
|
|
ray_expr = ~col("status").str.match_regex("act.*")
|
|
converted = ray_expr.to_pyarrow()
|
|
assert isinstance(converted, pc.Expression)
|
|
|
|
dataset = ds.dataset(test_table)
|
|
result = dataset.to_table(filter=converted)
|
|
assert all(
|
|
not bool(pc.match_substring_regex(s, "act.*"))
|
|
for s in result.column("status").to_pylist()
|
|
)
|
|
|
|
def test_pyarrow_compute_udf_as_dataset_filter(self, test_table):
|
|
"""Test that converted UDF expressions work as dataset scan filters."""
|
|
ray_expr = col("status").str.match_regex("^active$")
|
|
pa_expr = ray_expr.to_pyarrow()
|
|
|
|
dataset = ds.dataset(test_table)
|
|
result = dataset.to_table(filter=pa_expr)
|
|
assert all(s == "active" for s in result.column("status").to_pylist())
|
|
|
|
# ── Unsupported Expressions ──
|
|
|
|
def test_user_udf_expression_raises(self):
|
|
"""Test that user-defined UDF expressions raise TypeError."""
|
|
|
|
def dummy_fn(x):
|
|
return x
|
|
|
|
udf_expr = UDFExpr(
|
|
fn=dummy_fn,
|
|
args=[col("x")],
|
|
kwargs={},
|
|
data_type=DataType(int),
|
|
)
|
|
|
|
with pytest.raises(TypeError, match="UDF expressions cannot be converted"):
|
|
udf_expr.to_pyarrow()
|
|
|
|
def test_download_expression_raises(self):
|
|
"""Test that download expressions raise TypeError."""
|
|
with pytest.raises(TypeError, match="Download expressions cannot be converted"):
|
|
download("uri").to_pyarrow()
|
|
|
|
def test_star_expression_raises(self):
|
|
"""Test that star expressions raise TypeError."""
|
|
with pytest.raises(TypeError, match="Star expressions cannot be converted"):
|
|
star().to_pyarrow()
|
|
|
|
|
|
# ──────────────────────────────────────
|
|
# Iceberg Conversion Tests
|
|
# ──────────────────────────────────────
|
|
|
|
|
|
class TestIcebergExpressionVisitor:
|
|
"""Test conversion of Ray Data expressions to PyIceberg expressions."""
|
|
|
|
# ── Basic Expressions ──
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_expr,equivalent_iceberg_expr,description",
|
|
[
|
|
(col("age"), lambda: Reference("age"), "column reference"),
|
|
(lit(42), lambda: literal(42), "integer literal"),
|
|
(lit("active"), lambda: literal("active"), "string literal"),
|
|
],
|
|
ids=["col", "int_lit", "str_lit"],
|
|
)
|
|
def test_basic_expressions(self, ray_expr, equivalent_iceberg_expr, description):
|
|
"""Test conversion of basic expressions."""
|
|
visitor = _IcebergExpressionVisitor()
|
|
converted = visitor.visit(ray_expr)
|
|
expected = equivalent_iceberg_expr()
|
|
assert converted == expected
|
|
|
|
# ── Comparison Expressions ──
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_expr,equivalent_iceberg_expr,description",
|
|
[
|
|
(
|
|
col("age") > 18,
|
|
lambda: GreaterThan(Reference("age"), literal(18)),
|
|
"greater than",
|
|
),
|
|
(
|
|
col("age") >= 21,
|
|
lambda: GreaterThanOrEqual(Reference("age"), literal(21)),
|
|
"greater than or equal",
|
|
),
|
|
(
|
|
col("age") < 65,
|
|
lambda: LessThan(Reference("age"), literal(65)),
|
|
"less than",
|
|
),
|
|
(
|
|
col("age") <= 100,
|
|
lambda: LessThanOrEqual(Reference("age"), literal(100)),
|
|
"less than or equal",
|
|
),
|
|
(
|
|
col("status") == "active",
|
|
lambda: EqualTo(Reference("status"), literal("active")),
|
|
"equality",
|
|
),
|
|
(
|
|
col("status") != "inactive",
|
|
lambda: NotEqualTo(Reference("status"), literal("inactive")),
|
|
"not equal",
|
|
),
|
|
],
|
|
ids=["gt", "ge", "lt", "le", "eq", "ne"],
|
|
)
|
|
def test_comparison_expressions(
|
|
self, ray_expr, equivalent_iceberg_expr, description
|
|
):
|
|
"""Test conversion of comparison expressions."""
|
|
visitor = _IcebergExpressionVisitor()
|
|
converted = visitor.visit(ray_expr)
|
|
expected = equivalent_iceberg_expr()
|
|
assert converted == expected
|
|
|
|
# ── Boolean Expressions ──
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_expr,equivalent_iceberg_expr,description",
|
|
[
|
|
(
|
|
(col("age") > 18) & (col("age") < 65),
|
|
lambda: And(
|
|
GreaterThan(Reference("age"), literal(18)),
|
|
LessThan(Reference("age"), literal(65)),
|
|
),
|
|
"logical AND",
|
|
),
|
|
(
|
|
(col("is_member") == lit(True)) | (col("is_premium") == lit(True)),
|
|
lambda: Or(
|
|
EqualTo(Reference("is_member"), literal(True)),
|
|
EqualTo(Reference("is_premium"), literal(True)),
|
|
),
|
|
"logical OR",
|
|
),
|
|
(
|
|
~(col("deleted") == lit(True)),
|
|
lambda: Not(EqualTo(Reference("deleted"), literal(True))),
|
|
"logical NOT",
|
|
),
|
|
],
|
|
ids=["and", "or", "not"],
|
|
)
|
|
def test_boolean_expressions(self, ray_expr, equivalent_iceberg_expr, description):
|
|
"""Test conversion of boolean expressions."""
|
|
visitor = _IcebergExpressionVisitor()
|
|
converted = visitor.visit(ray_expr)
|
|
expected = equivalent_iceberg_expr()
|
|
assert converted == expected
|
|
|
|
# ── Predicate Expressions ──
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_expr,equivalent_iceberg_expr,description",
|
|
[
|
|
(
|
|
col("value").is_null(),
|
|
lambda: IsNull(Reference("value")),
|
|
"is_null check",
|
|
),
|
|
(
|
|
col("name").is_not_null(),
|
|
lambda: NotNull(Reference("name")),
|
|
"is_not_null check",
|
|
),
|
|
(
|
|
col("status").is_in(["active", "pending"]),
|
|
lambda: In(Reference("status"), ["active", "pending"]),
|
|
"is_in with list",
|
|
),
|
|
(
|
|
col("status").not_in(["inactive", "deleted"]),
|
|
lambda: NotIn(Reference("status"), ["inactive", "deleted"]),
|
|
"not_in with list",
|
|
),
|
|
],
|
|
ids=["is_null", "is_not_null", "is_in", "not_in"],
|
|
)
|
|
def test_predicate_expressions(
|
|
self, ray_expr, equivalent_iceberg_expr, description
|
|
):
|
|
"""Test conversion of predicate expressions."""
|
|
visitor = _IcebergExpressionVisitor()
|
|
converted = visitor.visit(ray_expr)
|
|
expected = equivalent_iceberg_expr()
|
|
assert converted == expected
|
|
|
|
# ── Complex Nested Expressions ──
|
|
|
|
def test_complex_nested_boolean(self):
|
|
"""Test complex nested boolean expression."""
|
|
ray_expr = (
|
|
(col("age") >= 21)
|
|
& (col("country") == "USA")
|
|
& col("verified").is_not_null()
|
|
)
|
|
expected = And(
|
|
And(
|
|
GreaterThanOrEqual(Reference("age"), literal(21)),
|
|
EqualTo(Reference("country"), literal("USA")),
|
|
),
|
|
NotNull(Reference("verified")),
|
|
)
|
|
|
|
visitor = _IcebergExpressionVisitor()
|
|
converted = visitor.visit(ray_expr)
|
|
assert converted == expected
|
|
|
|
def test_aliased_expression(self):
|
|
"""Test that alias expressions unwrap to inner expression."""
|
|
ray_expr = (col("age") > 18).alias("is_adult")
|
|
expected = GreaterThan(Reference("age"), literal(18))
|
|
|
|
visitor = _IcebergExpressionVisitor()
|
|
converted = visitor.visit(ray_expr)
|
|
assert converted == expected
|
|
|
|
# ── Unsupported Arithmetic ──
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_expr",
|
|
[
|
|
col("price") + 10,
|
|
col("quantity") * 2,
|
|
col("total") - col("discount"),
|
|
col("revenue") / col("count"),
|
|
col("items") // 5,
|
|
],
|
|
ids=["add", "mul", "sub", "div", "floordiv"],
|
|
)
|
|
def test_arithmetic_raises(self, ray_expr):
|
|
"""Test that arithmetic operations raise appropriate errors."""
|
|
visitor = _IcebergExpressionVisitor()
|
|
with pytest.raises(
|
|
ValueError, match="Unsupported binary operation for Iceberg"
|
|
):
|
|
visitor.visit(ray_expr)
|
|
|
|
# ── Unsupported Expressions ──
|
|
|
|
def test_udf_expression_raises(self):
|
|
"""Test that UDF expressions raise TypeError."""
|
|
|
|
def dummy_fn(x):
|
|
return x
|
|
|
|
udf_expr = UDFExpr(
|
|
fn=dummy_fn,
|
|
args=[col("x")],
|
|
kwargs={},
|
|
data_type=DataType(int),
|
|
)
|
|
|
|
visitor = _IcebergExpressionVisitor()
|
|
with pytest.raises(
|
|
TypeError, match="UDF expressions cannot be converted to Iceberg"
|
|
):
|
|
visitor.visit(udf_expr)
|
|
|
|
def test_download_expression_raises(self):
|
|
"""Test that download expressions raise TypeError."""
|
|
visitor = _IcebergExpressionVisitor()
|
|
with pytest.raises(
|
|
TypeError, match="Download expressions cannot be converted to Iceberg"
|
|
):
|
|
visitor.visit(download("uri"))
|
|
|
|
def test_star_expression_raises(self):
|
|
"""Test that star expressions raise TypeError."""
|
|
visitor = _IcebergExpressionVisitor()
|
|
with pytest.raises(
|
|
TypeError, match="Star expressions cannot be converted to Iceberg"
|
|
):
|
|
visitor.visit(star())
|
|
|
|
def test_is_in_requires_literal_list(self):
|
|
"""Test that IN/NOT_IN operations require literal lists."""
|
|
visitor = _IcebergExpressionVisitor()
|
|
|
|
# This should work - literal list
|
|
expr = col("status").is_in(["active", "pending"])
|
|
result = visitor.visit(expr)
|
|
assert isinstance(result, In)
|
|
|
|
# This should fail - column reference on right side
|
|
with pytest.raises(
|
|
ValueError, match="IN operation requires right operand to be a literal list"
|
|
):
|
|
invalid_expr = BinaryExpr(Operation.IN, col("a"), col("b"))
|
|
visitor.visit(invalid_expr)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|