Skip to content

Commit

Permalink
feat(selectors): support column references in column selector
Browse files Browse the repository at this point in the history
  • Loading branch information
gforsyth authored and cpcloud committed Aug 17, 2023
1 parent 5970cfe commit d4fae08
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 4 deletions.
10 changes: 7 additions & 3 deletions ibis/selectors.py
Original file line number Diff line number Diff line change
Expand Up @@ -365,9 +365,9 @@ def all_of(*predicates: str | Predicate) -> Predicate:


@public
def c(*names: str) -> Predicate:
def c(*names: str | ir.Column) -> Predicate:
"""Select specific column names."""
names = frozenset(names)
names = frozenset(col if isinstance(col, str) else col.get_name() for col in names)

def func(col: ir.Value) -> bool:
schema = col.op().table.schema
Expand Down Expand Up @@ -664,10 +664,14 @@ def all() -> Predicate:
return r[:]


def _to_selector(obj: str | Selector | Sequence[str | Selector]) -> Selector:
def _to_selector(
obj: str | Selector | ir.Column | Sequence[str | Selector | ir.Column],
) -> Selector:
"""Convert an object to a `Selector`."""
if isinstance(obj, Selector):
return obj
elif isinstance(obj, ir.Column):
return c(obj.get_name())
elif isinstance(obj, str):
return c(obj)
else:
Expand Down
6 changes: 5 additions & 1 deletion ibis/tests/expr/test_selectors.py
Original file line number Diff line number Diff line change
Expand Up @@ -157,8 +157,12 @@ def zscore(c):
s.across(s.numeric() & ~s.c("year"), (_ - _.mean()) / _.std())
),
lambda t: t.select(s.across(s.numeric() & ~s.c("year"), zscore)),
lambda t: t.select(
s.across(s.numeric() & ~s.c(t.year), (_ - _.mean()) / _.std())
),
lambda t: t.select(s.across(s.numeric() & ~s.c(t.year), zscore)),
],
ids=["deferred", "func"],
ids=["deferred", "func", "deferred-column-ref", "func-column-ref"],
)
def test_across_select(penguins, expr_func):
expr = expr_func(penguins)
Expand Down

0 comments on commit d4fae08

Please sign in to comment.