Make UnaryPredicate JSON serializable and add unit test

Author: Aniketsy

pyiceberg/expressions/__init__.py

@@ -39,7 +39,7 @@
     literal,
 )
 from pyiceberg.schema import Accessor, Schema
-from pyiceberg.typedef import L, StructProtocol
+from pyiceberg.typedef import IcebergBaseModel, L, StructProtocol
 from pyiceberg.types import DoubleType, FloatType, NestedField
 from pyiceberg.utils.singleton import Singleton
 
@@ -429,7 +429,20 @@ def bind(self, schema: Schema, case_sensitive: bool = True) -> BooleanExpression
     def as_bound(self) -> Type[BoundPredicate[L]]: ...
 
 
-class UnaryPredicate(UnboundPredicate[Any], ABC):
+class UnaryPredicate(UnboundPredicate[Any], IcebergBaseModel, ABC):
+    type: str
+    column: str
+
+    def __init__(self, term: Union[str, UnboundTerm[Any]]):
+        if isinstance(term, Reference):
+            term_name = term.name
+        elif isinstance(term, str):
+            term_name = term
+        else:
+            raise ValueError("term must be a string or Reference")
+        super().__init__(term=Reference(term_name))
+        self.column = term_name
+
     def bind(self, schema: Schema, case_sensitive: bool = True) -> BoundUnaryPredicate[Any]:
         bound_term = self.term.bind(schema, case_sensitive)
         return self.as_bound(bound_term)
@@ -488,6 +501,8 @@ def as_unbound(self) -> Type[NotNull]:
 
 
 class IsNull(UnaryPredicate):
+    type: str = "is-null"
+
     def __invert__(self) -> NotNull:
         """Transform the Expression into its negated version."""
         return NotNull(self.term)
@@ -498,6 +513,8 @@ def as_bound(self) -> Type[BoundIsNull[L]]:
 
 
 class NotNull(UnaryPredicate):
+    type: str = "not-null"
+
     def __invert__(self) -> IsNull:
         """Transform the Expression into its negated version."""
         return IsNull(self.term)
@@ -540,6 +557,8 @@ def as_unbound(self) -> Type[NotNaN]:
 
 
 class IsNaN(UnaryPredicate):
+    type: str = "is-nan"
+
     def __invert__(self) -> NotNaN:
         """Transform the Expression into its negated version."""
         return NotNaN(self.term)
@@ -550,6 +569,8 @@ def as_bound(self) -> Type[BoundIsNaN[L]]:
 
 
 class NotNaN(UnaryPredicate):
+    type: str = "not-nan"
+
     def __invert__(self) -> IsNaN:
         """Transform the Expression into its negated version."""
         return IsNaN(self.term)

tests/expressions/test_expressions.py

@@ -707,7 +707,7 @@ def test_and() -> None:
     assert and_ == pickle.loads(pickle.dumps(and_))
 
     with pytest.raises(ValueError, match="Expected BooleanExpression, got: abc"):
-        null & "abc"  # type: ignore
+        null & "abc"
 
 
 def test_or() -> None:
@@ -724,7 +724,7 @@ def test_or() -> None:
     assert or_ == pickle.loads(pickle.dumps(or_))
 
     with pytest.raises(ValueError, match="Expected BooleanExpression, got: abc"):
-        null | "abc"  # type: ignore
+        null | "abc"
 
 
 def test_not() -> None:
@@ -791,6 +791,16 @@ def test_not_null() -> None:
     assert non_null == pickle.loads(pickle.dumps(non_null))
 
 
+def test_serialize_is_null() -> None:
+    pred = IsNull(term="foo")
+    assert pred.model_dump_json() == '{"type":"is-null","term":"foo"}'
+
+
+def test_serialize_not_null() -> None:
+    pred = NotNull(term="foo")
+    assert pred.model_dump_json() == '{"type":"not-null","term":"foo"}'
+
+
 def test_bound_is_nan(accessor: Accessor) -> None:
     # We need a FloatType here
     term = BoundReference[float](