ENH: Method for selecting columns from DataFrameGroupBy (and maybe DataFrame)

[mwaskom]

Is your feature request related to a problem?

When using a "fluent"/"method chaining" style of programming with pandas, a common hiccup is the lack of a method for selecting columns from a DataFrameGroupBy object. For example, take this DataFrame

import pandas as pd
df = pd.DataFrame(dict(
    x=[1, 2, 3, 4, 5, 6],
    y=[2, 4, 6, 1, 2, 3],
    a=["a", "b", "c", "a", "b", "c"],
    b=["x", "x", "y", "z", "z", "z"],
))

To group by one column and then return the mean of another column within each group, one has to use multiple syntactical constructs:

(
    df
    .groupby("a")
    ["x"]
    .mean()
)

There is also dot-attribute access columns:

(
    df
    .groupby("a")
    .x
    .mean()
)

but that is not fully general: it won't work if your column name collides with an existing method, and it won't work if the column is defined by a variable.

Another option is to select the column first then groupby using a series, not a name:

(
    df["x"]
    .groupby(df["a"])
    .mean()
)

This is not so bad. Well, ideally each step in the pipeline would be on a new line, and it requires some duplicated typing. But the bigger issue is that it fails when you want to group by more than one column:

(
    df["x"]
    .groupby(df[["a", "b"]])
    .mean()
)

Traceback

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-14-edc38a1f02b1> in <module>
      1 (
----> 2     df["x"]
      3     .groupby(df[["a", "b"]])
      4     .mean()
      5 )

~/miniconda3/envs/py39/lib/python3.9/site-packages/pandas/core/series.py in groupby(self, by, axis, level, as_index, sort, group_keys, squeeze, observed, dropna)
   1689         axis = self._get_axis_number(axis)
   1690 
-> 1691         return SeriesGroupBy(
   1692             obj=self,
   1693             keys=by,

~/miniconda3/envs/py39/lib/python3.9/site-packages/pandas/core/groupby/groupby.py in __init__(self, obj, keys, axis, level, grouper, exclusions, selection, as_index, sort, group_keys, squeeze, observed, mutated, dropna)
    558             from pandas.core.groupby.grouper import get_grouper
    559 
--> 560             grouper, exclusions, obj = get_grouper(
    561                 obj,
    562                 keys,

~/miniconda3/envs/py39/lib/python3.9/site-packages/pandas/core/groupby/grouper.py in get_grouper(obj, key, axis, level, sort, observed, mutated, validate, dropna)
    826         # allow us to passing the actual Grouping as the gpr
    827         ping = (
--> 828             Grouping(
    829                 group_axis,
    830                 gpr,

~/miniconda3/envs/py39/lib/python3.9/site-packages/pandas/core/groupby/grouper.py in __init__(self, index, grouper, obj, name, level, sort, observed, in_axis, dropna)
    541                 if getattr(self.grouper, "ndim", 1) != 1:
    542                     t = self.name or str(type(self.grouper))
--> 543                     raise ValueError(f"Grouper for '{t}' not 1-dimensional")
    544                 self.grouper = self.index.map(self.grouper)
    545                 if not (

ValueError: Grouper for '<class 'pandas.core.frame.DataFrame'>' not 1-dimensional

(I actually feel like this really ought to work, independently of whether it's the best solution to this Feature Request, but let's leave it aside for now).

So in summary: There's no fully-general, non-awkward method for selecting a column from DataFrameGroupBy. While certainly not fatal, it does make fluent pandas harder to read[1] and write[2]. And this is not an esoteric application: groupby/select/apply must be one of the most common pandas workflows.

[1] The great advantage of this method chaining approach for understanding code is that you can read off the sequence of verbs that comprise the operation. Introducing the [column] syntax requires your brain to shift from comprehension mode to production mode — you need to generate a verb for what's happening — which is a costly operation that likely (on the margin) impairs comprehension.

[2] I really like fluent pandas but it currently can be a little bit annoying to fight with autoindent/autoformat tools when writing it. It likely would be easier to improve that tooling if one could assume each pipeline step begins with a dotted method access.

Describe the solution you'd like

The DataFrameGroupBy object could have a column selection method.

Possible names:

DataFrameGroupBy.get

This exists in DataFrame as a thin wrapper around self[arg]. Possibly adds confusion with DataFrameGroupBy.get_group

DataFrameGroupBy.select

This is the method I have reached for more than once. I'm broadly aware of the history of NDFrame.select; this name should now be available, though it could cause some confusion to reintroduce it with different semantics (but xref #26642).

It probably wouldn't make sense to add .select only to DataFrameGroupBy and not to DataFrame/NDFrame. So this would be more work (but also more benefit?)

In code, that might look like

(
    df
    .groupby("a")
    .select("x")
    .mean()
)

API breaking implications

No breakage, unless someone has very old code using DataFrame.select that completely missed the original deprecation cycle and gets resurrected now.

Adding new methods (especially new ways of indexing pandas objects) definitely has an API complexity cost. But I would argue that, by converging on "one way to do it" in method chains, it could provide a a net reduction in complexity from the user perspective.

Additional context

This originally got some discussion on twitter, note especially @TomAugspurger's response.

Other relevant context would be the select verb in dplyr.

[TomAugspurger]

I think I like the idea of ripping off dply'rs select (targeted to selecting rows columns). It feels a bit more natural than NDFrame.filter, which seems like it should operate on the rows (feels like a SQL where clause).

So overall I think +1 to adding select to NDFrameGroupBy and DataFrame. I think I like dplyr's select's API where the first argument can be a scalar, sequence, or callable, more than NDFrame.filter which spreads these out. I think that callables would be passed each scalar in the .columns and would be expected to return a bool indicated whether that key is selected.

[cbrnr]

Just to clarify, dplyr's select selects columns - what do you mean by "(targeted to selecting rows)" in your first sentence?

[TomAugspurger]

By "targeted to selecting rows" I obviously meant "targeted to selecting columns" :) Fixed.

[jreback]

why is

(
    df
    .groupby("a")
    ["x"]
    .mean()
)

just not the answer here? why is 'let's just add api' the answer?

I am -1 on adding almost anything at this point to an already very complicate interface.

[mwaskom]

Because either pandas supports method chaining as an interface or it doesn't.

If it does, then it should be possible to define a sequence of operations using a single syntactical construct. Needing to switch back and forth has a cost when learning pandas, when writing code and, and when reading code.

Now obviously there will be some edge cases that aren't possible through the pandas API, which is where .pipe comes in handy.

But this usecase (group by then select) is likely one of the most common operations in pandas.

I get that there is a maintenance cost to increasing the API surface, but users accrue real benefits from being able to do things in a consistent, predictable way.

[jreback]

Because either pandas supports method chaining as an interface or it doesn't.

again what is wrong with the first construct itself. This does support method chaining. What exactly do you want to do that is not possible?

[mwaskom]

It supports method chaining, then, oops, wait, gotta use the indexing operator, then ok some more method chaining.

That's not exactly the same thing.

I think it would be a mistake to blanket reject API additions that enhance consistency just because they don't add sufficient novel functionality.

(That said, some people feel excited about using a callable to select, which isn't possible in the current API).

[Batalex]

Hi,

the current pandas workflow incites to "extract" the bare minimum data, then apply transformations.
This proposal would make it easier to go the other way and avoid those types of method chaining:

df.query("sex == 'male'").groupby("species").apply(
    lambda x: x.filter(regex="_mm").mean()
)

That being said, this ⬆️ got the job done so far so I totally understand the point on the maintenance burden.

[mwaskom]

One thing I really like about Python is that usually, when you find yourself doing something that feels really awkward, it's a clue that there's a better way to do it.

That intuition has trapped me here many times. I write a pipeline that mixes method and index chaining, think "no way, that's too ugly, there must be a better syntax for this standard task" and then sink time trying to figure out what the right method is called.

[jreback]

It supports method chaining, then, oops, wait, gotta use the indexing operator, then ok some more method chaining.

so you are saying that python itself should not have [] operators and instead have .getitem() or whatever.

but we have these operators. I am not sure what you are actually suggesting, e.g. .select() sure, let's introduce a completely new function that has now precedence.

and please don't put words in my statement

I think it would be a mistake to blanket reject API additions that enhance consistency just because they don't add sufficient novel functionality.

I don't reject things off hand, instead they must have a really high bar. It doesn't need to be novel, just reduce complexity. The proposal may reduce complexity a small amount, at the enormous expense of adding non-standard usages.

[mwaskom]

When someone takes a while to write an informative feature request, with an explanation, and then just gets a "but why add this" response that doesn't engage at all with what was written, it suggests there is a kneejerk policy of rejection. (Which is fine too, you're the project lead, but it doesn't really incentivize thoughtful feature requests).

I don't think everything should be a method call. But I do think that when everything but one step in a chain of operations is a method call, it's ugly and confusing.

If pandas default were to operate in place, with an option for returning a copy, it would be reasonable to say "well we don't really support method-chaining as a first-class interface". But your design decisions have pushed users in that direction.

And the success of pandas — and dplyr pipelines in R — has proven that it's a really useful idiom for data transforms. So it's' painful when there's one (really common) step where it breaks down.

[attack68]

I don't think everything should be a method call. But I do think that when everything but one step in a chain of operations is a method call, it's ugly and confusing.

I disagree. I appreciate its subjective but I see no issue with the current syntax, in fact I prefer the simplicity of the operators to another method, don't find it ugly nor painful. I think it would be more confusing to have and maintain multiple syntax achieving the same result (there are other areas I would like to see focus/consistency on selecting data). However, its is a useful comment to raise discussion, its just not for me.