Skip to content

Commit 517b930

Browse files
AniketsyFokko
andauthored
Make SetPredicate and Subclasses JSON Serializable with Pydantic (#2557)
#2524 This PR addresses issue by making the `SetPredicate` class and its subclasses (In, NotIn) JSON serializable using Pydantic. - Added tests to verify JSON serialization of `In` and `NotIn` predicates. Please let me know if my approach or fix needs any improvements . I’m open to feedback and happy to make changes based on suggestions. Thank you ! --------- Co-authored-by: Fokko Driesprong <fokko@apache.org>
1 parent 3a9eab3 commit 517b930

File tree

2 files changed

+27
-5
lines changed

2 files changed

+27
-5
lines changed

pyiceberg/expressions/__init__.py

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
TypeVar,
3232
Union,
3333
)
34+
from typing import Literal as TypingLiteral
3435

3536
from pydantic import Field
3637

@@ -41,10 +42,15 @@
4142
literal,
4243
)
4344
from pyiceberg.schema import Accessor, Schema
44-
from pyiceberg.typedef import IcebergRootModel, L, StructProtocol
45+
from pyiceberg.typedef import IcebergBaseModel, IcebergRootModel, L, StructProtocol
4546
from pyiceberg.types import DoubleType, FloatType, NestedField
4647
from pyiceberg.utils.singleton import Singleton
4748

49+
try:
50+
from pydantic import ConfigDict
51+
except ImportError:
52+
ConfigDict = dict
53+
4854

4955
def _to_unbound_term(term: Union[str, UnboundTerm[Any]]) -> UnboundTerm[Any]:
5056
return Reference(term) if isinstance(term, str) else term
@@ -571,12 +577,14 @@ def as_bound(self) -> Type[BoundNotNaN[L]]:
571577
return BoundNotNaN[L]
572578

573579

574-
class SetPredicate(UnboundPredicate[L], ABC):
575-
literals: Set[Literal[L]]
580+
class SetPredicate(IcebergBaseModel, UnboundPredicate[L], ABC):
581+
model_config = ConfigDict(arbitrary_types_allowed=True)
582+
583+
type: TypingLiteral["in", "not-in"] = Field(default="in")
584+
literals: Set[Literal[L]] = Field(alias="items")
576585

577586
def __init__(self, term: Union[str, UnboundTerm[Any]], literals: Union[Iterable[L], Iterable[Literal[L]]]):
578-
super().__init__(term)
579-
self.literals = _to_literal_set(literals)
587+
super().__init__(term=_to_unbound_term(term), items=_to_literal_set(literals)) # type: ignore
580588

581589
def bind(self, schema: Schema, case_sensitive: bool = True) -> BoundSetPredicate[L]:
582590
bound_term = self.term.bind(schema, case_sensitive)
@@ -688,6 +696,8 @@ def as_unbound(self) -> Type[NotIn[L]]:
688696

689697

690698
class In(SetPredicate[L]):
699+
type: TypingLiteral["in"] = Field(default="in", alias="type")
700+
691701
def __new__( # type: ignore # pylint: disable=W0221
692702
cls, term: Union[str, UnboundTerm[Any]], literals: Union[Iterable[L], Iterable[Literal[L]]]
693703
) -> BooleanExpression:
@@ -710,6 +720,8 @@ def as_bound(self) -> Type[BoundIn[L]]:
710720

711721

712722
class NotIn(SetPredicate[L], ABC):
723+
type: TypingLiteral["not-in"] = Field(default="not-in", alias="type")
724+
713725
def __new__( # type: ignore # pylint: disable=W0221
714726
cls, term: Union[str, UnboundTerm[Any]], literals: Union[Iterable[L], Iterable[Literal[L]]]
715727
) -> BooleanExpression:

tests/expressions/test_expressions.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -873,6 +873,16 @@ def test_not_in() -> None:
873873
assert not_in == pickle.loads(pickle.dumps(not_in))
874874

875875

876+
def test_serialize_in() -> None:
877+
pred = In(term="foo", literals=[1, 2, 3])
878+
assert pred.model_dump_json() == '{"term":"foo","type":"in","items":[1,2,3]}'
879+
880+
881+
def test_serialize_not_in() -> None:
882+
pred = NotIn(term="foo", literals=[1, 2, 3])
883+
assert pred.model_dump_json() == '{"term":"foo","type":"not-in","items":[1,2,3]}'
884+
885+
876886
def test_bound_equal_to(term: BoundReference[Any]) -> None:
877887
bound_equal_to = BoundEqualTo(term, literal("a"))
878888
assert str(bound_equal_to) == f"BoundEqualTo(term={str(term)}, literal=literal('a'))"

0 commit comments

Comments
 (0)