feat: make LiteralPredicate serializable via internal IcebergBaseModel

Author: jaimeferj

pyiceberg/expressions/__init__.py

@@ -32,14 +32,16 @@
     Union,
 )
 
+from pydantic import Field
+
 from pyiceberg.expressions.literals import (
     AboveMax,
     BelowMin,
     Literal,
     literal,
 )
 from pyiceberg.schema import Accessor, Schema
-from pyiceberg.typedef import IcebergRootModel, L, StructProtocol
+from pyiceberg.typedef import IcebergBaseModel, IcebergRootModel, L, StructProtocol
 from pyiceberg.types import DoubleType, FloatType, NestedField
 from pyiceberg.utils.singleton import Singleton
 
@@ -732,6 +734,39 @@ def __init__(self, term: Union[str, UnboundTerm[Any]], literal: Union[L, Literal
         super().__init__(term)
         self.literal = _to_literal(literal)  # pylint: disable=W0621
 
+    # ---- JSON (Pydantic) serialization helpers ----
+
+    class _LiteralPredicateModel(IcebergBaseModel):
+        type: str = Field(alias="type")
+        term: str
+        value: Any
+
+    def _json_op(self) -> str:
+        mapping = {
+            EqualTo: "eq",
+            NotEqualTo: "not-eq",
+            LessThan: "lt",
+            LessThanOrEqual: "lt-eq",
+            GreaterThan: "gt",
+            GreaterThanOrEqual: "gt-eq",
+            StartsWith: "starts-with",
+            NotStartsWith: "not-starts-with",
+        }
+        for cls, op in mapping.items():
+            if isinstance(self, cls):
+                return op
+        raise ValueError(f"Unknown LiteralPredicate: {type(self).__name__}")
+
+    def model_dump(self, **kwargs: Any) -> dict:
+        term_name = getattr(self.term, "name", str(self.term))
+        return self._LiteralPredicateModel(type=self._json_op(), term=term_name, value=self.literal.value).model_dump(**kwargs)
+
+    def model_dump_json(self, **kwargs: Any) -> str:
+        term_name = getattr(self.term, "name", str(self.term))
+        return self._LiteralPredicateModel(type=self._json_op(), term=term_name, value=self.literal.value).model_dump_json(
+            **kwargs
+        )
+
     def bind(self, schema: Schema, case_sensitive: bool = True) -> BoundLiteralPredicate[L]:
         bound_term = self.term.bind(schema, case_sensitive)
         lit = self.literal.to(bound_term.ref().field.field_type)

tests/expressions/test_expressions.py

@@ -55,8 +55,10 @@
     NotIn,
     NotNaN,
     NotNull,
+    NotStartsWith,
     Or,
     Reference,
+    StartsWith,
     UnboundPredicate,
 )
 from pyiceberg.expressions.literals import Literal, literal
@@ -921,6 +923,7 @@ def test_bound_less_than_or_equal(term: BoundReference[Any]) -> None:
 
 def test_equal_to() -> None:
     equal_to = EqualTo(Reference("a"), literal("a"))
+    assert equal_to.model_dump_json() == '{"type":"eq","term":"a","value":"a"}'
     assert str(equal_to) == "EqualTo(term=Reference(name='a'), literal=literal('a'))"
     assert repr(equal_to) == "EqualTo(term=Reference(name='a'), literal=literal('a'))"
     assert equal_to == eval(repr(equal_to))
@@ -929,6 +932,7 @@ def test_equal_to() -> None:
 
 def test_not_equal_to() -> None:
     not_equal_to = NotEqualTo(Reference("a"), literal("a"))
+    assert not_equal_to.model_dump_json() == '{"type":"not-eq","term":"a","value":"a"}'
     assert str(not_equal_to) == "NotEqualTo(term=Reference(name='a'), literal=literal('a'))"
     assert repr(not_equal_to) == "NotEqualTo(term=Reference(name='a'), literal=literal('a'))"
     assert not_equal_to == eval(repr(not_equal_to))
@@ -937,6 +941,7 @@ def test_not_equal_to() -> None:
 
 def test_greater_than_or_equal_to() -> None:
     greater_than_or_equal_to = GreaterThanOrEqual(Reference("a"), literal("a"))
+    assert greater_than_or_equal_to.model_dump_json() == '{"type":"gt-eq","term":"a","value":"a"}'
     assert str(greater_than_or_equal_to) == "GreaterThanOrEqual(term=Reference(name='a'), literal=literal('a'))"
     assert repr(greater_than_or_equal_to) == "GreaterThanOrEqual(term=Reference(name='a'), literal=literal('a'))"
     assert greater_than_or_equal_to == eval(repr(greater_than_or_equal_to))
@@ -945,6 +950,7 @@ def test_greater_than_or_equal_to() -> None:
 
 def test_greater_than() -> None:
     greater_than = GreaterThan(Reference("a"), literal("a"))
+    assert greater_than.model_dump_json() == '{"type":"gt","term":"a","value":"a"}'
     assert str(greater_than) == "GreaterThan(term=Reference(name='a'), literal=literal('a'))"
     assert repr(greater_than) == "GreaterThan(term=Reference(name='a'), literal=literal('a'))"
     assert greater_than == eval(repr(greater_than))
@@ -953,14 +959,26 @@ def test_greater_than() -> None:
 
 def test_less_than() -> None:
     less_than = LessThan(Reference("a"), literal("a"))
+    assert less_than.model_dump_json() == '{"type":"lt","term":"a","value":"a"}'
     assert str(less_than) == "LessThan(term=Reference(name='a'), literal=literal('a'))"
     assert repr(less_than) == "LessThan(term=Reference(name='a'), literal=literal('a'))"
     assert less_than == eval(repr(less_than))
     assert less_than == pickle.loads(pickle.dumps(less_than))
 
 
+def test_starts_with() -> None:
+    starts_with = StartsWith(Reference("a"), literal("a"))
+    assert starts_with.model_dump_json() == '{"type":"starts-with","term":"a","value":"a"}'
+
+
+def test_not_starts_with() -> None:
+    not_starts_with = NotStartsWith(Reference("a"), literal("a"))
+    assert not_starts_with.model_dump_json() == '{"type":"not-starts-with","term":"a","value":"a"}'
+
+
 def test_less_than_or_equal() -> None:
     less_than_or_equal = LessThanOrEqual(Reference("a"), literal("a"))
+    assert less_than_or_equal.model_dump_json() == '{"type":"lt-eq","term":"a","value":"a"}'
     assert str(less_than_or_equal) == "LessThanOrEqual(term=Reference(name='a'), literal=literal('a'))"
     assert repr(less_than_or_equal) == "LessThanOrEqual(term=Reference(name='a'), literal=literal('a'))"
     assert less_than_or_equal == eval(repr(less_than_or_equal))