如何让Pandas稀疏DataFrame的groupby sum直接输出float32？

问题核心 #

合并超大型稀疏DataFrame时，使用numba引擎的**sum()**会默认返回float64类型，导致内存占用暴增；后续转float32时因顺序执行再次触发内存瓶颈，需要让聚合操作直接输出float32稀疏类型，同时优化整体流程。

当前方案的冗余点 #

• concat后重复执行astype(pd.SparseDtype)：输入的pdfs已经是稀疏float32，concat后类型不会丢失，无需重复转换
• sum()后两次执行astype：完全可以一次完成
• 索引排序的实现效率偏低：转series再处理的开销较大

优化方案 #

方案1：利用sum()的downcast参数强制输出float32 #

针对pandas 1.4+版本（numba引擎已支持downcast参数），可直接在聚合时指定输出类型，避免后续类型转换的内存开销：

import pandas as pd
import numpy as np

def get_random_df(row:int=10, col:int=7):
    np.random.seed(0)
    d=np.random.randint(low=0, high=10, size=(row,col)).astype(np.float32)
    d[d < 4] = np.nan
    df=pd.DataFrame(
        data=d,
        index=[f"ip{i}" for i in np.random.choice(range(max(row, 10)), row, replace=False) ],
        columns=[f"col_{c}" for c in np.random.choice(range(max(col, 10)), col, replace=False) ],
        dtype=pd.SparseDtype(dtype=np.float32),
    )
    df.index.name='usr'
    return df

def get_concat(pdfs):
    # 直接concat，无需重复转稀疏类型
    dfc = pd.concat(pdfs, axis=0, sort=True)
    # groupby聚合时直接指定downcast，输出float32
    dfc = dfc.groupby(level=0).sum(
        engine="numba",
        engine_kwargs={'nopython': True, 'parallel': True, 'nogil': False},
        downcast='float'
    )
    # 一次性转换为目标稀疏类型
    dfc = dfc.astype(pd.SparseDtype(dtype=np.float32, fill_value=0.0))
    # 优化索引排序：直接处理索引字符串，避免转series
    dfc = dfc.sort_index(key=lambda x: x.str[2:].astype(int))
    return dfc

方案2：自定义numba聚合函数强制保持float32精度 #

若numba引擎不支持downcast，可自定义聚合函数，确保计算全程使用float32：

import pandas as pd
import numpy as np
from numba import njit, prange

def get_random_df(row:int=10, col:int=7):
    np.random.seed(0)
    d=np.random.randint(low=0, high=10, size=(row,col)).astype(np.float32)
    d[d < 4] = np.nan
    df=pd.DataFrame(
        data=d,
        index=[f"ip{i}" for i in np.random.choice(range(max(row, 10)), row, replace=False) ],
        columns=[f"col_{c}" for c in np.random.choice(range(max(col, 10)), col, replace=False) ],
        dtype=pd.SparseDtype(dtype=np.float32),
    )
    df.index.name='usr'
    return df

# 自定义numba加速的sum函数，强制返回float32
@njit('float32[:](float32[:,:])', parallel=True)
def numba_float32_sum(arr):
    result = np.zeros(arr.shape[1], dtype=np.float32)
    for i in prange(arr.shape[1]):
        col = arr[:, i]
        mask = ~np.isnan(col)
        result[i] = col[mask].sum()
    return result

def get_concat(pdfs):
    dfc = pd.concat(pdfs, axis=0, sort=True)
    # 调用自定义聚合函数
    dfc = dfc.groupby(level=0).agg(numba_float32_sum)
    dfc = dfc.astype(pd.SparseDtype(dtype=np.float32, fill_value=0.0))
    dfc = dfc.sort_index(key=lambda x: x.str[2:].astype(int))
    return dfc

额外优化建议 #

• 合并前统一检查输入DataFrame的稀疏类型一致性，避免concat时自动触发类型转换
• 若内存仍紧张，采用分批次合并策略：每次合并N个DataFrame，聚合后再与下一批合并，降低中间内存峰值
• 可在生成DataFrame时提前按规则排序索引，避免后期大规模排序开销

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