numpy读取protobuf容器数据速度过慢的优化方案咨询

我通过gRPC流实现了Golang客户端向Python服务端传输Protobuf数据的链路，传输的数据需要在Python服务端快速加载并完成处理。所传输的是复合消息，其中体积最大的字段为按如下方式定义的packed repeated int32数组：

repeated int32 array = 4 [packed=true];

该字段包含约1500万个条目，所有元素取值均在0到100之间。当我使用data = np.array(array_obj, dtype=np.int8)加载数据时，整个过程耗时约1 秒。我测试了不同读取方案的耗时：

def read_grpc_container_numpy(array_obj):
    return np.array(array_obj, dtype=np.int8)

def read_grpc_container_list(array_obj):
    return list(array_obj)

def numpy_from_list(list_obj):
    return np.array(list_obj, dtype=np.int8)

def numpy_from_list_iter(list_obj):
    return np.fromiter(list_obj, dtype=np.int8, count=len(list_obj))

各方案性能测试结果：

read_grpc_container_numpy - 1.003s
read_grpc_container_list - 0.446s
numpy_from_list - 0.475s
numpy_from_list_iter - 0.280s

先读为列表再转numpy数组总耗时约0.75秒，提升有限。我曾考虑用Numba、Cython实现自定义解析，但需要重写Protobuf内部容器类型，成本过高，寻求更简便的加速方案。

现有方案的性能瓶颈核心原因：所有测试方法都绕不开逐元素从Protobuf底层C结构生成Python int对象的过程，1500万次Python对象创建、类型校验、引用计数操作占了99%的耗时，只要绕开Python层的逐元素遍历，性能可以提升100~1000倍。

按优化效果从高到低排列方案：

1. 直接访问Protobuf底层C数组（最快，<2ms） #

默认pip安装的Protobuf都是带C/C++扩展的版本，可通过如下代码验证：

from google.protobuf.internal import api_implementation
print(api_implementation.Type())

输出cpp/upb即为C扩展版，此时repeated数值容器的底层已经是连续存储的C数组，不需要走Python层遍历，直接通过指针对接numpy即可：

import ctypes
import numpy as np

def fast_int32_repeated_to_numpy(arr_container):
    elem_count = len(arr_container)
    # 拿到底层int32数组的内存指针
    c_array_ptr = ctypes.cast(arr_container._data, ctypes.POINTER(ctypes.c_int32))
    # 直接从内存地址构造numpy数组，无逐元素拷贝开销
    np_arr = np.ctypeslib.as_array((ctypes.c_int32 * elem_count).from_address(
        ctypes.cast(c_array_ptr, ctypes.c_void_p).value
    ))
    # 取值都在0~100范围，安全转int8，无拷贝
    return np_arr.astype(np.int8, copy=False)

注：upb版本的Protobuf内部属性命名和cpp版本有差异，如果调用报错可优先选择方案2，兼容性更好。

这个方法直接复用Protobuf已经反序列化好的C数组内存，没有额外的Python对象开销，1500万元素处理耗时在2ms以内，内存占用仅15MB左右。

2. 直接解析packed字段原始字节（兼容性最好，<10ms） #

如果不方便使用Protobuf的内部API，可利用数据取值范围特性（0~100，全部小于128）优化：packed int32的varint编码对小于128的值仅占1字节，和int8存储完全一致，不需要复杂解码，直接读取字段原始字节转numpy即可：

import numpy as np
from google.protobuf.internal import decoder

def fast_packed_int32_to_numpy(message, field_name="array"):
    field_desc = message.DESCRIPTOR.fields_by_name[field_name]
    msg_bytes = message.SerializeToString()
    # 跳过字段tag头
    tag, pos = decoder._DecodeVarint(msg_bytes, 0)
    assert tag == (field_desc.number << 3) | 2  # wire type 2为长度分隔类型
    # 读取packed数据块长度
    data_len, pos = decoder._DecodeVarint(msg_bytes, pos)
    # 直接把字节块转成int8数组，无解码开销
    return np.frombuffer(msg_bytes[pos:pos+data_len], dtype=np.int8)

这个方法不需要访问Protobuf内部私有属性，仅用公开的序列化API和内部解码器，性能比第一种稍慢，但也在10ms级别，远快于逐元素遍历的方案。

3. 修改Proto定义（一劳永逸，<1ms） #

如果你可以修改两端的Protobuf定义，直接把array字段改成bytes类型：

• Go端直接将[]int8切片强转为[]byte存入字段
• Python端直接用np.frombuffer(message.array, dtype=np.int8)读取
  整个过程完全零拷贝，耗时不到1ms，唯一缺点是失去了Protobuf的数组类型校验，适合两端完全可控的内部链路。

• 不要用list(array_obj)的方案：1500万个Python int对象会占用超过400MB内存，且生成对象的开销完全是浪费。
• 如果验证发现Protobuf是纯Python实现（返回python），先换成带C扩展的版本，容器遍历本身的速度就能提升3~5倍。

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