Skip to content

Commit

Permalink
feat(api): allow transmute-style select method
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud committed Aug 1, 2022
1 parent 2bd075f commit d5fc364
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 14 deletions.
49 changes: 35 additions & 14 deletions ibis/expr/types/relations.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import functools
import itertools
import operator
import warnings
from functools import cached_property
from typing import IO, TYPE_CHECKING, Any, Iterable, Literal, Mapping, Sequence

Expand Down Expand Up @@ -112,13 +113,13 @@ def __getitem__(self, what):
return what._to_semi_join(self)[self]
elif isinstance(what, (list, tuple, Table)):
# Projection case
return self.projection(what)
return self.select(what)
elif isinstance(what, BooleanColumn):
# Boolean predicate
return self.filter([what])
elif isinstance(what, Column):
# Projection convenience
return self.projection(what)
return self.select(what)
else:
raise NotImplementedError(
'Selection rows or columns with {} objects is not '
Expand Down Expand Up @@ -572,11 +573,12 @@ def mutate(
exprs.append(value.name(name))

mutation_exprs = an.get_mutation_exprs(exprs, self)
return self.projection(mutation_exprs)
return self.select(mutation_exprs)

def select(
self,
exprs: ir.Value | str | Sequence[ir.Value | str],
*exprs: ir.Value | str | Iterable[ir.Value | str],
**named_exprs: ir.Value | str,
) -> Table:
"""Compute a new table expression using `exprs`.
Expand All @@ -585,6 +587,10 @@ def select(
automatically constructs a window function expression. See the examples
section for more details.
For backwards compatibility the keyword argument `exprs` is reserved
and cannot be used to name an expression. This behavior will be removed
in v4.
Parameters
----------
exprs
Expand All @@ -601,9 +607,8 @@ def select(
Simple projection
>>> import ibis
>>> fields = [('a', 'int64'), ('b', 'double')]
>>> t = ibis.table(fields, name='t')
>>> proj = t.projection([t.a, (t.b + 1).name('b_plus_1')])
>>> t = ibis.table(dict(a="int64", b="double"), name='t')
>>> proj = t.select(t.a, b_plus_1=t.b + 1)
>>> proj
r0 := UnboundTable[t]
a int64
Expand All @@ -612,13 +617,13 @@ def select(
selections:
a: r0.a
b_plus_1: r0.b + 1
>>> proj2 = t[t.a, (t.b + 1).name('b_plus_1')]
>>> proj2 = t.select("a", b_plus_1=t.b + 1)
>>> proj.equals(proj2)
True
Aggregate projection
>>> agg_proj = t[t.a.sum().name('sum_a'), t.b.mean().name('mean_b')]
>>> agg_proj = t.select(sum_a=t.a.sum(), mean_b=t.b.mean())
>>> agg_proj
r0 := UnboundTable[t]
a int64
Expand All @@ -635,7 +640,7 @@ def select(
The purpose of this expression rewrite is to make it easy to write
column/scalar-aggregate operations like
>>> t[(t.a - t.a.mean()).name('demeaned_a')]
>>> t.select(demeaned_a=t.a - t.a.mean())
r0 := UnboundTable[t]
a int64
b float64
Expand All @@ -645,8 +650,24 @@ def select(
"""
import ibis.expr.analysis as an

if isinstance(exprs, (Expr, str)):
exprs = [exprs]
if backcompat_exprs := named_exprs.pop("exprs", []):
warnings.warn(
"Passing `exprs` as a keyword argument is deprecated"
" and will be removed in 4.0. Pass the value(s) as"
" positional arguments.",
FutureWarning,
)

exprs = list(
itertools.chain(
itertools.chain.from_iterable(map(util.promote_list, exprs)),
util.promote_list(backcompat_exprs),
(
self._ensure_expr(expr).name(name)
for name, expr in named_exprs.items()
),
)
)

projector = an.Projector(self, exprs)
op = projector.get_result()
Expand Down Expand Up @@ -681,7 +702,7 @@ def relabel(self, substitutions: Mapping[str, str]) -> Table:
if c not in observed:
raise KeyError(f'{c!r} is not an existing column')

return self.projection(exprs)
return self.select(exprs)

def drop(self, fields: str | Sequence[str]) -> Table:
"""Remove fields from a table.
Expand Down Expand Up @@ -905,7 +926,7 @@ def set_column(self, name: str, expr: ir.Value) -> Table:
else:
proj_exprs.append(self[key])

return self.projection(proj_exprs)
return self.select(proj_exprs)

def join(
left: Table,
Expand Down
8 changes: 8 additions & 0 deletions ibis/tests/expr/test_table.py
Original file line number Diff line number Diff line change
Expand Up @@ -1518,3 +1518,11 @@ def test_materialize_no_op():
expr = left.inner_join(right, "id")
with pytest.warns(FutureWarning):
expr.materialize()


def test_exprs_to_select():
t = ibis.table(dict(a="string"))
exprs = [t.a.length().name("len")]
with pytest.warns(FutureWarning, match="Passing `exprs`"):
result = t.select(exprs=exprs)
assert result.equals(t.select(len=t.a.length()))

0 comments on commit d5fc364

Please sign in to comment.