核心内容

通过选择合适的数据类型(category、int32、float32、sparse、Arrow)显著降低内存占用并提升运算速度。

一、category — 分类数据

适合**低基数(唯一值少)**的字符串或整型列。

import pandas as pd
import numpy as np
 
df = pd.DataFrame({
    'city': np.random.choice(['北京', '上海', '广州', '深圳'], 1_000_000),
    'value': np.random.rand(1_000_000)
})
 
# 转换前内存占用
print(df.memory_usage(deep=True))
# city      ~58 MB (object)
 
# 转换为 category
df['city'] = df['city'].astype('category')
print(df.memory_usage(deep=True))
# city      ~1 MB (category)
 
# 查看内部编码
df['city'].cat.codes.head()

适用条件

唯一值数量 / 总行数 比值越低收益越大(一般建议 < 50%)。

二、int32 / float32 — 降精度

默认 pandas 使用 int64 / float64,数据范围允许时降低精度可减半内存。

df['int_col'] = df['int_col'].astype('int32')
df['float_col'] = df['float_col'].astype('float32')
 
# 指定 dtype 读取
pd.read_csv('large.csv', dtype={'id': 'int32', 'price': 'float32'})
类型字节取值范围适用场景
int81-128 ~ 127小整数、等级
int162-32768 ~ 32767小型 ID
int324±21 亿一般 ID、计数
int648极大大整数、时间戳
float3247 位有效数字精度要求不高的指标
float64815 位有效数字默认,高精度

溢出风险

int32 转为 int8 超出范围会抛出 OverflowError,转换前检查 min() / max()。

三、sparse — 稀疏数据

适合大量为 0 或缺失的数据,仅存储非零值。

from pandas.arrays import SparseArray
 
# 创建稀疏 Series
s = pd.Series(SparseArray([0, 0, 1, 0, 0, 2, 0]))
print(s.sparse.density)   # 非零密度
 
# DataFrame 中单独列指定稀疏 dtype
df['sparse_col'] = pd.arrays.SparseArray(np.zeros(1_000_000))
print(df.memory_usage(deep=True))
 
# 密度阈值参考:< 10% 时收益明显
# 从稠密转稀疏
s_dense = pd.Series(np.zeros(100_000))
s_sparse = s_dense.astype(pd.SparseDtype('float64', fill_value=0.0))

四、Arrow — PyArrow 后端

pandas 2.x 支持 PyArrow 扩展 dtype,提供更好的字符串性能与更丰富的数据类型。

# 读取时直接指定
df = pd.read_csv('large.csv', dtype_backend='pyarrow')
 
# 转换列
df['name'] = df['name'].astype('string[pyarrow]')
df['id'] = df['id'].astype('int64[pyarrow]')
 
# 或全局开启
pd.set_option('future.infer_string', True)
对比项object 字符串string[pyarrow]
内存高(每元素独立对象)低(连续缓冲区)
速度慢快 2-10 倍
缺失值np.nan / Nonepd.NA
NA 传播不稳定一致

五、综合优化流程

def optimize_dtypes(df):
    for col in df.columns:
        if df[col].dtype == 'object':
            nunique = df[col].nunique()
            if nunique / len(df) < 0.5:
                df[col] = df[col].astype('category')
        elif df[col].dtype == 'int64':
            c_min, c_max = df[col].min(), df[col].max()
            for dtype in ['int8', 'int16', 'int32', 'int64']:
                if c_min > np.iinfo(dtype).min and c_max < np.iinfo(dtype).max:
                    df[col] = df[col].astype(dtype)
                    break
        elif df[col].dtype == 'float64':
            df[col] = df[col].astype('float32')
    return df

相关笔记