From e941541a3d12f355aaa99ba8064811ce7dc03673 Mon Sep 17 00:00:00 2001 From: Nick Crews Date: Mon, 4 Mar 2024 15:50:55 -0900 Subject: [PATCH] feat: support nargs on Table.group_by() Fixes https://github.com/ibis-project/ibis/issues/8504 --- ibis/expr/types/relations.py | 6 ++---- ibis/tests/expr/test_table.py | 8 ++++++++ 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/ibis/expr/types/relations.py b/ibis/expr/types/relations.py index 76be407e3cb8..07c0cb60c420 100644 --- a/ibis/expr/types/relations.py +++ b/ibis/expr/types/relations.py @@ -922,7 +922,7 @@ def schema(self) -> sch.Schema: def group_by( self, - by: str | ir.Value | Iterable[str] | Iterable[ir.Value] | None = (), + *by: str | ir.Value | Iterable[str] | Iterable[ir.Value] | None, **key_exprs: str | ir.Value | Iterable[str] | Iterable[ir.Value], ) -> GroupedTable: """Create a grouped table expression. @@ -978,9 +978,7 @@ def group_by( """ from ibis.expr.types.groupby import GroupedTable - if by is None: - by = () - + by = tuple(v for v in by if v is not None) groups = bind(self, (by, key_exprs)) return GroupedTable(self, groups) diff --git a/ibis/tests/expr/test_table.py b/ibis/tests/expr/test_table.py index 537700e1bed5..961c330024d8 100644 --- a/ibis/tests/expr/test_table.py +++ b/ibis/tests/expr/test_table.py @@ -799,6 +799,14 @@ def test_group_by_kwargs(table): assert_equal(expr, expected) +def test_group_by_nargs(table): + t = table + by = ["f", t.h] + e1 = t.group_by(by, z="g").aggregate(t.d.mean().name("foo")) + e2 = t.group_by(*by, z="g").aggregate(t.d.mean().name("foo")) + assert_equal(e1, e2) + + def test_compound_aggregate_expr(table): # See ibis #24 compound_expr = (table["a"].sum() / table["a"].mean()).name("foo")