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

Closes #42864

New indexable attribute GroupBy.iloc to perform most grouped slices analagous to DataFrame.iloc

Row indexes can be integers or slices with +ve step
Row slices preserve the DataFrame order and grouping and extend the existing GroupBy head() and tail() methods
All column indexes are handled as for DataFrame.iloc

-ve slice step is not currently handled, and would have to reverse the order.
There is no plan to implement Integer lists since these do not preserve the order at all and cannot be easily vectorized.

Note
----
Neither GroupBy.nth() nor GroupBy.take() take a slice argument and neither of them preserve the original DataFrame order and index. They are both slow for large integer lists and take() is very slow for large group counts.

Example use case
--------------------
Supose that we have a multi-indexed DataFrame with a large primary index and a secondary sorted
in a different order for each primary index. 
To reduce the DataFrame to a middle slice of each secondary, we can groupby the primary index and then use iloc. 
This preserves the original DataFrame's order and indexing.
(See tests/groupby/test_groupby_iloc)

Labels:
name: Enhancement
name: Groupby

Comments:
Hello @johnzangwill! Thanks for updating this PR. We checked the lines you've touched for [PEP 8](https://www.python.org/dev/peps/pep-0008) issues, and found:











There are currently no PEP 8 issues detected in this Pull Request. Cheers! :beers: 

##### Comment last updated at 2021-10-13 21:39:55 UTC
I have updated the docstring with more detail and made some corrections. I haven't had much luck building the documentation locally so far, so it is dificult for me to see how it renders... Help would be appreciated!
I have a lot more information regarding outputs and timings of the various alternatives, but I don't feel that the docstring is the right place for that. I am happy to share anything that will help review this proposed addition.
I added iloc to groupby,rst so that it appears in the documentation.
I improved the multiindex test to make iloc's value clearer: g.iloc[3:-3] cannot be provided by any of the other
methods , especially in the case of a varying group size:
![image](https://user-images.githubusercontent.com/8927258/130102057-9b9b7c58-0bbf-48c4-8e92-c1e82eac6468.png)

Also, even though iloc is pure Python, it is much faster than nth() for large slices:
![image](https://user-images.githubusercontent.com/8927258/130101622-8cc98797-49a7-4377-bd1a-233806e16623.png)

I thought to add the multiindex test to the docstring, but I feel that it is too long. Please explain how I should document this better to get it successfully reviewed.
Thanks!
so my main question here is why can't we just fix `.nth` itself? we could in theory deprecate it too in favor of this. but i'd rather not have 2 methods which do about the same thing.
I agree that the current situation is quite tragic!

For positional indexing:
- DataFrame has head, tail, iat, iloc, first, last and take.
- GroupBy has head, tail, first, last, nth and take.

Confusingly, GroupBy.first and last are completely different from DataFrame.first and last.
GroupBy.nth and take have the same main argument, but handle the list orders differently.

My GroupBy.iloc is a simple extension of head and tail and does something that none of the other methods can, namely take a slice relative to the end of end of each group. (e.g. grouped.iloc[3: -3])
Therefore, I think that iloc is needed. In fact, it is a response to a specific request from a user in quantitative finance.
It is lightweight and efficient and we don't lose anything by adding it.

Alternatively, you could add, say, grouped.nth(slice(3, -3), preserve_index=True), but iloc[3: -3] is the way that DataFrame would do it. GroupBy.nth is already complicated and includes a filtering function (dropna). iloc is just a simple slice.

Resolving first, last, nth and take is a separate question...
@johnzangwill my main questions is can we deprecate `nth` and simply use `.iloc` as a *complete* replacement (yes I guess we'd have to remove the `dropna=` arg as `iloc` is an indexer and not a function. We would also want to implement `.head/.tail` as well.
`.take` doesn't matter as its pseudo private.
### Tests
I will cut the tests down as you request (but  when I get back from vacation).
The multiindex test was written primarily to demonstrate the proposed functionality and performance. Is there another place for such content or shall I just keep it to myself?

### Replacing nth
I wanted to avoid implementing iloc[[1, 3, 2]] since changing the order is not easy to make efficient, and my goal with iloc was always to preserve the underlying df order. However, the existing grouped.nth([1, 3, 2]) does does not honor the list order and preserves the original df order (the list is treated as a set), whereas df.iloc[[1, 3, 2]] (and grouped.take([1, 3, 2]) for that matter) returns the rows in the order 1, 3, 2.

Maybe calling it iloc is only making problems. If I called it "rows", say, (i.e. grouped.rows[3: -3]) then there would be no expectation that it was the same as DataFrame.iloc and I would be free to implement lists as does nth.

Assuming that, then I propose adding grouped.rows[[1, 3, 2]] (and grouped.rows[{1, 3, 2}]) which will behave the same way as nth, except for retaining the DataFrame's index structure and being much faster.
> ### Tests
> I will cut the tests down as you request (but when I get back from vacation).
> The multiindex test was written primarily to demonstrate the proposed functionality and performance. Is there another place for such content or shall I just keep it to myself?
> 
> ### Replacing nth
> I wanted to avoid implementing iloc[[1, 3, 2]] since changing the order is not easy to make efficient, and my goal with iloc was always to preserve the underlying df order. However, the existing grouped.nth([1, 3, 2]) does does not honor the list order and preserves the original df order (the list is treated as a set), whereas df.iloc[[1, 3, 2]] (and grouped.take([1, 3, 2]) for that matter) returns the rows in the order 1, 3, 2.
> 
> Maybe calling it iloc is only making problems. If I called it "rows", say, (i.e. grouped.rows[3: -3]) then there would be no expectation that it was the same as DataFrame.iloc and I would be free to implement lists as does nth.
> 
> Assuming that, then I propose adding grouped.rows[[1, 3, 2]] (and grouped.rows[{1, 3, 2}]) which will behave the same way as nth, except for retaining the DataFrame's index structure and being much faster.

@johnzangwill the problem is adding another method is immediately -1 from me. if you want to add and call it `.iloc` (and it matches the expections about what iloc does then let's call it that. however my question still stands, can we *fully* deprecate `nth` in that case? or if not, can you simply implement `nth` in *terms* of `iloc`.

alternatively can you use your implementation and make .nth faster? i don't mind clarifying that a list is not actually doing what its suppoed to (e.g. its a bug), you can replicate that behavior as obviously its not tested.

@jreback Thank you for your patience!
At risk of being downvoted into oblivion, let me explain why I have implemented this as a new method, provisionally called GroupBy.rows.
### Why not iloc?
- iloc does not currently exist on GroupBy, so it would be adding a new method anyway.
- The method would resemble DataFrame.iloc, but confusingly would have a different signature and behaviour (as do first and last)
- **Signature**: iloc handles 2 dimensions, but you rightly suggested that my method should only handle rows.
- **Behaviour**. iloc (and take) preserve the list content and order whereas head, tail, and rows preserve the frame order. E.g. df.iloc[[1,1]] returns 2 rows, whereas gp.nth([1,1]) and gp.rows[1,1] return one. My use case does not require that the order of the list be preserved. If it was then gp.rows[1, -2] would be unacceptably expensive to implement, since it would have to swap or duplicate some rows depending on each group length.
### Why not nth?
-	nth accepts integers and lists. The would be no way change its signature to accept a slice (apart from gp.nth(slice(1,2,3)), which I don’t think is very pandas-like…)
-	Without slices then it is not possible to extract, say [3:-3], since the length of the list for each group would depend on the group length.
-	nth currently preserves neither the frame order nor the frame indexing, although one could always add keyword arguments to change that.
### What is rows?
-	rows is a simple and very cheap (just numpy masks) extension of the current gp.head and tail.
-	The exact order and indexing of the original frame are preserved.
-	rows accepts integers, slices and lists.
-	rows can also accept multiple slices (The very useful “slices with gaps”), but if we think that the world is not ready for that then it can always be omitted from the documentation…
### What of nth would be missing?
-	There would be no way to add the dropna argument.
-	nth seems to add an alphabetic sort of the groups as an index. There would be no way to specify that behaviour.
### Previous requests
The following are requests for GroupBy.head to handle negative arguments. gp.head(-3) can now be implemented as gp.rows[:-3] (This is not possible with nth since the group lengths may not be constant)
- https://stackoverflow.com/questions/33783072/head-and-tail-with-negative-indexes-on-pandas-groupby-object
- #9214
- #30192
- #30556
@jreback My CI builds usually (but not always) timeout at 60 minutes. I would rather have 5 ticks than exclamations and crosses! 
Question: Does it matter and is there anything that I can do about it?
@jreback We can't fix nth because rows[1:-2:2] cannot be written as nth(1:-2:2) or nth([1:-2:2]). It is not legal Python. You could implement it as nth([slice(1:-2:2)]), but I would not regard that as a step forwards!
It turns out that the dropna keyword argument of nth only applies to single integer arguments (not to lists.) So nth is already overloaded with two entirely different types of operation.
I have looked at test_nth, and I believe that most, if not all, of the non-dropna tests will be passed by rows without change.
Therefore, I suggest considering deprecating the list functionality of nth but keeping the single integer with dropna functionality, which I assume is being actively used.
> @jreback We can't fix nth because rows[1:-2:2] cannot be written as nth(1:-2:2) or nth([1:-2:2]). It is not legal Python. You could implement it as nth([slice(1:-2:2)]), but I would not regard that as a step forwards!
> It turns out that the dropna keyword argument of nth only applies to single integer arguments (not to lists.) So nth is already overloaded with two entirely different types of operation.
> I have looked at test_nth, and I believe that most, if not all, of the non-dropna tests will be passed by rows without change.
> Therefore, I suggest considering deprecating the list functionality of nth but keeping the single integer with dropna functionality, which I assume is being actively used.

I don't know why you are insisiting on slices here? this doesnt' have a great benefit and IMHO is actually confusing.

bottom line is we are NOT adding a new method w/o first deprecating one. We already have too many. and this is NOT a good argument at all for 'let's just add another one'. I have said again, the choice basically is:

- deprecate nth entirely and add rows (sure this would remove the ability to use dropna, maybe we don't care about this, its probably ok).
- fix nth w/o adding rows publicly.
- just move on

those are the options. 

@jreback Sorry to continue insisting! Slice with negative numbers gives you functionality not available by any other means.  If the group size is varying (as in my use-case) then [3:-3] is the _only_ performant way to take a central chunk of each group. You cannot do it with a list, because the range would change from group to group.
Yes, you can get the functionality with groupby.apply. But that not a vectorized operation and is extremely slow. rows is well over 100 times faster with my test frame.

If we really cannot add a method, and if we can tolerate the slice() notation (rather than the neater [:]) then I suggest that I hack the list functionality out of nth and replace it with my code. We then get faster lists, including slices, without losing the dropna function with single integers.

The combined head and tail slice then looks like `gp.nth([slice(None, 3), slice(-3, None)])`. Not great. but works.

And the much asked for `gp.head(-3) ` becomes `gp.nth([slice(None, -3)])`

Shall I do it?
I have made rows[] 100% compatible with nth(list) and temporarily made it the implementation. It passes all the unit tests.
Looking at the code I see that nth(n) and nth(list) are two entirely different things with separate code. nth(n) takes the dropna keyword argument, whereas nth(list) does not. So reimplementing nth(list) should not break anything. 
I would still like to replace it by rows[], since the slice syntax is much neater. But, failing that, we can just use what I have done in nth, which at least gives the needed functionality.
@jreback @simonjayhawkins @jbrockmendel @mroeschke @Dr-Irv
### Review of groupby slice following September 8th meeting
### Usecase
A large time-series with a variable amount of commodity data per date and grouped by date. The need was to make a df with a middle slice of commodities for each date.
### Implemented so far
A groupby indexer gp.rows (currently private) used to implement gp.nth when dropna is absent.
### Options

- apply(lambda(x: x.iloc[n:-n])): slow and messes up the  index. Extremely slow with default as_index=True. The “semi-private” take() has same performance and behaviour, but only handles lists.

- head(-n), tail(-n) Already considered and abandoned: #9214, #30192 #30556. Note: head and tail do _not_ honour as_index, therefore very fast.

- Add groupby.iloc[n:-n] Pro: similar to df, could extend head and tail. Con: list semantics iloc[[a,b,c]] honours list order and so could swap or duplicate rows depending in group length (as does groupby.take([a,b,c])). This is neither required nor desirable and would be slow.

- nth(slice(n, -n)) This is implemented in my recent commit with no change to existing nth functionality or performance. Note: nth does honour as_index, which defaults to True which is slow. Setting both as_index=False and dropna=False in the df.groupby call leaves behaviour and performance exacly as for head and tail.

- rows[n:-n] new proposed indexer. rows[n] and rows[list] behave exactly as nth. This is included but undocumented in my recent commit. Treats lists as unordered as does nth. 
### Future Options

1. As is, with `rows[]` hidden.
2. As is, but with `rows[]` documented.
3. Deprecate `nth` altogether and lose the grouped `nth(n, dropna=...)` option:
`df.groupby("A").nth(n,dropna="any")` 
would have to be written as
 `df.dropna(how="any").groupby("A").rows[n]`
Notes: 
a) There is a slight difference in behaviour in the (pathological) case when _all_ of one group is `nan`.
b) There is no performance loss. In fact, it is slightly faster chaining through `dropna`.
4. Deprecate `nth` altogether and add another keyword in `df.groupby()` to add `dropna="any"` or `"all"` behaviour to `rows[]` and all other `GroupBy` methods. (Unfortunately `df.groupby()` already has a `dropna `keyword that does something different(!): filtering out `nan` from the group keys.)
5. Deprecate `nth(list)` and keep `nth(n)` with or without `dropna`.
6. Deprecate both `nth(list)` and `nth(n)`  but keep `nth(n, dropna="any"` or `"all")`.

1 and 2 do not break anything. For the others, users would have to recode some or all instances of `nth`. 5 would be the least impact because it would just mean changing the less common `nth(list)` cases.
4 would be a lot of work and testing for little benefit.
looks good. some followups
- make sure we have asv coverage for the additions to `.nth`
- rewrite head/tail impl to use `._rows` (I think `.nth`) is already set

cc @rhshadrach @jbrockmendel if any comments.
> looks good. some followups
> 
> * make sure we have asv coverage for the additions to `.nth`

**I can't see any asv tests even for nth(list) . They are all for nth(0, dropna=...)**

> * rewrite head/tail impl to use `._rows` (I think `.nth`) is already set

**You would finally get negative indexes!**
You would have to decide whether or not to change the existing behaviour.
Options:

- Implement `head(n)` simply as `rows[:n]`: This would change the behaviour. To get the current behaviour would need to add `df.groupby("A", as_index=False, dropna=False).head(n)`.
- Add some code to make `head(n)` behave as it does now (preserving index and order as in the docstring, and preserving nan group index)
- Add an argument `head(n, preserve=True)`, say. (adding `as_index `and `dropna `would be confusing, since `dropna `means something completely different in `nth()`...)

> 
> cc @rhshadrach @jbrockmendel if any comments.


> Add some code to make head(n) behave as it does now (preserving index and order as in the docstring, and preserving nan group index)

this is what i would do for head/tail (note that you may need/want to modify rows so it doesn't do *anything* with indices and instead this is handled at a higher level.
> > Add some code to make head(n) behave as it does now (preserving index and order as in the docstring, and preserving nan group index)
> 
> this is what i would do for head/tail (note that you may need/want to **modify rows so it doesn't do _anything_ with indices** and instead this is handled at a higher level.

That was what I wanted to do in the first place! (Hence the title of the request). Then `head() `and `tail() `would have been easy. 
If I now take out the code to emulate `nth()`, then it will not be so easy to use it to re-code `nth()` and will be impossible to functionally replace it in the future.
Personally, I would rather do that and not worry about deprecating `nth()`, but that is not for me to decide.
> Also - I don't see any tests for using groupby arguments with `nth` (so arguments to DataFrame.groupby - not passed to the nth method of GroupBy) - i.e. with `nth` when using `axis=1`, `as_index=True`, `sort=False`, categorical with `observed=True/False`, and `dropna=False` (in the groupby, there are tests for the argument in `nth`). Are these supported?

Yes, there seems to be a combinatorial explosion of options. There seem to be many tests and they certainly run for a long time. Which does not mean that they test every combination...
I do not currently know what all the options do.
My original `gp.rows[]` was raw (like `head()` and `tail())`. Jeff suggested that I make it behave like `nth()`, with a view to possibly eliminating `nth()` in the future. So I just copied in the needed nth code which seems to now work as before. But I would be happy to go back to how it was.
There are some tests for `_rows` with `as_index=True`, `sort=False`. These are okay to have, but we need to be testing `nth` directly with slices. Can also add to other nth tests if that would be easier. Also need tests using observed/dropna.

> I do not currently know what all the options do.

If we can be of any assistance here, let us know.
> There are some tests for `_rows` with `as_index=True`, `sort=False`. These are okay to have, but we need to be testing `nth` directly with slices. Can also add to other nth tests if that would be easier. Also need tests using observed/dropna.
> 
> > I do not currently know what all the options do.
> 
> If we can be of any assistance here, let us know.

@rhshadrach Thanks for alerting me to that nest of worms!
It turns out that `nth` makes no attempt to implement `axis=1` and fails ungracefully if you try.
Whereas, I need to handle `axis=1` to provide an implementation for groupby `head` and `tail`.
I decided to go back to my original plan of providing a new groupby filter, extending `head` and `tail` and allowing slices and negative indexes, but not attempting to provide a full implementation of `nth()`.
Since it handles both axes now I renamed it from `_rows` to `_body`, but I am open to other suggestions, hoping to make it public one day...
The common code in `head`, `tail`, `nth` and  `_body` is now factored out into `_make_mask()`
I have added a few tests but, please, I really do not want to get bogged down in patching up nth and its tests in this pull request!
> we need to be testing `nth` directly with slices

I had added it to `test_nth.py` as `test_slice()`.
And I have now changed the `head` and `tail` tests in `test_nth.py` to give the right results for negative index and for both axes. 
@jreback Question: Now that the renamed  _body[] is private and not included in the documentation, is there any reason why I cannot put back its docstring?
> * can you add a whatsnew sub section in 1.4 which shows the before / after for the changes in .head/.tail; also can show the .nth slice changes if you'd like

I have added a section to v1.4.0 Enhancements
thanks @johnzangwill really nice (and you have been great about rebasing!)
Did we have a list of follows? (maybe not any)?