如何切片DataFrame至mask首次匹配行（包含该行）？

问题场景 #

需要对Pandas DataFrame进行切片，保留到mask首次匹配的行（包含该行）。

示例数据 #

import pandas as pd
import numpy as np
df = pd.DataFrame(
    {
        'a': [np.nan, np.nan, np.nan, 20, 12, 42, 33, 32, 31],
        'b': [np.nan, np.nan, np.nan, np.nan, 2333, np.nan, np.nan, 12323, np.nan]
    }
)

Mask定义 #

mask = (
    (df.a.notna()) &
    (df.b.notna())
)

预期输出 #

a        b
0   NaN      NaN
1   NaN      NaN
2   NaN      NaN
3  20.0      NaN
4  12.0   2333.0

mask首次匹配的是索引为4的行，目标是切片至该行。

尝试方案分析 #

方案1（可行） #

# attempt 1
idx = df.loc[mask.cumsum().eq(1) & mask].index[0]
df = df.loc[:idx]
print(df)

这个方法逻辑正确：mask.cumsum()统计累计匹配次数，当累计次数等于1且当前mask为True时，定位到首次匹配的行，取其索引后用loc[:idx]切片，能正确保留到首次匹配行（包含该行）。

方案2（不符合需求） #

# attempt 2
out = df[~mask.cummax()]

mask.cummax()会在首次匹配后将所有后续行标记为True，取反后会排除首次匹配行及之后的所有行，最终结果仅保留索引0-3的行，不符合需求。

更简洁的写法 #

可以用mask.idxmax()直接获取首次匹配的索引，代码更简洁：

idx = mask.idxmax()
df = df.loc[:idx]

idxmax()会返回第一个True值对应的索引，刚好满足需求。

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