Add CAST expression support to parser

Author: nssalian

pyiceberg/expressions/parser.py

@@ -16,6 +16,7 @@
 #  under the License.
 import re
 from decimal import Decimal
+from typing import Any
 
 from pyparsing import (
     CaselessKeyword,
@@ -66,6 +67,14 @@
     LongLiteral,
     StringLiteral,
 )
+from pyiceberg.transforms import (
+    DayTransform,
+    HourTransform,
+    MonthTransform,
+    Transform,
+    UnboundTransform,
+    YearTransform,
+)
 from pyiceberg.typedef import L
 from pyiceberg.types import strtobool
 
@@ -80,6 +89,8 @@
 NAN = CaselessKeyword("nan")
 LIKE = CaselessKeyword("like")
 BETWEEN = CaselessKeyword("between")
+CAST = CaselessKeyword("cast")
+AS = CaselessKeyword("as")
 
 unquoted_identifier = Word(alphas + "_", alphanums + "_$")
 quoted_identifier = QuotedString('"', esc_quote="\\", unquote_results=True)
@@ -265,7 +276,51 @@ def _evaluate_like_statement(result: ParseResults) -> BooleanExpression:
         return EqualTo(result.column, StringLiteral(literal_like.value.replace("\\%", "%")))
 
 
-predicate = (between | comparison | in_check | null_check | nan_check | starts_check | boolean).set_results_name("predicate")
+# CAST expression support: CAST(column AS type) maps to Iceberg transforms
+_CAST_TYPE_TO_TRANSFORM: dict[str, Transform[Any, Any]] = {
+    "date": DayTransform(),
+    "year": YearTransform(),
+    "month": MonthTransform(),
+    "hour": HourTransform(),
+}
+
+cast_type = Word(alphas, alphanums + "_")
+cast_term = Suppress(CAST) + Suppress("(") + column + Suppress(AS) + cast_type + Suppress(")")
+
+
+@cast_term.set_parse_action
+def _(result: ParseResults) -> UnboundTransform:
+    ref = result[0]
+    target = str(result[1]).lower()
+    if target not in _CAST_TYPE_TO_TRANSFORM:
+        raise ValueError(f"Unsupported CAST target type: {target}")
+    return UnboundTransform(ref, _CAST_TYPE_TO_TRANSFORM[target])
+
+
+cast_left_ref = cast_term + comparison_op + literal
+
+
+@cast_left_ref.set_parse_action
+def _(result: ParseResults) -> BooleanExpression:
+    term = result[0]
+    if result.op == "<":
+        return LessThan(term, result.literal)
+    elif result.op == "<=":
+        return LessThanOrEqual(term, result.literal)
+    elif result.op == ">":
+        return GreaterThan(term, result.literal)
+    elif result.op == ">=":
+        return GreaterThanOrEqual(term, result.literal)
+    if result.op in ("=", "=="):
+        return EqualTo(term, result.literal)
+    if result.op in ("!=", "<>"):
+        return NotEqualTo(term, result.literal)
+    raise ValueError(f"Unsupported operation type: {result.op}")
+
+
+predicate = (between | cast_left_ref | comparison | in_check | null_check | nan_check | starts_check | boolean).set_results_name(
+    "predicate"
+)
 
 
 def handle_not(result: ParseResults) -> Not:

pyiceberg/transforms.py

@@ -44,6 +44,7 @@
     BoundNotIn,
     BoundNotStartsWith,
     BoundPredicate,
+    BoundReference,
     BoundSetPredicate,
     BoundStartsWith,
     BoundTerm,
@@ -58,6 +59,7 @@
     Reference,
     StartsWith,
     UnboundPredicate,
+    UnboundTerm,
 )
 from pyiceberg.expressions.literals import (
     DateLiteral,
@@ -67,7 +69,7 @@
     TimestampLiteral,
     literal,
 )
-from pyiceberg.typedef import IcebergRootModel, L
+from pyiceberg.typedef import IcebergRootModel, L, StructProtocol
 from pyiceberg.types import (
     BinaryType,
     DateType,
@@ -92,6 +94,8 @@
 if TYPE_CHECKING:
     import pyarrow as pa
 
+    from pyiceberg.schema import Schema
+
     ArrayLike = TypeVar("ArrayLike", pa.Array, pa.ChunkedArray)
 
 S = TypeVar("S")
@@ -1152,3 +1156,49 @@ class BoundTransform(BoundTerm):
     def __init__(self, term: BoundTerm, transform: Transform[Any, Any]):
         self.term: BoundTerm = term
         self.transform = transform
+
+    def ref(self) -> BoundReference:
+        """Return the bound reference of the underlying term."""
+        return self.term.ref()
+
+    def eval(self, struct: StructProtocol) -> Any:
+        """Evaluate the transform on the struct value."""
+        return self.transform.transform(self.term.ref().field.field_type)(self.term.eval(struct))
+
+
+class UnboundTransform(UnboundTerm):
+    """An unbound transform expression that pairs a column reference with a transform.
+
+    When bound to a schema, validates that the transform is compatible with the
+    source column type and produces a BoundTransform.
+    """
+
+    transform: Transform[Any, Any]
+
+    def __init__(self, term: UnboundTerm, transform: Transform[Any, Any]):
+        self.term: UnboundTerm = term
+        self.transform = transform
+
+    def bind(self, schema: "Schema", case_sensitive: bool = True) -> BoundTransform:
+        """Bind this unbound transform to a schema, returning a BoundTransform."""
+        bound_term = self.term.bind(schema, case_sensitive)
+        source_type = bound_term.ref().field.field_type
+        if not self.transform.can_transform(source_type):
+            raise TypeError(f"Cannot apply {self.transform} to {source_type}")
+        return BoundTransform(bound_term, self.transform)
+
+    @property
+    def as_bound(self) -> type[BoundTransform]:
+        return BoundTransform
+
+    def __eq__(self, other: object) -> bool:
+        """Return the equality of two instances of the UnboundTransform class."""
+        return isinstance(other, UnboundTransform) and self.term == other.term and self.transform == other.transform
+
+    def __repr__(self) -> str:
+        """Return the string representation of the UnboundTransform class."""
+        return f"UnboundTransform(term={self.term}, transform={self.transform})"
+
+    def __hash__(self) -> int:
+        """Return the hash of the UnboundTransform."""
+        return hash((self.term, self.transform))

tests/expressions/test_parser.py

@@ -42,7 +42,14 @@
     Reference,
     StartsWith,
 )
-from pyiceberg.expressions.literals import DecimalLiteral, LongLiteral, literal
+from pyiceberg.expressions.literals import DecimalLiteral, LongLiteral, StringLiteral, literal
+from pyiceberg.transforms import (
+    DayTransform,
+    HourTransform,
+    MonthTransform,
+    UnboundTransform,
+    YearTransform,
+)
 
 
 def test_always_true() -> None:
@@ -272,3 +279,86 @@ def test_valid_between_with_numerics() -> None:
     ) == parser.parse("foo between '2025-01-01T00:00:00.000000' and '2025-01-10T12:00:00.000000'")
 
     assert parser.parse("foo between 1 and 3") == parser.parse("1 <= foo and foo <= 3")
+
+
+def test_cast_date_comparison() -> None:
+    expected = EqualTo(
+        UnboundTransform(Reference("created_at"), DayTransform()),
+        StringLiteral("2024-01-01"),
+    )
+    assert expected == parser.parse("CAST(created_at AS date) = '2024-01-01'")
+
+
+def test_cast_year_comparison() -> None:
+    expected = GreaterThan(
+        UnboundTransform(Reference("ts"), YearTransform()),
+        LongLiteral(2020),
+    )
+    assert expected == parser.parse("CAST(ts AS year) > 2020")
+
+
+def test_cast_month_comparison() -> None:
+    expected = EqualTo(
+        UnboundTransform(Reference("ts"), MonthTransform()),
+        LongLiteral(6),
+    )
+    assert expected == parser.parse("CAST(ts AS month) = 6")
+
+
+def test_cast_hour_comparison() -> None:
+    expected = GreaterThanOrEqual(
+        UnboundTransform(Reference("ts"), HourTransform()),
+        LongLiteral(12),
+    )
+    assert expected == parser.parse("CAST(ts AS hour) >= 12")
+
+
+def test_cast_case_insensitive() -> None:
+    """CAST keyword and type name should be case-insensitive."""
+    expected = EqualTo(
+        UnboundTransform(Reference("created_at"), DayTransform()),
+        StringLiteral("2024-01-01"),
+    )
+    assert expected == parser.parse("cast(created_at as DATE) = '2024-01-01'")
+
+
+def test_cast_nested_field() -> None:
+    """CAST should work with dotted column references."""
+    expected = EqualTo(
+        UnboundTransform(Reference("event.timestamp"), DayTransform()),
+        StringLiteral("2024-01-01"),
+    )
+    assert expected == parser.parse("CAST(event.timestamp AS date) = '2024-01-01'")
+
+
+def test_cast_unsupported_type() -> None:
+    with pytest.raises(ValueError, match="Unsupported CAST target type"):
+        parser.parse("CAST(col AS foobar) = 5")
+
+
+def test_cast_not_equal() -> None:
+    expected = NotEqualTo(
+        UnboundTransform(Reference("ts"), YearTransform()),
+        LongLiteral(2020),
+    )
+    assert expected == parser.parse("CAST(ts AS year) != 2020")
+
+
+def test_cast_less_than_or_equal() -> None:
+    expected = LessThanOrEqual(
+        UnboundTransform(Reference("ts"), MonthTransform()),
+        LongLiteral(6),
+    )
+    assert expected == parser.parse("CAST(ts AS month) <= 6")
+
+
+def test_cast_with_and() -> None:
+    """CAST predicates compose with AND/OR via infix_notation."""
+    result = parser.parse("CAST(ts AS date) = '2024-01-01' and status = 'active'")
+    assert isinstance(result, And)
+
+
+def test_cast_with_not() -> None:
+    """NOT should negate a CAST predicate."""
+    result = parser.parse("not CAST(ts AS date) = '2024-01-01'")
+    assert isinstance(result, Not)