BUG: Fixed groupby quantile for listlike q

[TomAugspurger]

Closes #27526

Just a WIP right now. This is under tested for

• groupby(sort=False)
• non-sorted q
• groupby(multiple levels)

[TomAugspurger]

Typo: pandas

[TomAugspurger]

Multiple -> list-like

[TomAugspurger]

Would welcome improvements to this. Essentially, we have two things to do after concating the results

1. Reorder the levels. By default concat places the keys on the outermost level, we need it on the innermost
2. Re-sort things. Initially we're sorted by quantile (the q-th quantile for every group is together, then the next quantile). We need to be grouped by group key, then quantiles. We can't just use sort_index, since we need to preserve the user-provided group key order with sort=False and the order of the quintile list.

[jbrockmendel]

yah this isn't fun. just checking: we can have mixed dtypes here? if we had a single dtype then stack/unstack would be cheap and there might be a prettier workaround

[jbrockmendel]

concat(results, axis=1, keys=q).stack(0) gives the same answer for at least one of the test cases

[TomAugspurger]

Unfortunately that sorts the axis. I theory, we should be able to follow it up with a loc to restore the correct order on that level,

diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py
index c00bd8398..617466b8d 100644
--- a/pandas/core/groupby/groupby.py
+++ b/pandas/core/groupby/groupby.py
@@ -1928,23 +1928,10 @@ class GroupBy(_GroupBy):
                 for qi in q
             ]
             # fix levels to place quantiles on the inside
-            result = concat(results, axis=0, keys=q)
-            order = np.roll(list(range(result.index.nlevels)), -1)
-            result = result.reorder_levels(order)
-            result = result.reindex(q, level=-1)
-
-            # fix order.
-            hi = len(q) * self.ngroups
-            arr = np.arange(0, hi, self.ngroups)
-            arrays = []
-
-            for i in range(self.ngroups):
-                arr = arr + i
-                arrays.append(arr)
-
-            indices = np.concatenate(arrays)
-            assert len(indices) == len(result)
-            return result.take(indices)
+            result = concat(results, axis=1, keys=q).stack(0)
+            slices = [slice(None)] * result.index.ndim
+            result = result.loc[tuple(slices), :]
+            return result
 
     @Substitution(name="groupby")
     def ngroup(self, ascending=True):

But that runs into the (I think known) issue with loc & a MultiIndex

In [21]: s = pd.Series(list(range(12)), index=pd.MultiIndex.from_product([['a', 'b', 'c'], [1, 2, 3, 4]]))

In [22]: s.loc[pd.IndexSlice[:, [3, 1, 2, 4]]]
Out[22]:
a  1     0
   2     1
   3     2
   4     3
b  1     4
   2     5
   3     6
   4     7
c  1     8
   2     9
   3    10
   4    11
dtype: int64

[jbrockmendel]

Darn. Can you add a my-idea-doesnt-work case to the tests

[TomAugspurger]

@WillAyd @jbrockmendel do you have time to take a look at this?

[jbrockmendel]

will look in AM

[TomAugspurger]

@jbrockmendel I added your alternative implementation as a TODO once the MultiIndex issue (#10710) is fixed.

Darn. Can you add a my-idea-doesnt-work case to the tests

It's covered.

[TomAugspurger]

CI failure is just codecov.

[TomAugspurger]

Merging in 1 hour if there aren't any objections (only CI failure is codecov).