Skip to content

Commit

Permalink
refactor(ir): use @Annotated decorator to coerce Selection.order_by a…
Browse files Browse the repository at this point in the history
…nd Aggregation.order_by arguments
  • Loading branch information
kszucs authored and cpcloud committed Aug 7, 2023
1 parent fedd4b1 commit 8b841c1
Showing 1 changed file with 9 additions and 12 deletions.
21 changes: 9 additions & 12 deletions ibis/expr/operations/relations.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,14 @@
import ibis.expr.datatypes as dt
import ibis.expr.operations as ops
from ibis import util
from ibis.common.annotations import attribute
from ibis.common.annotations import annotated, attribute
from ibis.common.collections import FrozenDict # noqa: TCH001
from ibis.common.grounds import Immutable
from ibis.common.patterns import Coercible, Pattern
from ibis.common.typing import VarTuple
from ibis.common.patterns import Coercible
from ibis.common.typing import VarTuple # noqa: TCH001
from ibis.expr.deferred import Deferred
from ibis.expr.operations.core import Column, Named, Node, Scalar, Value
from ibis.expr.operations.sortkeys import SortKey
from ibis.expr.operations.sortkeys import SortKey # noqa: TCH001
from ibis.expr.schema import Schema

if TYPE_CHECKING:
Expand Down Expand Up @@ -458,17 +458,16 @@ def __init__(self, table, selections, predicates, sort_keys, **kwargs):
**kwargs,
)

def order_by(self, sort_exprs):
@annotated
def order_by(self, keys: VarTuple[SortKey]):
from ibis.expr.analysis import shares_all_roots, sub_immediate_parents

p = Pattern.from_typehint(VarTuple[SortKey])
keys = p.validate(sort_exprs, {})

if not self.selections:
if shares_all_roots(keys, table := self.table):
sort_keys = tuple(self.sort_keys) + tuple(
sub_immediate_parents(key, table) for key in keys
)

return Selection(
table,
self.selections,
Expand Down Expand Up @@ -536,12 +535,10 @@ def schema(self):
types.append(value.dtype)
return Schema.from_tuples(zip(names, types))

def order_by(self, sort_exprs):
@annotated
def order_by(self, keys: VarTuple[SortKey]):
from ibis.expr.analysis import shares_all_roots, sub_immediate_parents

p = Pattern.from_typehint(VarTuple[SortKey])
keys = p.validate(sort_exprs, {})

if shares_all_roots(keys, table := self.table):
sort_keys = tuple(self.sort_keys) + tuple(
sub_immediate_parents(key, table) for key in keys
Expand Down

0 comments on commit 8b841c1

Please sign in to comment.