Python itertools 实战:用 groupby、chain、islice 优雅处理大数据流
写数据处理脚本时,你是不是经常这么干:先把整个文件读进一个 list,再用 for 循环套 if 分组、拼接、切片?数据量小没事,一旦文件几个 G,内存直接爆掉。Python 标准库里的 itertools 就是专门解决这个的------它让你像流水线一样处理数据,内存占用几乎恒定。这篇挑三个最实用的函数,配真实场景讲清楚怎么用、坑在哪。
场景一:groupby 分组,但它有个致命前提
假设你有一份访问日志,要按用户 ID 统计每人的访问次数。很多人第一反应是用 defaultdict,但如果日志已经按用户排好序,groupby 更省内存------它不需要把所有数据攒在字典里。
先看一个新手必踩的坑:
python
from itertools import groupby
logs = [
("alice", "/home"),
("alice", "/profile"),
("bob", "/home"),
("alice", "/settings"), # 注意:alice 又出现了,但和前面不连续
]
for user, group in groupby(logs, key=lambda x: x[0]):
print(user, len(list(group)))输出是:
alice 2
bob 1
alice 1alice 被拆成了两组!因为 groupby 只会合并「相邻」的相同 key ,它不是全局分组。这是它和 SQL 的 GROUP BY 最大的区别,也是最容易出错的地方。
正确用法是先排序,再分组:
python
from itertools import groupby
logs = [
("alice", "/home"),
("alice", "/profile"),
("bob", "/home"),
("alice", "/settings"),
]
# 关键:排序的 key 必须和 groupby 的 key 一致
logs.sort(key=lambda x: x[0])
for user, group in groupby(logs, key=lambda x: x[0]):
paths = [path for _, path in group]
print(f"{user}: {len(paths)} 次访问 -> {paths}")输出正确了:
alice: 3 次访问 -> ['/home', '/profile', '/settings']
bob: 1 次访问 -> ['/home']还有一个隐蔽坑:group 是个迭代器,只能消费一次 。如果你先 len(list(group)) 再想遍历它,第二次拿到的就是空的。要么存成 list 复用,要么一次性处理完。
场景二:chain 把多个数据源接成一条流
你有三个日志文件要合并处理,朴素写法是读进三个 list 再相加:
python
# 朴素写法:三个文件全读进内存再拼接,内存 = 三个文件之和
all_lines = list(open("a.log")) + list(open("b.log")) + list(open("c.log"))
for line in all_lines:
process(line)chain 能把多个可迭代对象「首尾相接」成一个,全程惰性求值,任意时刻内存里只有一行:
python
from itertools import chain
files = ["a.log", "b.log", "c.log"]
# chain.from_iterable 接收「可迭代对象的可迭代对象」
# 每个文件对象本身就是行的迭代器,读一行处理一行
lines = chain.from_iterable(open(f) for f in files)
for line in lines:
process(line.strip())chain(a, b, c) 是把参数一个个接起来,chain.from_iterable(iterables) 则是把一个装着多个可迭代对象的序列拉平------处理动态数量的数据源时用后者,不用手动 * 解包。
场景三:islice 给无限流「切一刀」
itertools 里很多函数返回的是无限迭代器,比如 count()(从 0 数到无穷)。想取前 N 个怎么办?你不能用 list(count())[:10],那会把内存撑爆。用 islice------它是「迭代器版的切片」:
python
from itertools import count, islice
# count(1, 2) 生成 1, 3, 5, 7... 无限奇数流
odds = count(1, 2)
# islice 取前 5 个,不会真的去遍历无限流
first_five = list(islice(odds, 5))
print(first_five) # [1, 3, 5, 7, 9]islice 最实用的场景是分页读大文件。比如跳过表头、每次取 100 行做批处理:
python
from itertools import islice
def read_batches(file_path, batch_size=100):
with open(file_path) as f:
next(f) # 跳过表头
while True:
# islice 从文件迭代器里切出 batch_size 行
batch = list(islice(f, batch_size))
if not batch: # 切不出东西了,文件读完
break
yield batch
for batch in read_batches("huge.csv", batch_size=100):
bulk_insert(batch) # 每 100 行批量入库一次,内存恒定注意 islice 不支持负索引(islice(it, -1) 会报错),因为它无法预知迭代器长度。需要「最后 N 个」时得换 collections.deque(it, maxlen=n)。
组合起来:一条真正的处理流水线
三个函数串起来,就是一条不落地、内存恒定的数据管道。比如「合并多个已排序日志、按用户分组、只看前 3 个用户」:
python
from itertools import chain, groupby, islice
def merged_sorted_logs(files):
# 假设每个文件内部已按 user 排好序,这里简化为直接 chain
# 生产环境跨文件全局有序应该用 heapq.merge
lines = chain.from_iterable(open(f) for f in files)
return (line.strip().split(",") for line in lines)
rows = merged_sorted_logs(["u1.log", "u2.log"])
# 按第一列(user)分组,只取前 3 组
for user, group in islice(groupby(rows, key=lambda r: r[0]), 3):
print(user, sum(1 for _ in group))整条链路里没有一次 list(全部数据),内存只跟「当前处理的一行 + 一组」有关。这就是 itertools 的价值:用惰性迭代把「空间换时间」翻转成「几乎不占空间」。
小结
groupby只合并相邻元素,不是全局分组------用前必须先按同样的 key 排序,否则同一 key 会被拆成多组。返回的 group 是一次性迭代器。chain/chain.from_iterable把多个数据源惰性接成一条流,合并多文件时内存恒定;动态数量的源用from_iterable。islice是迭代器版切片 ,能安全地从无限流或大文件里取前 N 个、做分页批处理;不支持负索引,要「最后 N 个」用deque(maxlen=n)。- 三者可自由串联,组成不落地的处理流水线。
一句话记忆:itertools 处理的是「流」不是「集合」,凡是能一行行流着处理的,就别整个读进内存。