BUG: Fix return type of loc/iloc

[sanggon6107]

• closes BUG: loc[] returns object type instead of float #60600
• Tests added and passed if fixing a bug or adding a new feature
• All code checks passed.
• Added type annotations to new arguments/methods/functions.
• Added an entry in the latest doc/source/whatsnew/vX.X.X.rst file if fixing a bug or adding a new feature.

Description of Linked Issue

loc/iloc inconsistently returns dtype. For example,

>>> import pandas as pd
>>> df = pd.DataFrame([['a', 1., 2.], ['b', 3., 4.]])

>>> df.loc[0, [1, 2]]
1    1.0
2    2.0
Name: 0, dtype: object

>>> df[[1, 2]].loc[0]
1    1.0
2    2.0
Name: 0, dtype: float64

>>> df.loc[[0, 1], 1]
0    1.0
1    3.0
Name: 1, dtype: float64

This behaviour seems to happen following the below sequence:

1. For axis=0, BlockManager.fast_xs() returns a cross-section of df, determining the dtype as object, since df.loc[0,:] is supposed to include 'a'.
2. For axis=1, NDFrame._reindex_with_indexers() returns the result, not additionally inferring the dtype of the result.

Proposed Solution

Based on the above examples, we can conclude that this issue only apprears where axis[0]=int amd axis[1]=list/slice - loc[int/slice].
Therefore, I'd like to propose to add the below codes to additionally infer the dtype after the column selection.

    @final
    def _getitem_lowerdim(self, tup: tuple):

...

                # This is an elided recursive call to iloc/loc
                out = getattr(section, self.name)[new_key]
                # Re-interpret dtype of out.values for loc/iloc[int, list/slice].
                # GH60600
                if (
                    i == 0
                    and isinstance(key, int)
                    and isinstance(new_key, (list, slice))
                ):
                    out = out.infer_objects()
                return out

Thanks!

[rhshadrach]

Thanks for the PR!

[rhshadrach]

This will infer non-object dtype even on data that is object-dtype, no? I do not think this is the right change.

[sanggon6107]

Thanks @rhshadrach !
I've also confirmed what you've mentioned.

>>>import pandas as pd
>>> df = pd.DataFrame({1: [11, 22], 2: [33, 44], "a": [55, 66]}, dtype=object)

>>>df.loc[0, [1, 2]]
1	11
2	33
Name: 0, dtype: int64

>>>df[[1, 2]].loc[0]
1	11
2	33
Name: 0, dtype: object

I'll do some further investigation to fix this.

[jbrockmendel]

I think the way to do it would be by changing getitem_lowerdim to reverse the order in which we do the indexing (i think the loop on L1069). Would need to try it to see if that breaks anything. Best guess would be MultiIndex cases (possibly their perf?)

I think a bunch of getitem_lowerdim dates back to when Panel and PanelND existed, so it may be more general than it needs to be.