如何通过axis参数便捷获取pandas DataFrame的行或列切片？

我希望通过axis参数指定行或列，来获取pandas.DataFrame的切片，也就是实现下方代码中_loc_axis函数的功能，请问有没有更简便的实现方式？

原实现代码：

import pandas as pd

def _loc_axis(df, idx, axis):
    """Call loc along proper axis"""
    # Axis as integer, error if axis is invalid
    axis = {0: 0, 'index': 0,
            1: 1, 'columns': 1}[axis]
    return (df.loc[idx, :] if axis==0 else
            df.loc[:, idx])

df = pd.DataFrame(
    data=[[1,2,3], [4,5,6], [7,8,9]],
    index=['a','b','c'],
    columns=['D','E','F'],
)

print(df)
print(_loc_axis(df, ['a','b'], axis='index'))
print(_loc_axis(df, ['F','D'], axis='columns'))

原代码运行输出：

D  E  F
a  1  2  3
b  4  5  6
c  7  8  9

   D  E  F
a  1  2  3
b  4  5  6

   F  D
a  3  1
b  6  4
c  9  7

几种更简便的实现方式 #

1. 用reindex原生方法（最简洁）

pandas.DataFrame.reindex()本身支持axis参数，可直接传入行/列标签列表，完全覆盖原函数需求，还能自动处理标签不存在的情况（可选丢弃或补NaN）：

def loc_axis(df, idx, axis):
    return df.reindex(idx, axis=axis)

测试后输出和原函数完全一致，且无需自行处理axis的字符串转整数逻辑，pandas原生支持axis='index'或axis='columns'。

2. 简化loc的调用逻辑

如果坚持用loc，可以通过构造索引元组简化代码，避免分支判断：

def loc_axis(df, idx, axis):
    # 根据axis构造loc的索引参数
    loc_args = (idx, slice(None)) if axis in [0, 'index'] else (slice(None), idx)
    return df.loc[loc_args]

这里用slice(None)替代:，效果完全相同，代码更紧凑。

3. 用take方法（适合按位置切片）

如果你的idx是位置索引而非标签，take方法更高效，同样支持axis参数：

def loc_axis(df, idx, axis):
    return df.take(idx, axis=axis)

如果需要从标签转位置，可配合get_indexer：

def loc_axis(df, idx, axis):
    pos = df.index.get_indexer(idx) if axis in [0, 'index'] else df.columns.get_indexer(idx)
    return df.take(pos, axis=axis)

内容的提问来源于stack exchange，提问作者LogZ