用最小堆做日志分片的 k 路归并：Python heapq 流式实现

问题场景

日志采集常按实例或按小时切分写入多个分片文件，单个分片内部按时间戳递增，但分片之间没有全局顺序。下游做审计回放、按时间窗口聚合时，需要一条全局有序的流。把全部分片读进内存再排序，日志量大时内存不可控；而且多数场景只需要边读边写，不需要保留全量数据。

k 路归并正好适配这个约束：每个分片只保留当前一行，用最小堆在 k 个候选中选出时间戳最小的一条输出，然后从该分片推进一行。堆的大小恒为分片数 k，总时间复杂度 O(N log k)，N 为总行数。

最小实现

约定日志行格式为 时间戳\t级别\t消息，时间戳格式 %Y-%m-%d %H:%M:%S，示例内容均为占位值。

import heapq
from datetime import datetime

TS_FMT = '%Y-%m-%d %H:%M:%S'


def parse_line(line):
    ts_text, level, msg = line.rstrip('\n').split('\t', 2)
    return datetime.strptime(ts_text, TS_FMT), level, msg


def iter_shard(path):
    """逐行读取一个分片，跳过空行，产出 (时间戳, 原始行)。"""
    with open(path, 'r', encoding='utf-8') as f:
        for raw in f:
            if not raw.strip():
                continue
            ts, _level, _msg = parse_line(raw)
            yield ts, raw.rstrip('\r\n')


def merge_shards(paths, out):
    heap = []
    for idx, path in enumerate(paths):
        it = iter_shard(path)
        try:
            ts, line = next(it)
        except StopIteration:
            continue  # 空分片直接跳过
        heapq.heappush(heap, (ts, idx, line, it))

    count = 0
    while heap:
        ts, idx, line, it = heapq.heappop(heap)
        out.write(line + '\n')
        count += 1
        try:
            next_ts, next_line = next(it)
        except StopIteration:
            continue
        heapq.heappush(heap, (next_ts, idx, next_line, it))
    return count

调用入口：

if __name__ == '__main__':
    import sys

    paths = sys.argv[1:]
    with open('merged.log', 'w', encoding='utf-8') as out:
        total = merge_shards(paths, out)
    print('merged lines:', total)

运行方法：

mkdir -p logs
printf '2024-05-01 10:00:00\tINFO\ta1\n2024-05-01 10:00:03\tINFO\ta2\n' > logs/a.log
printf '2024-05-01 10:00:01\tWARN\tb1\n2024-05-01 10:00:03\tINFO\tb2\n' > logs/b.log
python merge_logs.py logs/a.log logs/b.log
cat merged.log

输出中两条 10:00:03 的记录会按传入的分片顺序排列，这正是次级键 idx 的作用。

关键细节

堆元素是四元组 (ts, idx, line, it)。比较时先比时间戳，相同则比 idx。idx 在同一分片内唯一，且堆中每个分片最多只有一个元素，因此前两项已经能区分所有元素，line 与迭代器对象不会再被比较，避免把不可比较的文件对象放进比较链。

迭代器的惰性很关键：iter_shard 是生成器，next() 才真正读下一行，with open 在生成器耗尽时才释放句柄。堆里只保存 k 行，不会随日志总量增长。

边界与限制

• 空分片：next(it) 抛 StopIteration 后直接跳过，不向堆里放哨兵；全部分片为空时循环不执行，返回 0。
• 相同时间戳：只能保证按 idx 稳定，不能还原跨分片的真实先后。如果日志行里带纳秒或单调序号，应把 (ts, seq) 一起放进堆键的前两位。
• 时区：datetime.strptime 得到的是无时区对象，不同分片若时区不同，比较结果没有意义，需要先统一到 UTC。
• 输入有序是前提：单个分片内部必须按时间戳非递减，否则归并结果不全局有序；分片内乱序要先排序。
• 解析失败：当前实现遇到格式不符的行会抛 ValueError 中断。生产中可在 iter_shard 里用 try/except 计数并跳过，但要区分“坏行”和“时间戳倒退”，后者应报错。
• 文件句柄：同时打开 k 个文件。k 很大（例如上千）时可能触及 ulimit -n，需要分批归并成中间文件，再做多轮归并。
• 性能：复杂度 O(N log k)，但 Python 逐行 strptime 往往是瓶颈。若时间戳是固定宽度、可按字典序比较的 ISO 格式，可直接用字符串做键，省掉解析开销。
• 输出顺序不是全局稳定排序，只是“按时间戳 + 分片序号”，不要把它当成原始写入顺序。

小结

k 路归并的核心是把“全局排序”拆成“每个分片各出一个候选”，用大小为 k 的堆持续选出最小时间戳。Python 的 heapq 足够轻量，配合生成器即可做到流式、低内存，适合日志合并、外部排序和多路有序数据汇聚等场景。