REGR: Fixed AssertionError in groupby

[TomAugspurger]

Closes #31522

cc @jbrockmendel. Just raising a TypeError when that assert failed didn't work. The finally still runs, which raised an assertion error.

It seemed easier to try to just support this case. IIUC, it only occurs when an (P, n_rows) input block gets split into P result blocks. I believe that

1. The result blocks should all have the same dtype
2. The input block must not have been an extension block, since it's 2d

So it should be safe to just cast the result values into an ndarray. Hopefully...

Are there any edge cases I'm not considering? Some kind of agg that returns a result that can't be put in a 2D block? Even something like .agg(lambda x: pd.Period())
won't hit this, since it has to be a Cython function.

[jbrockmendel]

The result blocks should all have the same dtype

I'm not clear on how we end up with multiple blocks if they all have the same dtype. if so, could we just consolidate?

cc @WillAyd you put a lot of time into this code late last year.

[TomAugspurger]

Pushed up a change to consolidate if needed (and fixed a bug in my test. Had transposed some values on accident).

[WillAyd]

Hmm seems like the actual issue is somewhere else; do we know where we are getting multiple blocks in the first place? I think the consolidate call should be coupled tighter with whatever that function is

[TomAugspurger]

The split happens as a result of result._convert(datetime=True) after the aggregation:

pandas/pandas/core/groupby/generic.py

Line 984 in 73ea6ca

return result._convert(datetime=True)

I'm not exactly sure why that's called, but it seems necessary to split object columns, since only some columns within a block might be converted? Indeed, that reveals a problem even with this patched version

In [44]: df = pd.DataFrame({"A": pd.date_range("2000", periods=4), "B": ['a', 'b', 'c', 'd']}).astype(object)

In [45]: df.groupby([0, 0, 0, 1]).min()
ValueError: Wrong number of items passed 1, placement implies 2

[TomAugspurger]

@jbrockmendel @WillAyd can you take another look when you get a chance? The basic problem is described in #31616 (comment), but the tldr is that we deliberately split the result of object blocks in aggregate. The test at https://github.com/pandas-dev/pandas/pull/31616/files#diff-98307885a4959790d371ce5c886d6039R404 demonstrates a case where this is probably useful.

Regardless, all the code in DataFrameGroupBy._cython_agg_blocks is assuming that operation is 1:1 in terms of the number of blocks. I've chosen to handle split blocks after the fact, similar to deleted_blocks. I think it's slightly less ugly than trying to handle them inside the main for loop...

[TomAugspurger]

Oh, whoops, that last commit is going to fail.

Is there a reason the block at https://github.com/pandas-dev/pandas/pull/31616/files#diff-bfee1ba9e7cb79839776fac1a57ed940R1082-R1084 needs to be in a finally? That makes that section run even when we continue. I would think just dedenting it would achieve the same effect.

[jbrockmendel]

could we avoid some of this by changing the agg_blocks.append to agg_blocks.extend? and construct these separate blocks up in 1069-1076?

[TomAugspurger]

I tried that but it didn't look promising so I abandoned it. Several things work against that

1. The conversion to ndarray. We need to avoid that since we have mixed types
2. The construction of the new block with block.make_block would need to be handle specially, since we have multiple blocks and the origin block's locs aren't correct anymore, since it's been split.

[jbrockmendel]

Hmm this might be out of scope, but i think if we used block.apply it would handle both the make_block and potential splitting

[jreback]

I guess this is ok; the whole groupby operating on blocks has been a mess for quite a while. @TomAugspurger pls rebase.

[jorisvandenbossche]

Rebased

[TomAugspurger]

CI is passing now.

[TomAugspurger]

Will give @jreback and @WillAyd another hour or so for feedback, but I think this is OKish, at least for 1.0.1.

[TomAugspurger]

OK, merging. Apologies for rushing things along, but I think this is an improvement on master.

Happy to work through followups if people have suggestions.

I'll start with #31616 (comment), though immediate see where block.apply would be used. In place of s.aggregate?