Make SetPredicate and Subclasses JSON Serializable with Pydantic

pyiceberg/expressions/__init__.py

@@ -42,6 +42,8 @@
 from pyiceberg.typedef import L, StructProtocol
 from pyiceberg.types import DoubleType, FloatType, NestedField
 from pyiceberg.utils.singleton import Singleton
+from pyiceberg.typedef import IcebergBaseModel
+from pydantic import Field
 
 
 def _to_unbound_term(term: Union[str, UnboundTerm[Any]]) -> UnboundTerm[Any]:
@@ -559,8 +561,9 @@ def as_bound(self) -> Type[BoundNotNaN[L]]:
         return BoundNotNaN[L]
 
 
-class SetPredicate(UnboundPredicate[L], ABC):
-    literals: Set[Literal[L]]
+class SetPredicate(UnboundPredicate[L], IcebergBaseModel, ABC):
+    type: str = Field(default="in", alias="type")
+    literals: Set[Literal[L]] = Field(alias="items")
 
     def __init__(self, term: Union[str, UnboundTerm[Any]], literals: Union[Iterable[L], Iterable[Literal[L]]]):
         super().__init__(term)
@@ -676,6 +679,8 @@ def as_unbound(self) -> Type[NotIn[L]]:
 
 
 class In(SetPredicate[L]):
+    type: Literal["in"] = Field(default="in", alias="type")
+
     def __new__(  # type: ignore  # pylint: disable=W0221
         cls, term: Union[str, UnboundTerm[Any]], literals: Union[Iterable[L], Iterable[Literal[L]]]
     ) -> BooleanExpression:
@@ -698,6 +703,8 @@ def as_bound(self) -> Type[BoundIn[L]]:
 
 
 class NotIn(SetPredicate[L], ABC):
+    type: str = Field(default="not-in", alias="type")
+
     def __new__(  # type: ignore  # pylint: disable=W0221
         cls, term: Union[str, UnboundTerm[Any]], literals: Union[Iterable[L], Iterable[Literal[L]]]
     ) -> BooleanExpression:
@@ -712,7 +719,7 @@ def __new__(  # type: ignore  # pylint: disable=W0221
 
     def __invert__(self) -> In[L]:
         """Transform the Expression into its negated version."""
-        return In[L](self.term, self.literals)
+        return NotIn[L](self.term, self.literals)
 
     @property
     def as_bound(self) -> Type[BoundNotIn[L]]:

tests/expressions/test_expressions.py

@@ -868,6 +868,13 @@ def test_not_in() -> None:
     assert not_in == eval(repr(not_in))
     assert not_in == pickle.loads(pickle.dumps(not_in))
 
+def test_serialize_in() -> None:
+    pred = In(term="foo", literals=[1, 2, 3])
+    assert pred.model_dump_json() == '{"type":"in","term":"foo","value":[1,2,3]}'
+
+def test_serialize_not_in() -> None:
+    pred = NotIn(term="foo", literals=[1, 2, 3])
+    assert pred.model_dump_json() == '{"type":"not-in","term":"foo","value":[1,2,3]}'
 
 def test_bound_equal_to(term: BoundReference[Any]) -> None:
     bound_equal_to = BoundEqualTo(term, literal("a"))