21.4 数据类型优化

核心内容

通过选择合适的数据类型（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'})

类型	字节	取值范围	适用场景
int8	1	-128 ~ 127	小整数、等级
int16	2	-32768 ~ 32767	小型 ID
int32	4	±21 亿	一般 ID、计数
int64	8	极大	大整数、时间戳
float32	4	7 位有效数字	精度要求不高的指标
float64	8	15 位有效数字	默认，高精度

溢出风险

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 / None	pd.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

相关笔记

• 21.3 内存分析与分块读取 - 量化内存优化效果
• Pandas-八-分类数据 - category 的完整操作
• Pandas-十三-Arrow 与高性能数据类型 - PyArrow 深度集成