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