ENH: A new GroupBy method to slice rows preserving index and order

[johnzangwill]

Is your feature request related to a problem?

Pandas provides DataFrameGroupBy.head() and tail(), which efficiently slice the beginning and end of each group while preserving the order and index. I would like to be able to do a general row slice with the same properties. DataFrame has head(), tail() and iloc that behave in a compatible way. There is no corresponding DataFrameGroupBy.iloc.

Describe the solution you'd like

Provide a new DataFrameGroupBy method to slice rows per group

API breaking implications

None

Describe alternatives you've considered

The following are existing ways to extract, say, the second and third entry of each group, assuming that there are a large number of rows in each group (~10000):

1. grouped.apply(lambda x: x.iloc[1:3, :]) - Extremely slow. Does not preserve the order or indexing.
2. grouped.take([1, 2]) - Extremely slow. Does not preserve the order or indexing.
3. grouped.nth([1, 2]) - Quite fast for a small list. Does not preserve the order or indexing.
4. grouped.head(3).groupby('...').tail(2) - Quite fast. Does preserve index and ordering.
5. grouped._selected_obj[mask] where mask is built from grouped.cumcount() - Very fast. Does preserve index and ordering. But uses private attribute of DataFrameGroupBy and takes several lines of code.

Additional context

There are three options:

1. Add an option to an existing method to force it to preserve index and order. But take() is very slow and nth() is quite slow. Neither accept a slice argument, so a range list has to be provided.
2. Easiest: Add a new method taking a slice as an argument and implementing it as in 5 above.
3. Most logical and complete: Add a new iloc attribute analogous to DataFrame.iloc

[jreback]

.nth already does this

[johnzangwill]

As I pointed out, grouped.nth() does not preserve the order or index of the original df, particularly if it has a multiindex. So it is not compaible with head() or tail() and it takes a list rather than a slice argument.

Also, the speed of grouped.nth(list_of_ints) grows with the length of the list. It is over 10 times slower than alternative 5 when working with a large slice.

[johnzangwill]

I appreciate that this has not yet been triaged, but I can propose a solution for GroupBy.iloc that addresses this issue. So I would like to take this.

[jreback]

ok would be ok with .iloc as long as it's clear how this is different than nth head and tail - eg the usecases r clear in the docs and api

[johnzangwill]

take

[johnzangwill]

I have implemented #42947 and submitted a pull request. I'm not sure what happens next (this is my first contribution...)