BUG: BlockManager.references_same_values not working with object-dtype block that has reference to an Index

[jorisvandenbossche]

On main, the following snippet fails:

df = pd.DataFrame({'A': ['a', 'b', 'c', 'd'], 'B': [1, 2, 3, 4], 'C': ['a', 'a', 'b', 'b']})
df = df.set_index("A")
key = df["C"]
df.groupby(key, observed=True)

with

File ~/scipy/pandas/pandas/core/internals/managers.py:342, in BaseBlockManager.references_same_values(self, mgr, blkno)
    337 """
    338 Checks if two blocks from two different block managers reference the
    339 same underlying values.
    340 """
    341 ref = weakref.ref(self.blocks[blkno])
--> 342 return ref in mgr.blocks[blkno].refs.referenced_blocks

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

I don't yet fully understand why the list containment check results in that kind of error (I would expect it just does == checks on the individual weakref objects and then checks if any of them is True)

But it seems to be triggered by at least 1) the index being constructed from a column (if the df is created with an index from scratch (passing index=pd.Index(['a', 'b', 'c', 'd'])), it doesn't trigger this), and 2) the index being object dtype (when trying with an integer column set as index, it also doesn't happen).

So following two examples pass:

df = pd.DataFrame({'B': [1, 2, 3, 4], 'C': ['a', 'a', 'b', 'b']}, index=pd.Index(['a', 'b', 'c', 'd'], name="A"))
key = df["C"]
df.groupby(key, observed=True)


df = pd.DataFrame({'A': [1, 2, 3, 4], 'B': [1, 2, 3, 4], 'C': ['a', 'a', 'b', 'b']})
df = df.set_index("A")
key = df["C"]
df.groupby(key, observed=True)

(this is also all without enabling CoW, we pass by that code path regardless)

cc @phofl I suppose this is originally caused by #51442, but we never noticed it. The same snippet also raises another error on 2.2.0 (ValueError: buffer source array is read-only in comp_method_OBJECT_ARRAY)

[jorisvandenbossche]

I don't yet fully understand why the list containment check results in that kind of error (I would expect it just does == checks on the individual weakref objects and then checks if any of them is True)

Ah, but that is because the weakref object equality defers to the referenced objects if they are still alive. For Blocks that is fine (they don't have a custom __eq__, so equality falls back to identity), but not for Index objects, which do element-wise equality. The same is true for hashing of weakref objects, which I now remember was the reason we can't just use a set of weakrefs, because Index is unhashable (and thus so is a weakref to an Index object).

In [23]: %debug
> /home/joris/scipy/pandas/pandas/core/internals/managers.py(343)references_same_values()
    342         ref = weakref.ref(self.blocks[blkno])
--> 343         return ref in mgr.blocks[blkno].refs.referenced_blocks

ipdb> ref
<weakref at 0x7f7322527ec0; to 'NumpyBlock' at 0x7f73221c8290>
ipdb> mgr.blocks[blkno].refs.referenced_blocks
[<weakref at 0x7f73205b3920; dead>, <weakref at 0x7f7323f33a10; dead>, <weakref at 0x7f73205b3b00; dead>, <weakref at 0x7f7320c2e520; to 'Index' at 0x7f7320a17990>, <weakref at 0x7f7320c2e700; to 'NumpyBlock' at 0x7f73205ca1b0>, <weakref at 0x7f7322527ec0; to 'NumpyBlock' at 0x7f73221c8290>, <weakref at 0x7f7321765030; to 'NumpyBlock' at 0x7f7320618530>]
ipdb> ref == mgr.blocks[blkno].refs.referenced_blocks[0]
False
ipdb> ref == mgr.blocks[blkno].refs.referenced_blocks[3]
array([False, False, False, False])
ipdb> ref == mgr.blocks[blkno].refs.referenced_blocks[5]
True

Not yet fully sure why this doesn't happen for a numeric index, though. Are we adding the Index weakref wrongly?

Probably a separate issue, but the following difference seems a bit suspicious:

# column A set as index is object dtype
In [2]: df = pd.DataFrame({'A': ['a', 'b', 'c', 'd'], 'B': [1, 2, 3, 4], 'C': ['a', 'a', 'b', 'b']})
   ...: df = df.set_index("A")
   ...: key = df["C"]

In [4]: key._mgr._block.refs.referenced_blocks
Out[4]: 
[<weakref at 0x7f33be6df9c0; dead>,
 <weakref at 0x7f336a8a2c50; dead>,
 <weakref at 0x7f336a8a3920; dead>,
 <weakref at 0x7f336a8e4f90; to 'Index' at 0x7f336a83b310>,
 <weakref at 0x7f336a8a0180; to 'NumpyBlock' at 0x7f336a6acd10>,
 <weakref at 0x7f336a8e72e0; to 'NumpyBlock' at 0x7f33be4b3f50>]

# column A set as index is now integer dtype
In [5]: df = pd.DataFrame({'A': [1, 2, 3, 4], 'B': [1, 2, 3, 4], 'C': ['a', 'a', 'b', 'b']})
   ...: df = df.set_index("A")
   ...: key = df["C"]

In [6]: key._mgr._block.refs.referenced_blocks
Out[6]: 
[<weakref at 0x7f33636aa7f0; dead>,
 <weakref at 0x7f336a8a3970; dead>,
 <weakref at 0x7f33636aa160; to 'NumpyBlock' at 0x7f336322af90>,
 <weakref at 0x7f3363158130; to 'NumpyBlock' at 0x7f336a6acbf0>]