Python嵌套循环性能优化求助：十万级数据迭代提速

针对你处理10万行数据耗时100秒的问题，咱们来一步步拆解瓶颈并优化。先分析下当前代码的主要性能损耗点：Python层面的嵌套循环开销、每次循环的列表操作（list(k)+append）、异常捕获的额外成本，还有rowgetter函数内的分支判断在循环中重复执行。

下面是几个针对性的优化方案，按实现复杂度和收益排序：

1. 基础优化：消除循环内的冗余操作与低效代码 #

这一步不需要引入额外依赖，仅通过调整代码逻辑就能获得显著提升：

优化后的代码 #

import operator

def optimized_generator(source, key_indices, variables, variables_indices):
    # 预定义getkey函数，把分支判断移到循环外
    key_len = len(key_indices)
    if key_len == 0:
        getkey = lambda row: ()
    elif key_len == 1:
        idx = key_indices[0]
        getkey = lambda row: (row[idx],)
    else:
        getkey = operator.itemgetter(*key_indices)
    
    # 预打包变量与索引对，避免每次循环重复zip
    var_idx_pairs = list(zip(variables, variables_indices))
    # 计算需要的最大索引，提前过滤无效行
    all_indices = key_indices + variables_indices
    max_required_idx = max(all_indices) if all_indices else 0

    for row in source:
        # 提前跳过长度不足的行，替代低效的try-except
        if len(row) <= max_required_idx:
            continue
        k = getkey(row)
        # 用生成器表达式+yield from替代显式嵌套循环，底层C实现更快
        yield from (k + (v, row[i]) for v, i in var_idx_pairs)

关键优化点解释 #

• 预计算常量与函数：把getkey的分支判断移到循环外部，避免每次迭代都执行分支逻辑；提前将variables和variables_indices打包成列表，消除循环内重复zip的开销。
• 替换try-except为提前过滤：异常捕获在Python中是高成本操作，提前检查行长度是否满足索引要求，直接跳过无效行，比捕获IndexError高效得多。
• 元组拼接替代列表操作：原来的list(k)+两次append改成直接元组拼接k + (v, row[i])，元组是不可变对象，拼接操作的底层实现比列表修改更高效，还避免了额外的内存分配。
• yield from + 生成器表达式：生成器表达式的迭代由Python的C层处理，比纯Python写的嵌套for循环快很多；yield from直接转发生成器结果，减少了一层Python层面的循环开销。

2. 进阶优化：用NumPy处理大规模结构化数据 #

如果你的数据是规整的结构化数据（每行长度一致），可以用NumPy来进一步提速，NumPy的数组操作都是C级别的，比纯Python循环快几个数量级：

示例代码 #

import numpy as np

def numpy_based_generator(source, key_indices, variables, variables_indices):
    # 转换为NumPy结构化数组（假设source[1]是表头，根据你的数据源调整）
    headers = source[1]
    dtype = [(h, 'U20') for h in headers]
    data = np.array(source[2:], dtype=dtype)
    
    # 提取键列
    key_cols = [headers[i] for i in key_indices]
    keys = data[key_cols]
    
    # 遍历变量与索引，生成结果
    for v, i in zip(variables, variables_indices):
        col_name = headers[i]
        values = data[col_name]
        # 组合键与变量、值
        for key, val in zip(keys, values):
            yield tuple(key) + (v, val)

注意事项 #

• 这个方案适合数据规整的场景，如果你的数据源存在大量长度不一致的行，转换NumPy数组的开销可能抵消收益。
• 转换时要注意数据类型的匹配，避免类型转换错误。

3. 极端优化：用NumJIT编译（Numba） #

如果不想切换到NumPy，还可以用Numba将Python函数编译为机器码，进一步降低循环开销：

示例代码 #

from numba import njit

# Numba对生成器支持有限，这里改为返回列表，也可调整为自定义迭代器
@njit
def numba_process_rows(rows, key_indices, variables, variables_indices):
    result = []
    key_len = len(key_indices)
    var_count = len(variables)
    max_idx = max(key_indices + variables_indices)
    
    for row in rows:
        if len(row) <= max_idx:
            continue
        # 提取键
        k = tuple(row[i] for i in key_indices)
        # 遍历变量与索引
        for idx in range(var_count):
            v = variables[idx]
            i = variables_indices[idx]
            result.append(k + (v, row[i]))
    return result

# 调用时传入数据行（跳过表头部分）
processed = numba_process_rows(source[2:], key_indices, variables, variables_indices)

注意事项 #

• Numba对Python的一些高级特性（如生成器、lambda）支持有限，可能需要调整代码结构。
• 首次编译会有开销，但后续重复调用速度极快。

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