Skip to content

Commit

Permalink
feat(api): allow selectors in order_by
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud authored and kszucs committed Oct 6, 2023
1 parent dd056f7 commit 359fd5e
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 8 deletions.
24 changes: 16 additions & 8 deletions ibis/expr/types/relations.py
Original file line number Diff line number Diff line change
Expand Up @@ -1296,7 +1296,13 @@ def head(self, n: int = 5) -> Table:

def order_by(
self,
by: str | ir.Column | Sequence[str] | Sequence[ir.Column] | None,
by: str
| ir.Column
| s.Selector
| Sequence[str]
| Sequence[ir.Column]
| Sequence[s.Selector]
| None,
) -> Table:
"""Sort a table by one or more expressions.
Expand Down Expand Up @@ -1421,19 +1427,21 @@ def order_by(
│ 2 │ B │ 6 │
└───────┴────────┴───────┘
"""
import ibis.selectors as s

sort_keys = []
for item in util.promote_list(by):
if isinstance(item, tuple):
if len(item) != 2:
raise ValueError(
"Tuple must be of length 2, got {}".format(len(item))
)
item = (bind_expr(self, item[0]), item[1])
used_tuple_syntax = True
raise ValueError(f"Tuple must be of length 2, got {len(item):d}")
sort_keys.append(bind_expr(self, item[0]), item[1])
elif isinstance(item, s.Selector):
sort_keys.extend(item.expand(self))
else:
item = bind_expr(self, item)
sort_keys.append(item)
sort_keys.append(bind_expr(self, item))

if not sort_keys:
raise com.IbisError("At least one sort key must be provided")
return self.op().order_by(sort_keys).to_expr()

def union(self, table: Table, *rest: Table, distinct: bool = False) -> Table:
Expand Down
15 changes: 15 additions & 0 deletions ibis/tests/expr/test_selectors.py
Original file line number Diff line number Diff line change
Expand Up @@ -475,3 +475,18 @@ def test_c_error_on_misspelled_column(penguins):
sel = s.any_of(s.c("island", "inland"))
with pytest.raises(exc.IbisInputError, match=match):
penguins.select(sel)


def test_order_by_with_selectors(penguins):
expr = penguins.order_by(s.of_type("string"))
assert tuple(key.name for key in expr.op().sort_keys) == (
"species",
"island",
"sex",
)

expr = penguins.order_by(s.all())
assert tuple(key.name for key in expr.op().sort_keys) == tuple(expr.columns)

with pytest.raises(exc.IbisError):
penguins.order_by(~s.all())

0 comments on commit 359fd5e

Please sign in to comment.