Python itertools.compress 怎么按布尔掩码筛选数据流

我在处理一批需要“按标记保留”的记录时，最先遇到的不是语法问题，而是输入规模：数据来自生成器，选择标记也在逐条到达，如果先把两边都转成列表，内存峰值和等待时间都会一起上升。这个场景可以用 itertools.compress(data, selectors) 解决：它按位置配对两个可迭代对象，只把选择器为真值的数据项交给下游。

Python 官方文档：https://docs.python.org/3/library/itertools.html#itertools.compress

如果数据和布尔选择器可以顺序消费，且不需要随机访问，itertools.compress 适合做轻量的流式筛选。要牢记一个边界：两条迭代器只要有一条耗尽，输出就立即停止。

先按负载判断：流式筛选还是列表物化

这个选择可以先从业务负载倒推。若数据来自文件读取、分页结果、消息消费或生成器，后续只需要顺序处理保留下来的项目，流式筛选能把中间结果交给下游，避免额外保存整批数据。若数据量很小、需要多次遍历，或者要显示完整的筛选结果，列表反而更容易调试。

用 data 和 selectors 做一次位置配对

compress 不会根据值查找对应关系，也不会判断两个序列的业务主键。它只读取当前位置的数据和当前位置的选择器；选择器经过布尔判断为真时，才把该数据项产出。因此，两个输入的顺序必须由调用方保证。

from itertools import compress


def selected_orders(order_stream, keep_stream):
    # 让数据和选择器保持同一位置，返回可继续向下游消费的迭代器。
    return compress(order_stream, keep_stream)


orders = (f"订单-{number}" for number in range(1, 6))
keep_flags = (number % 2 == 1 for number in range(1, 6))

# 只有真正遍历时才消费两个生成器，避免提前构造完整结果列表。
for order in selected_orders(orders, keep_flags):
    print(order)

这里会依次保留“订单-1”“订单-3”“订单-5”。函数返回的是迭代器，不是已经装满结果的列表；如果下游只需要取前几条，可以继续用 islice 截断，而不必让整个输入跑完。

长度不一致时，停止边界比补齐更重要

官方文档给出的等价逻辑是基于 zip(data, selectors) 的生成器表达式，所以它遵循“任一输入耗尽即停止”的规则。data 比 selectors 长时，尾部数据没有对应的选择器，不会被默认保留；反过来，多出来的选择器也不会单独产生输出。

from itertools import compress


def show_boundary(data, selectors):
    # 用 list 只为了让这个边界示例一次性展示结果，真实大流量场景可直接返回 compress。
    return list(compress(data, selectors))


# selectors 只有三个位置，第四、第五个数据不会被隐式补上 True。
short_flags = [True, False, True]
print(show_boundary(["A", "B", "C", "D", "E"], short_flags))

# data 只有三个位置，多出的选择器没有对应数据，也不会生成额外结果。
long_flags = [True, False, True, True, True]
print(show_boundary(["A", "B", "C"], long_flags))

两次结果都只由前三个位置决定。这个行为适合“选择器明确描述当前批次”的系统；如果业务要求数据和标记长度必须一致，就应该在进入 compress 前单独建立长度或批次一致性约束，而不是期待它替你报警。

把它接入批处理时的推荐结构

在批量导入或消息预处理里，我会把“产生数据”“产生标记”“执行筛选”拆开，让每个来源都能单独替换。数据流可以来自文件解析器，选择器可以来自规则判断或前置服务；只要两边按同一顺序输出，筛选层无需知道底层存储。

from itertools import compress


def accepted_records(records, decisions):
    # 只负责筛选，不在这里缓存全部记录；下游可以继续逐条写入目标系统。
    for record in compress(records, decisions):
        yield record


def read_records():
    # 示例数据源用生成器表达真实的逐条读取，不代表某个具体文件格式。
    for record in ({"id": index, "valid": index % 3 != 0} for index in range(1, 7)):
        yield record


records = read_records()
decisions = (record["valid"] for record in records)
# 注意：这里不能把同一个 records 生成器同时作为 data 和 decisions 的来源。

上面的最后一行特意保留为提醒：一个生成器通常只能顺序消费一次。如果选择器需要读取每条记录的字段，那么应当在同一层生成“记录流”和“选择器流”，或直接使用带条件的生成器表达式，避免 data 与 selectors 争抢同一个输入。

常见误区与落地清单

• 不要把 compress 当作主键连接：它按位置配对，不会检查订单号或 ID。
• 不要假设长度不一致会自动补齐：需要补齐时，应明确使用 zip_longest 或预先校验。
• 不要在需要二次遍历时忘记物化：迭代器消费完后不能从头读取。
• 不要把动态条件硬拆成两条流：条件依赖数据字段时，filter 或生成器表达式通常更直观。
• 接入生产批处理前，明确谁负责发现数据与选择器不同批、不同顺序或提前结束。

相关问题

selectors 可以使用整数吗？

可以。选择器按真值判断，非零整数会被视为真，零会被视为假；生产代码最好保持类型含义清楚，避免把业务数值误当成布尔标记。

compress 和 filter 有什么区别？

compress 接收已经准备好的选择器流，适合外部标记与数据按位置对齐；filter 由函数根据每个数据项计算条件，适合条件本身依赖数据内容的情况。

怎样获得列表结果？

直接使用 list(compress(data, selectors)) 即可，但这会一次性物化所有保留项。数据量较大时，优先让下游继续消费迭代器。

最终的选择可以归结为一句话：数据和选择器都是有序流、只需消费一次，就用 itertools.compress 保持轻量；需要随机访问、重复遍历或严格长度检查，就在筛选前后补上明确的工程约束。