Skip to content

Python: Refactor unary and set expressions #5362

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
danielcweeks merged 9 commits into from
Jul 29, 2022
Merged

Python: Refactor unary and set expressions #5362

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
d6f31a6
Python: Refactor unary expressions.
rdblue Jul 26, 2022
03b622a
Python: Add SetPredicate and BoundSetPredicate to simply in cases.
rdblue Jul 26, 2022
c0416c2
Update for review comments.
rdblue Jul 27, 2022
ed6f679
Fix type checks.
rdblue Jul 27, 2022
5eaafff
Fix linting.
rdblue Jul 27, 2022
2bf12be
Fix mypy issues.
rdblue Jul 28, 2022
aa73775
Add LiteralPredicate and refactor.
rdblue Jul 28, 2022
d8aa0b7
Rename inequalities.
rdblue Jul 28, 2022
947056d
Python: Ignore false positive pylint W0221 for 3.8.
rdblue Jul 29, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
281 changes: 152 additions & 129 deletions python/pyiceberg/expressions/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,58 +19,22 @@
from abc import ABC, abstractmethod
from dataclasses import dataclass
from functools import reduce, singledispatch
from typing import Generic, TypeVar
from typing import (
ClassVar,
Generic,
Literal,
Comment thread
rdblue marked this conversation as resolved.
Outdated
TypeVar,
)

from pyiceberg.files import StructProtocol
from pyiceberg.schema import Accessor, Schema
from pyiceberg.types import NestedField
from pyiceberg.types import DoubleType, FloatType, NestedField
from pyiceberg.utils.singleton import Singleton

T = TypeVar("T")
B = TypeVar("B")


class Literal(Generic[T], ABC):
"""Literal which has a value and can be converted between types"""

def __init__(self, value: T, value_type: type):
if value is None or not isinstance(value, value_type):
raise TypeError(f"Invalid literal value: {value} (not a {value_type})")
self._value = value

@property
def value(self) -> T:
return self._value # type: ignore

@abstractmethod
def to(self, type_var) -> Literal:
... # pragma: no cover

def __repr__(self):
return f"{type(self).__name__}({self.value})"

def __str__(self):
return str(self.value)

def __eq__(self, other):
return self.value == other.value

def __ne__(self, other):
return not self.__eq__(other)

def __lt__(self, other):
return self.value < other.value

def __gt__(self, other):
return self.value > other.value

def __le__(self, other):
return self.value <= other.value

def __ge__(self, other):
return self.value >= other.value


class BooleanExpression(ABC):
"""Represents a boolean expression tree."""

Expand Down Expand Up @@ -105,6 +69,10 @@ class BaseReference(Generic[T], Term, ABC):
class BoundTerm(Bound[T], Term):
"""Represents a bound term."""

@abstractmethod
def ref(self) -> BoundReference[T]:
...


class UnboundTerm(Unbound[T, BoundTerm[T]], Term):
"""Represents an unbound term."""
Expand All @@ -131,6 +99,9 @@ def eval(self, struct: StructProtocol) -> T:
"""
return self.accessor.get(struct)

def ref(self) -> BoundReference[T]:
return self


@dataclass(frozen=True)
class Reference(UnboundTerm[T], BaseReference[T]):
Expand Down Expand Up @@ -342,151 +313,197 @@ def __str__(self) -> str:


@dataclass(frozen=True)
class AlwaysTrue(BooleanExpression, ABC, Singleton):
class AlwaysTrue(BooleanExpression, Singleton):

@rdblue rdblue Jul 26, 2022

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These aren't actually abstract.

"""TRUE expression"""

def __invert__(self) -> AlwaysFalse:
return AlwaysFalse()


@dataclass(frozen=True)
class AlwaysFalse(BooleanExpression, ABC, Singleton):
class AlwaysFalse(BooleanExpression, Singleton):
"""FALSE expression"""

def __invert__(self) -> AlwaysTrue:
return AlwaysTrue()


class IsNull(UnboundPredicate[T]):
def __invert__(self) -> NotNull:
return NotNull(self.term)
@dataclass(frozen=True)
Comment thread
rdblue marked this conversation as resolved.
class UnaryPredicate(Unbound[T, BooleanExpression], BooleanExpression, ABC):
as_bound: ClassVar[type]
term: UnboundTerm[T]

def _validate_literals(self): # pylint: disable=W0238
if self.literals is not None:
raise AttributeError("Null is a unary predicate and takes no Literals.")
@abstractmethod
def __invert__(self) -> UnaryPredicate:
...

def bind(self, schema: Schema, case_sensitive: bool) -> BoundIsNull[T]:
bound_ref = self.term.bind(schema, case_sensitive)
return BoundIsNull(bound_ref)
def bind(self, schema: Schema, case_sensitive: bool = True) -> BooleanExpression:

@samredai samredai Jul 29, 2022

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this return type be BoundUnaryPredicate? Same question for SetPredicate.bind(...) and LiteralPredicate.bind(...) which can have return types of BoundSetPredicate and BoundLiteralPredicate.

@rdblue rdblue Jul 29, 2022

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Binding may produce a different type of expression. For example, if IsNull is bound to a required type, binding will return AlwaysFalse

bound_term = self.term.bind(schema, case_sensitive)
return self.as_bound(bound_term)


class BoundIsNull(BoundPredicate[T]):
def __invert__(self) -> BoundNotNull:
return BoundNotNull(self.term)
@dataclass(frozen=True)
class BoundUnaryPredicate(Bound[T], BooleanExpression, ABC):
term: BoundTerm[T]

def _validate_literals(self): # pylint: disable=W0238
if self.literals:
raise AttributeError("Null is a unary predicate and takes no Literals.")
@abstractmethod
def __invert__(self) -> BoundUnaryPredicate:
...


class NotNull(UnboundPredicate[T]):
def __invert__(self) -> IsNull:
return IsNull(self.term)
class BoundIsNull(BoundUnaryPredicate[T]):
def __new__(cls, term: BoundTerm[T]):
if term.ref().field.required:
return AlwaysFalse()
return super().__new__(cls)

def _validate_literals(self): # pylint: disable=W0238
if self.literals:
raise AttributeError("NotNull is a unary predicate and takes no Literals.")
def __invert__(self) -> BoundNotNull:
return BoundNotNull(self.term)

def bind(self, schema: Schema, case_sensitive: bool) -> BoundNotNull[T]:
bound_ref = self.term.bind(schema, case_sensitive)
return BoundNotNull(bound_ref)

class BoundNotNull(BoundUnaryPredicate[T]):
def __new__(cls, term: BoundTerm[T]):

@samredai samredai Jul 29, 2022 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

pylint for the 3.8 CI seems to not like that the argument signature is different between this and the default super().__new__. You can change this to match the *args, **kwargs from super, but then you'd have to pull out the kwargs and it won't look as clean plus won't allow positional arguments.

@dataclass(frozen=True)
class BoundNotNull(BoundUnaryPredicate[T]):
    def __new__(cls, *args, **kwargs):
        term: BoundTerm[T] = kwargs["term"]
        if term.ref().field.required:
            return AlwaysTrue()
        return super().__new__(cls)

This is just a linting error though and the assumptions made when using the super type hold well so I think it'd be better to just ignore it for these lines.

@dataclass(frozen=True)
class BoundNotNull(BoundUnaryPredicate[T]):
    def __new__(cls, term: BoundTerm[T]):  # pylint: disable=W0221

@rdblue rdblue Jul 29, 2022

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! I suppressed the false positives.

if term.ref().field.required:
return AlwaysTrue()
return super().__new__(cls)

class BoundNotNull(BoundPredicate[T]):
def __invert__(self) -> BoundIsNull:
return BoundIsNull(self.term)

def _validate_literals(self): # pylint: disable=W0238
if self.literals:
raise AttributeError("NotNull is a unary predicate and takes no Literals.")

class IsNull(UnaryPredicate[T]):
as_bound = BoundIsNull

class IsNaN(UnboundPredicate[T]):
def __invert__(self) -> NotNaN:
return NotNaN(self.term)
def __invert__(self) -> NotNull:
return NotNull(self.term)

def _validate_literals(self): # pylint: disable=W0238
if self.literals:
raise AttributeError("IsNaN is a unary predicate and takes no Literals.")

def bind(self, schema: Schema, case_sensitive: bool) -> BoundIsNaN[T]:
bound_ref = self.term.bind(schema, case_sensitive)
return BoundIsNaN(bound_ref)
class NotNull(UnaryPredicate[T]):
as_bound = BoundNotNull

def __invert__(self) -> IsNull:
return IsNull(self.term)


class BoundIsNaN(BoundUnaryPredicate[T]):
def __new__(cls, term: BoundTerm[T]):
bound_type = term.ref().field.field_type

@Fokko Fokko Jul 28, 2022

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like how this is done, by checking this in the __new__ method.

@rdblue rdblue Jul 28, 2022

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm actually not sure about this. There are mypy errors if you declare the return type from __new__ to be correct. This is from In.__new__ for example:

451: error: Incompatible return type for "__new__" (returns "BooleanExpression", but must return a subtype of "BoundIn[Any]")

@Fokko Fokko Jul 28, 2022

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, the __new__ expects to return to the class it belongs to. I think we can ignore this one because I like the trick of returning a different type very much.

if type(bound_type) in {FloatType, DoubleType}:
return super().__new__(cls)
return AlwaysFalse()

class BoundIsNaN(BoundPredicate[T]):
def __invert__(self) -> BoundNotNaN:
return BoundNotNaN(self.term)

def _validate_literals(self): # pylint: disable=W0238
if self.literals:
raise AttributeError("IsNaN is a unary predicate and takes no Literals.")

class BoundNotNaN(BoundUnaryPredicate[T]):
def __new__(cls, term: BoundTerm[T]):
bound_type = term.ref().field.field_type
if type(bound_type) in {FloatType, DoubleType}:
return super().__new__(cls)
return AlwaysTrue()

def __invert__(self) -> BoundIsNaN:
return BoundIsNaN(self.term)


class NotNaN(UnboundPredicate[T]):
class IsNaN(UnaryPredicate[T]):
as_bound = BoundIsNaN

def __invert__(self) -> NotNaN:
return NotNaN(self.term)


class NotNaN(UnaryPredicate[T]):
as_bound = BoundNotNaN

def __invert__(self) -> IsNaN:
return IsNaN(self.term)

def _validate_literals(self): # pylint: disable=W0238
if self.literals:
raise AttributeError("NotNaN is a unary predicate and takes no Literals.")

def bind(self, schema: Schema, case_sensitive: bool) -> BoundNotNaN[T]:
bound_ref = self.term.bind(schema, case_sensitive)
return BoundNotNaN(bound_ref)
@dataclass(frozen=True)
class SetPredicate(Unbound[T, BooleanExpression], BooleanExpression, ABC):
as_bound: ClassVar[type]
term: UnboundTerm[T]
literals: tuple[Literal[T], ...]

@abstractmethod
def __invert__(self) -> BooleanExpression:
...

def bind(self, schema: Schema, case_sensitive: bool = True) -> BooleanExpression:
bound_term = self.term.bind(schema, case_sensitive)
return self.as_bound(bound_term, {lit.to(bound_term.ref().field.field_type) for lit in self.literals})

class BoundNotNaN(BoundPredicate[T]):
def __invert__(self) -> BoundIsNaN:
return BoundIsNaN(self.term)

def _validate_literals(self): # pylint: disable=W0238
if self.literals:
raise AttributeError("NotNaN is a unary predicate and takes no Literals.")
@dataclass(frozen=True)
class BoundSetPredicate(Bound[T], BooleanExpression, ABC):
term: BoundTerm[T]
literals: set[Literal[T]]

@abstractmethod
def __invert__(self) -> BooleanExpression:
...

class BoundIn(BoundPredicate[T]):
def _validate_literals(self): # pylint: disable=W0238
if not self.literals:
raise AttributeError("BoundIn must contain at least 1 literal.")

def __invert__(self) -> BoundNotIn[T]:
return BoundNotIn(self.term, *self.literals)
class BoundIn(BoundSetPredicate[T]):
def __new__(cls, term: BoundTerm[T], literals: set[Literal[T]]) -> BooleanExpression:

@Fokko Fokko Jul 28, 2022

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could also consider wrapping it into positional arguments:

class BoundIn(BoundSetPredicate[T]):
    def __new__(cls, term: BoundTerm[T], *literals: Literal[T]) -> BooleanExpression:
        literals_set = set(literals)
        count = len(literals)
        if count == 0:
            return AlwaysFalse()
        if count == 1:
            return BoundEq(term, literals.pop())
        else:
            return super().__new__(cls)

    def __invert__(self) -> BooleanExpression:
        return BoundNotIn(self.term, self.literals)

This way we also don't have to care about mutating the set. See comment below 👍🏻

This way we don't have to pass in a set:

BoundIn[str](
    base.BoundReference(
        field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False),
        accessor=Accessor(position=0, inner=None),
    ),
    StringLiteral("foo"), StringLiteral("bar"), StringLiteral("baz"),
)

Instead of:

BoundIn[str](
    base.BoundReference(
        field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False),
        accessor=Accessor(position=0, inner=None),
    ),
    {StringLiteral("foo"), StringLiteral("bar"), StringLiteral("baz")},
)

count = len(literals)
if count == 0:
return AlwaysFalse()
if count == 1:
Comment thread
rdblue marked this conversation as resolved.
Outdated
return BoundEq(term, literals.pop())
Comment thread
rdblue marked this conversation as resolved.
Outdated
else:
return super().__new__(cls)

def __invert__(self) -> BooleanExpression:
return BoundNotIn(self.term, self.literals)

class In(UnboundPredicate[T]):
def _validate_literals(self): # pylint: disable=W0238
if not self.literals:
raise AttributeError("In must contain at least 1 literal.")

def __invert__(self) -> NotIn[T]:
return NotIn(self.term, *self.literals)
class BoundNotIn(BoundSetPredicate[T]):
def __new__(cls, term: BoundTerm[T], literals: set[Literal[T]]) -> BooleanExpression:
count = len(literals)
if count == 0:
return AlwaysTrue()
if count == 1:
Comment thread
rdblue marked this conversation as resolved.
Outdated
return BoundNotEq(term, literals.pop())
else:
return super().__new__(cls)

def bind(self, schema: Schema, case_sensitive: bool) -> BoundIn[T]:
bound_ref = self.term.bind(schema, case_sensitive)
return BoundIn(bound_ref, *tuple(lit.to(bound_ref.field.field_type) for lit in self.literals)) # type: ignore
def __invert__(self) -> BooleanExpression:
return BoundIn(self.term, self.literals)


class BoundNotIn(BoundPredicate[T]):
def _validate_literals(self): # pylint: disable=W0238
if not self.literals:
raise AttributeError("BoundNotIn must contain at least 1 literal.")
class In(SetPredicate[T]):
as_bound = BoundIn

def __invert__(self) -> BoundIn[T]:
return BoundIn(self.term, *self.literals)
def __new__(cls, term: UnboundTerm[T], literals: tuple[Literal[T], ...]) -> BooleanExpression:

@Fokko Fokko Jul 28, 2022

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about accumulating the positional arguments like:

Suggested change
def __new__(cls, term: UnboundTerm[T], literals: tuple[Literal[T], ...]) -> BooleanExpression:
def __new__(cls, term: UnboundTerm[T], *literals: Literal[T]) -> BooleanExpression:

This way we can write:

def test_in_to_eq():
    assert base.In(
        base.Reference("x"),
        literal(34.56),
        literal(19.25)
    ) == base.Eq(base.Reference("x"), literal(34.56))

instead of:

def test_in_to_eq():
    assert base.In(
        base.Reference("x"),
        (literal(34.56), literal(19.25))
    ) == base.Eq(base.Reference("x"), literal(34.56))

I think this makes the API a bit friendlier and less verbose.

You can easily expand the arguments when you want to supply a list:

class NotIn(SetPredicate[T]):
    as_bound = BoundNotIn
...
    def __invert__(self) -> BooleanExpression:
        return In(self.term, *self.literals)

@rdblue rdblue Jul 28, 2022

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wanted to do this, but I didn't know how to make it work with @dataclass. I looked it up and saw that you have to implement __init__. But maybe we don't want to use @dataclass anyway? It seems to have strange mypy errors:

327: error: Only concrete class can be given where "Type[UnaryPredicate[T]]" is expected

count = len(literals)
if count == 0:
return AlwaysFalse()
if count == 1:
Comment thread
rdblue marked this conversation as resolved.
Outdated
return Eq(term, literals[0])
else:
return super().__new__(cls)

def __invert__(self) -> BooleanExpression:
return NotIn(self.term, self.literals)

class NotIn(UnboundPredicate[T]):
def _validate_literals(self): # pylint: disable=W0238
if not self.literals:
raise AttributeError("NotIn must contain at least 1 literal.")

def __invert__(self) -> In[T]:
return In(self.term, *self.literals)
class NotIn(SetPredicate[T]):
as_bound = BoundNotIn

def bind(self, schema: Schema, case_sensitive: bool) -> BoundNotIn[T]:
bound_ref = self.term.bind(schema, case_sensitive)
return BoundNotIn(bound_ref, *tuple(lit.to(bound_ref.field.field_type) for lit in self.literals)) # type: ignore
def __new__(cls, term: UnboundTerm[T], literals: tuple[Literal[T], ...]) -> BooleanExpression:

@Fokko Fokko Jul 28, 2022

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
def __new__(cls, term: UnboundTerm[T], literals: tuple[Literal[T], ...]) -> BooleanExpression:
def __new__(cls, term: UnboundTerm[T], literals: *Literal[T]) -> BooleanExpression:

See the comment at class In(SetPredicate[T]):

count = len(literals)
if count == 0:
return AlwaysTrue()
if count == 1:
Comment thread
rdblue marked this conversation as resolved.
Outdated
return NotEq(term, literals[0])
else:
return super().__new__(cls)

def __invert__(self) -> BooleanExpression:
return In(self.term, self.literals)


class BoundEq(BoundPredicate[T]):
Expand Down Expand Up @@ -680,6 +697,12 @@ def _(obj: In, visitor: BooleanExpressionVisitor[T]) -> T:
return visitor.visit_unbound_predicate(predicate=obj)


@visit.register(UnboundPredicate)
def _(obj: UnboundPredicate, visitor: BooleanExpressionVisitor[T]) -> T:
"""Visit an In boolean expression with a concrete BooleanExpressionVisitor"""
return visitor.visit_unbound_predicate(predicate=obj)


@visit.register(Or)
def _(obj: Or, visitor: BooleanExpressionVisitor[T]) -> T:
"""Visit an Or boolean expression with a concrete BooleanExpressionVisitor"""
Expand Down
Loading