[enh]: More permissive `group_by().agg(...)` for `pyarrow`

[dangotbanned]

Related

• Focused on introducing expressions in group_by(...).
  • [Enh]: Ability to use nw.col / Expr for .group_by #1385
  • feat: Allow expressions as group_by keys #2325
• General
  • api: formalise expression expansion in group-by #2225
  • refactor: pass ExprMetadata down to compliant Exprs from narwhals.Expr #1848
  • fix: pyarrow unique in group_by context #1076
  • [Enh]: mode in grouped context #981
• feat: Fully spec TableGroupBy.aggregate zen-xu/pyarrow-stubs#197
• feat: Fully spec TableGroupBy.aggregate zen-xu/pyarrow-stubs#197 (comment)

Description

While working on (zen-xu/pyarrow-stubs#197), I learned some new things about pyarrow.Table.group_by.aggregate.

We could allow a wider range of functions

Currently we allow 10 different Expr aggregation methods:

narwhals/narwhals/_compliant/group_by.py

Line 52 in 1642744

"sum", "mean", "median", "max", "min", "std", "var", "len", "n_unique", "count"

For pyarrow, they utilize only 9/21 native functions, that are allowed in the native context.

9 native used

• pc.approximate_median
• pc.count
• pc.count_distinct
• pc.max
• pc.mean
• pc.min
• pc.stddev
• pc.sum
• pc.variance

How are they mapped over?

narwhals/narwhals/_arrow/group_by.py

Lines 29 to 41 in 1642744

	class ArrowGroupBy(EagerGroupBy["ArrowDataFrame", "ArrowExpr"]):
	_REMAP_AGGS: ClassVar[Mapping[NarwhalsAggregation, Any]] = {
	"sum": "sum",
	"mean": "mean",
	"median": "approximate_median",
	"max": "max",
	"min": "min",
	"std": "stddev",
	"var": "variance",
	"len": "count",
	"n_unique": "count_distinct",
	"count": "count",
	}

Some also accept options for a small amount of flexibility.
We seem to be using pc.CountOptions and pc.VarianceOptions which is good 🙂.
But pc.ScalarAggregateOptions and pc.TDigestOptions may also be useful.

Projection

Now this was where it got interesting for me.
The docs provide an "annotation":

aggregations type?

class TableGroupBy
    def aggregate(self, aggregations):
        """
        Perform an aggregation over the grouped columns of the table.

        Parameters
        ----------
        aggregations : list[tuple(str, str)] or \
list[tuple(str, str, FunctionOptions)]
        """

Which we have been following in:

narwhals/narwhals/_arrow/group_by.py

Line 60 in 1642744

aggs: list[tuple[str, str, Any]] = []

However, the description in the docs immediately contradicts that by stating:

The column name can be a string, an empty list or a list of column names

Meaning instead of a single str column, we can do:

Updated stubs

UnarySelector: TypeAlias = str # <---------------- We only use this at the moment
NullarySelector: TypeAlias = tuple[()]
NarySelector: TypeAlias = list[str] | tuple[str, ...]

ColumnSelector: TypeAlias = UnarySelector | NullarySelector | NarySelector

Summary

I'm hoping together this could mean:

1. We don't need to consider as many cases (for pyarrow) as complex
2. Expressions that expand to multiple outputs can be performed more efficiently (natively)

[MarcoGorelli]

thanks for looking into this

do you have an example of a complex aggregation being expressible in pyarrow? say .agg((nw.col('a') + nw.col('b')).mean()) for example

[dangotbanned]

Thanks @MarcoGorelli

I think I've been using a different definition of complex, since this issue hasn't looked at chaining.

Maybe I've misunderstood how we've implemented group_by, but I was under the impression that we limit expressions to these 10:

narwhals/narwhals/_compliant/group_by.py

Line 52 in 1642744

"sum", "mean", "median", "max", "min", "std", "var", "len", "n_unique", "count"

I'm terms of complex, I meant that I think we may be able to grow that list for pyarrow - using just the native tools.

I'm happy to be wrong here 😅, but I don't think this would help with your example

Important

I haven't tested any of this yet btw, the stubs PR only got accepted this morning

[dangotbanned]

Alright this might be something we can work with for pyarrow>=13.0.0.

The docs seem to be super outdated.

UDF support is limited to scalar functions.
A scalar function is a function which executes elementwise operations on arrays or scalars.

Well, turns out that they support 4 kinds (including aggregate!) of functions:

• Their signatures are in this part of the stubs
• Also 13.0.0 cython source

Strangely, I can't find any examples beyond their tests:

• https://github.com/apache/arrow/blob/f7bc27132ea84c4639528d73a9c289c7a1db154f/python/pyarrow/conftest.py#L334
• https://github.com/apache/arrow/blob/f7bc27132ea84c4639528d73a9c289c7a1db154f/python/pyarrow/tests/test_udf.py