如何用Pandas对时长数据框分桶？大内存数据优化方案咨询

问题背景 #

需要对包含时间戳（ts）和时长（dur）的DataFrame按跨度3进行分桶，统计每个分桶内的总时长，同时解决大时间范围（如纳秒单位）下的内存溢出问题。

示例数据：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'ts':   [1, 5, 10, 12],
    'dur':  [1, 2,  6, 6],
})

预期输出：

bin1, 1
bin2, 2
bin3, 0
bin4, 4
bin5, 6
bin6, 2

基础实现方案（小数据量适用） #

核心思路是生成完整分桶区间，计算每条数据与分桶的交集时长后汇总：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'ts':   [1, 5, 10, 12],
    'dur':  [1, 2,  6, 6],
})

# 计算每条数据的时间起止区间
df['start'] = df['ts']
df['end'] = df['ts'] + df['dur']

# 定义分桶参数
bin_size = 3
min_start = df['start'].min()
max_end = df['end'].max()

# 生成所有分桶的区间索引
bin_starts = np.arange(min_start, max_end + bin_size, bin_size)
bins = pd.IntervalIndex.from_arrays(bin_starts[:-1], bin_starts[1:], closed='left')

# 计算单条数据与所有分桶的交集时长
def calculate_overlap(row):
    row_start, row_end = row['start'], row['end']
    overlap_start = np.maximum(bins.left, row_start)
    overlap_end = np.minimum(bins.right, row_end)
    return np.maximum(0, overlap_end - overlap_start)

# 汇总所有分桶的总时长
overlap_matrix = df.apply(calculate_overlap, axis=1)
bin_totals = overlap_matrix.sum(axis=0)

# 格式化输出结果
result = pd.Series(bin_totals.values, index=[f'bin{i+1}' for i in range(len(bin_totals))])
print(result.to_csv(sep=',', header=False))

大内存优化方案（稀疏处理） #

当时间单位为纳秒、数据量较大时，全量分桶会生成巨量数组导致内存溢出，此时采用稀疏处理，仅计算有数据交集的分桶：

import pandas as pd
import numpy as np
from collections import defaultdict

df = pd.DataFrame({
    'ts':   [1, 5, 10, 12],
    'dur':  [1, 2,  6, 6],
})

bin_size = 3
df['start'] = df['ts']
df['end'] = df['ts'] + df['dur']

min_start = df['start'].min()
max_end = df['end'].max()

# 计算分桶的最小/最大索引
min_bin_idx = 0
max_bin_idx = ((max_end - min_start - 1) // bin_size)

# 用字典累加分桶时长（仅存储有数据的分桶）
bin_duration = defaultdict(int)

for _, row in df.iterrows():
    s, e = row['start'], row['end']
    if s >= e:
        continue
    
    # 计算当前区间覆盖的分桶索引范围
    idx_start = ((s - min_start) // bin_size)
    idx_end = ((e - min_start - 1) // bin_size)
    
    if idx_start == idx_end:
        # 仅覆盖单个分桶
        bin_duration[idx_start] += e - s
    else:
        # 累加第一个分桶的部分时长
        first_bin_end = min_start + (idx_start + 1)*bin_size
        bin_duration[idx_start] += first_bin_end - s
        # 累加中间完整分桶的时长
        for idx in range(idx_start + 1, idx_end):
            bin_duration[idx] += bin_size
        # 累加最后一个分桶的部分时长
        last_bin_start = min_start + idx_end*bin_size
        bin_duration[idx_end] += e - last_bin_start

# 生成包含空分桶的完整结果
for idx in range(min_bin_idx, max_bin_idx + 1):
    print(f'bin{idx+1}, {bin_duration.get(idx, 0)}')

进一步内存优化 #

若需存储结果，可转为Pandas稀疏序列，仅保存非0值：

idx_list = list(range(min_bin_idx, max_bin_idx + 1))
val_list = [bin_duration.get(idx, 0) for idx in idx_list]
sparse_series = pd.SparseSeries(val_list, index=[f'bin{idx+1}' for idx in idx_list], fill_value=0)

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