多索引Pandas DataFrame切片：Awkward Array转换后的三类查询需求

先给出原始数据准备代码及生成的DataFrame：

import awkward as ak
import pandas as pd

ak_arr = ak.Array([
    {
      'jet_pt': [2.33e+05, 1.1e+04, 1.47e+05, 1.33e+04, 1.73e+05, 1.07e+04], 
      'jet_num': 6, 
      'bb_dR': [0.83e-01, 0.56e-01, 0.98e-01, 0.32e-01, 0.21e-01, 0.66e-01], 
      'hh_m': 3.25e+05
    }, 
    {
      'jet_pt': [1.48e+05, 2.06e+04, 9.93e+04, 1.29e+04], 
      'jet_num': 4, 
      'bb_dR': [0.12e-1, 0.32e-01, 0.45e-01, 0.76e-01, 0.33e-01, 0.54e-01], 
      'hh_m': 2.87e+05
    }
])

df = ak.to_dataframe(ak_arr, how='outer')

生成的DataFrame：

jet_pt  jet_num  bb_dR      hh_m
entry subentry                                    
0     0         233000.0      6    0.083  325000.0
      1          11000.0      6    0.056  325000.0
      2         147000.0      6    0.098  325000.0
      3          13300.0      6    0.032  325000.0
      4         173000.0      6    0.021  325000.0
      5          10700.0      6    0.066  325000.0
1     0         148000.0      4    0.012  287000.0
      1          20600.0      4    0.032  287000.0
      2          99300.0      4    0.045  287000.0
      3          12900.0      4    0.076  287000.0
      4              NaN      4    0.033  287000.0
      5              NaN      4    0.054  287000.0

1. 根据jet_num筛选jet_pt，保留对应数量子条目 #

需求：按每个entry的jet_num值，保留该entry下对应数量的jet_pt条目
期望结果：

jet_pt
entry subentry                                    
0     0         233000.0
      1          11000.0
      2         147000.0
      3          13300.0
      4         173000.0
      5          10700.0
1     0         148000.0
      1          20600.0
      2          99300.0
      3          12900.0

原低效实现：

jet_num = df['jet_num'].max(level=0)
jet_z = df['jet_z'].groupby(level=0).apply(lambda x: x[:jet_num[x.name]]).droplevel(0)

优化方案：利用向量化索引操作，避免groupby.apply的循环开销

# 提取每个entry对应的jet_num值
jet_num_map = df['jet_num'].groupby(level='entry').first()
# 生成布尔筛选条件：subentry < 当前entry的jet_num
mask = df.index.get_level_values('subentry') < df.index.get_level_values('entry').map(jet_num_map)
# 筛选jet_pt列
filtered_jet_pt = df.loc[mask, 'jet_pt']

2. 筛选每个entry前4个bb_dR条目 #

需求：保留每个entry下前4条bb_dR数据
期望结果：

bb_dR
entry subentry                                    
0     0         0.083
      1         0.056
      2         0.098
      3         0.032
1     0         0.012
      1         0.032
      2         0.045
      3         0.076

原低效实现：

df['bb_dR'].groupby(level=0).apply(lambda x: x[:4]).droplevel(0)

优化方案：使用Pandas内置的head方法，它是向量化实现，比自定义lambda快得多

filtered_bb_dR = df['bb_dR'].groupby(level='entry').head(4)

3. 提取每个entry第一行的hh_m并移除索引 #

需求：获取每个entry的subentry=0行的hh_m值，同时移除多级索引
期望结果：

0    325000.0
1    287000.0
Name: hh_m, dtype: float64

优化方案：直接用xs定位子条目，再重置索引

# 提取subentry=0的hh_m，然后移除索引
hh_m_first = df.xs(0, level='subentry')['hh_m'].reset_index(drop=True)

或者用索引定位：

hh_m_first = df.loc[(slice(None), 0), 'hh_m'].reset_index(drop=True)

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