通过传入字符串或其他对象，使用 loc 选择 Pandas DataFrame 中的所有行答案

【问题标题】：Select all rows in a Pandas DataFrame with loc by passing in a string or other object通过传入字符串或其他对象，使用 loc 选择 Pandas DataFrame 中的所有行
【发布时间】：2021-11-21 14:52:45
【问题描述】：

我正在编写一个从 pandas DataFrame 中选择行子集的函数。

函数长这样，

def get_predictions(df: pd.DataFrame, subset: str) -> pd.DataFrame:
    return df['properties', 'prediction'].loc[subset]

我希望这个函数能够处理我想要选择 DataFrame 中所有行的情况。一种解决方案是让子集参数默认为无，如果子集参数设置为无，则返回整个 DataFrame。

def get_predictions(df: pd.DataFrame, subset: str) -> pd.DataFrame:
    if subset is None:
        return df['properties', 'prediction']
    else:
        return df['properties', 'prediction'].loc[subset]

我不喜欢这个解决方案，因为我复制了很多代码。有没有更好的解决方案，不涉及重复。具体来说，是否有一个对象可以传递给 .loc[]，它会返回 DataFrame 中的所有行？

这是我正在寻找的理想解决方案，

def get_predictions(df: pd.DataFrame, subset=MysteryObject) -> pd.DataFrame:
    return df['properties', 'prediction'].loc[MysteryObject]

是否有MysteryObject 可以实现这种期望的行为？

【问题讨论】：

标签： python pandas dataframe slice

【解决方案1】：

让我们尝试将默认值设置为slice(None) 而不仅仅是None：

def get_predictions(
        df: pd.DataFrame, subset: str = slice(None)
) -> pd.DataFrame:
    return df[['properties', 'prediction']].loc[subset]

虽然在一个步骤中对两个轴进行子集化会更好：

def get_predictions(
        df: pd.DataFrame, subset: str = slice(None)
) -> pd.DataFrame:
    return df.loc[subset, ['properties', 'prediction']]

slice(None) 等价于:，只是它可以分配给变量。

df.loc[:, 'col'] == df.loc[slice(None), 'col']

测试代码：

test_df = pd.DataFrame({'properties': [1, 2, 3],
                        'prediction': [4, 5, 6],
                        'other': [7, 8, 9]},
                       index=['a', 'b', 'c'])

print('Subset \'a\'')
print(get_predictions(test_df, 'a'))
print('No Subset')
print(get_predictions(test_df))

输出：

Subset 'a'
properties    1
prediction    4
Name: a, dtype: int64

No Subset
   properties  prediction
a           1           4
b           2           5
c           3           6

【讨论】：

谢谢，切片（无）对我有用。我对我的列和索引都使用了 MultiIndex。因此，ds.raw_df[('properties', 'prediction')].loc[slice(None)] 有效，但 ds.raw_df[(slice(None)),('properties', 'prediction')] 无效。你明白为什么吗？
您在第二次尝试中缺少 loc。 ds.raw_df.loc[slice(None), ('properties', 'prediction')]
太好了，谢谢。那是我的错误。

【解决方案2】：

直接进来

subset = df.index

此外，最好使用 .loc 对行和列进行子集化。这样，您可以查看子集，而不是先生成列的副本。所以就这样做

df.loc[subset, ['properties', 'prediction']]

【讨论】：