diff --git a/pyiceberg/expressions/__init__.py b/pyiceberg/expressions/__init__.py index ece0db82db..f8529d7df4 100644 --- a/pyiceberg/expressions/__init__.py +++ b/pyiceberg/expressions/__init__.py @@ -698,19 +698,18 @@ class SetPredicate(UnboundPredicate, ABC): def __init__( self, term: str | UnboundTerm, literals: Iterable[Any] | Iterable[LiteralValue] | None = None, **kwargs: Any ) -> None: - if literals is None and "values" in kwargs: - literals = kwargs["values"] - - if literals is None: - literal_set: set[LiteralValue] = set() - else: - literal_set = _to_literal_set(literals) + # __new__ already built the set, and may have used up a one-shot iterator doing it. + literal_set = self.__dict__.pop("_literals_from_new", None) + if literal_set is None: + if literals is None and "values" in kwargs: + literals = kwargs["values"] + literal_set = set() if literals is None else _to_literal_set(literals) super().__init__(term=_to_unbound_term(term), values=literal_set) def bind(self, schema: Schema, case_sensitive: bool = True) -> BoundSetPredicate: bound_term = self.term.bind(schema, case_sensitive) - literal_set = self.literals - return self.as_bound(bound_term, {lit.to(bound_term.ref().field.field_type) for lit in literal_set}) # type: ignore + field_type = bound_term.ref().field.field_type + return self.as_bound(bound_term, {lit.to(field_type) for lit in self.literals}) # type: ignore def __str__(self) -> str: """Return the string representation of the SetPredicate class.""" @@ -844,7 +843,9 @@ def __new__( # pylint: disable=W0221 elif count == 1: return EqualTo(term, next(iter(literals_set))) else: - return super().__new__(cls) + predicate = super().__new__(cls) + object.__setattr__(predicate, "_literals_from_new", literals_set) + return predicate def __invert__(self) -> NotIn: """Transform the Expression into its negated version.""" @@ -879,7 +880,9 @@ def __new__( # pylint: disable=W0221 elif count == 1: return NotEqualTo(term, next(iter(literals_set))) else: - return super().__new__(cls) + predicate = super().__new__(cls) + object.__setattr__(predicate, "_literals_from_new", literals_set) + return predicate def __invert__(self) -> In: """Transform the Expression into its negated version.""" diff --git a/tests/expressions/test_expressions.py b/tests/expressions/test_expressions.py index 8ce48a6897..a12726ecbe 100644 --- a/tests/expressions/test_expressions.py +++ b/tests/expressions/test_expressions.py @@ -327,6 +327,14 @@ def test_in_list() -> None: assert In(Reference("foo"), ["a", "bc", "def"]).literals == {literal("a"), literal("bc"), literal("def")} +def test_in_generator() -> None: + assert In(Reference("foo"), (v for v in ["a", "bc", "def"])).literals == {literal("a"), literal("bc"), literal("def")} + + +def test_not_in_generator() -> None: + assert NotIn(Reference("foo"), (v for v in ["a", "bc", "def"])).literals == {literal("a"), literal("bc"), literal("def")} + + def test_not_in_empty() -> None: assert NotIn(Reference("foo"), ()) == AlwaysTrue()