如何高效重构含3000万条数据的2D NumPy数组（按时间&ID合并）

问题背景 #

我正在处理一个包含约3000万条记录的数据集，每条记录包含时间戳（timestamp）、ID、描述（Description）和浮点型数值（value），原始2D NumPy数组结构如下：

[
[Time 1, ID 1, D 1_1, V 1_1],
[Time 1, ID 1, D 1_2, V 1_2],
...
[Time 2, ID 1, D 2_1, V 2_1],
[Time 2, ID 1, D 2_2, V 2_2],
...
[Time X, ID 2, D X_1, V X_1],
...
]

需要将数组压缩为两种目标格式：

1. 合并同时间戳、同ID的记录：

[
[Time 1, ID 1, D 1_1, V 1_1, D 1_2, V 1_2, ...],
[Time 2, ID 1, D 2_1, V 2_1, D 2_2, V 2_2, ...],
[Time X, ID 2, D X_1, V X_1, ...],
...
]

2. 理想输出为字典数组，但直接构建字典耗时超100小时，需更高效方案：

[{'time': Time1, 'ID': ID1, 'D1_1': V1_1, 'D1_2': V1_2, ...}...]

高效解决方案 #

方案一：使用Pandas（优先推荐） #

Pandas的分组操作基于C实现，对大数据量的处理效率远高于纯Python循环：

1. 将NumPy数组转换为DataFrame：

import pandas as pd

# 假设原始数组名为arr，指定列名对应字段
df = pd.DataFrame(arr, columns=['timestamp', 'id', 'desc', 'value'])

2. 按时间戳和ID分组，合并同组的描述与数值为键值对：

# 分组后将每组的desc和value转为字典，再展开为列
grouped_df = df.groupby(['timestamp', 'id']).apply(
    lambda x: pd.Series(dict(zip(x['desc'], x['value'])))
).reset_index()

3. 转换为目标格式：

• 生成压缩后的2D数组：

compressed_arr = grouped_df.to_numpy()

• 生成字典数组：

dict_arr = grouped_df.to_dict('records')

方案二：纯NumPy操作（无第三方依赖） #

如果不想引入Pandas，可利用NumPy的排序和分组索引实现：

1. 先按时间戳和ID排序，确保同组记录连续：

import numpy as np

# 按ID（第1列）和时间戳（第0列）排序，保证同组记录相邻
sorted_arr = arr[np.lexsort((arr[:, 1], arr[:, 0]))]

2. 定位分组的边界：

# 标记分组起始位置：当时间戳或ID变化时，视为新组
group_flags = np.concatenate([[True], 
                             (sorted_arr[1:, 0] != sorted_arr[:-1, 0]) | 
                             (sorted_arr[1:, 1] != sorted_arr[:-1, 1])])
group_starts = np.where(group_flags)[0]
# 计算每个组的长度
group_lengths = np.diff(np.concatenate([group_starts, [len(sorted_arr)]]))

3. 构建压缩后的2D数组：
   由于每组长度不固定，最终为object类型的不规则数组：

compressed_arr = []
for start, length in zip(group_starts, group_lengths):
    group = sorted_arr[start:start+length]
    # 拼接时间戳、ID，以及所有描述和数值
    compressed_row = np.concatenate([group[0, :2], group[:, 2:].flatten()])
    compressed_arr.append(compressed_row)
compressed_arr = np.array(compressed_arr, dtype=object)

4. 优化版字典数组构建：

dict_arr = []
for start, length in zip(group_starts, group_lengths):
    group = sorted_arr[start:start+length]
    record = {'time': group[0, 0], 'ID': group[0, 1]}
    # 批量添加描述-数值键值对，比逐个添加高效
    record.update(dict(zip(group[:, 2], group[:, 3])))
    dict_arr.append(record)

核心优化点 #

• 避免Python级逐元素循环：用Pandas/NumPy的底层C实现替代，速度提升几十到上百倍
• 先排序再分组：确保同组记录连续，减少分组时的查找开销
• 字典批量更新：使用dict.update()替代逐个添加键值对，降低构建开销

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