基于None值分组列生成多DataFrame的numpy优化实现问询

核心思路 #

靠Numpy的向量化操作定位分组边界，直接用数组切片批量处理数据，少写冗余循环，最大化发挥Numpy的性能优势，写法更贴合Python/Numpy风格。

• 定位分组分割点
  先找出第0行（id列）和第1行（period列）中包含None的列索引，这些位置就是分组的分隔线。记得补充数组首尾的索引，确保所有列都能被覆盖到分组里。

  import numpy as np
  import pandas as pd
  
  # 示例numpy数组，date列单独准备
  arr = np.array([
      ['A', None, 'B', None, 'C'],
      ['2023Q1', None, '2023Q2', None, '2023Q3'],
      [100, np.nan, 200, np.nan, 300],  # price字段
      [5000, np.nan, 6000, np.nan, 7000],  # volume字段
      [5e8, np.nan, 1.2e9, np.nan, 1.5e9]   # mktcap字段
  ])
  date_arr = pd.date_range(start='2023-01-01', periods=3, freq='Q')  # 示例日期列
  
  # 生成掩码定位None位置：字符串类型用==，数值型NaN用np.isnan
  mask = (arr[0] == 'None') | (arr[1] == 'None')
  split_indices = np.where(mask)[0]
  # 补充首尾索引，方便后续切片
  split_points = np.concatenate([[-1], split_indices, [arr.shape[1]]])
  

• 批量处理分组并生成DataFrame
  遍历分割点区间，对每个分组的列进行切片，提取有效id、period和字段数据，结合date列构造DataFrame，最后一次性合并所有结果。

  df_list = []
  # 逐个处理每个分组区间
  for i in range(len(split_points)-1):
      start = split_points[i] + 1
      end = split_points[i+1]
      # 跳过空分组（比如连续出现None的情况）
      if start >= end:
          continue
      # 取分组内第一个有效id和period（同一分组内该字段值一致）
      current_id = arr[0, start]
      current_period = arr[1, start]
      # 切取当前分组的字段数据
      price = arr[2, start:end]
      volume = arr[3, start:end]
      mktcap = arr[4, start:end]
      # 构造DataFrame：若分组含多列数据，date需重复对应次数，可用np.tile实现
      current_df = pd.DataFrame({
          'date': date_arr[i],
          'id': current_id,
          'period': current_period,
          'price': price,
          'volume': volume,
          'mktcap': mktcap
      })
      df_list.append(current_df)
  
  # 一次性合并所有DataFrame，比循环分步合并效率更高
  final_df = pd.concat(df_list, ignore_index=True)
  

• 关键优化点

  • 用np.where+向量化掩码替代循环查找边界，效率提升明显。
  • 直接用Numpy数组切片提取数据，完全贴合向量化操作的风格，避免逐元素处理的冗余。
  • 先收集所有小DataFrame再统一concat，减少内存碎片化，提升合并效率。

注意事项 #

• 如果数组中的None是数值类型的np.nan，需将掩码判断替换为np.isnan(arr[0]) | np.isnan(arr[1])。
• 若单个分组包含多列数据，date列要与数据长度匹配，可使用np.tile(date_arr[i], end - start)生成对应次数的重复日期。

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