如何在PyArrow中针对字段元组编写isin表达式

要筛选(first_name, last_name)元组在给定列表中的行，你可以通过结构体字段匹配或哈希值匹配两种方式实现，以下是具体步骤和代码示例：

1. 基础准备 #

先导入依赖并构造测试数据集（模拟大型数据集场景）：

import pyarrow as pa
import pyarrow.compute as pc
from pyarrow.dataset import dataset

# 构造测试数据并保存为Parquet数据集
test_data = pa.table({
    "first_name": ["Alice", "Bob", "Charlie", "David"],
    "last_name": ["Smith", "Jones", "Smith", "Brown"],
    "age": [30, 25, 35, 40]
})
test_data.write_dataset("user_dataset", format="parquet")
user_ds = dataset("user_dataset")

# 定义需要保留的(first_name, last_name)元组列表
target_pairs = [("Alice", "Smith"), ("Charlie", "Smith")]

2. 方法一：结构体字段匹配（直观可靠） #

将first_name和last_name组合成结构体字段，再直接用isin匹配目标元组转换后的结构体数组：

# 把目标元组转换为PyArrow结构体数组
target_struct_array = pa.array(target_pairs, type=pa.struct([
    ("first_name", pa.string()),
    ("last_name", pa.string())
]))

# 构造过滤表达式：组合字段为结构体，判断是否在目标数组内
filter_expr = pc.struct({
    "first_name": pc.field("first_name"),
    "last_name": pc.field("last_name")
}).isin(target_struct_array)

# 执行过滤并输出结果
filtered_result = user_ds.to_table(filter=filter_expr)
print(filtered_result.to_pandas())

这种方法语义清晰，完全贴合需求逻辑，几乎不会出错，适合大多数场景。

3. 方法二：哈希值匹配（性能优先） #

如果数据集规模极大，结构体匹配性能不够理想，可以通过哈希组合两个字段的值，比较哈希值是否在目标列表中来实现过滤：

# 计算目标元组的哈希值列表
target_hashes = pc.hash_combine(
    pa.array([pair[0] for pair in target_pairs]),
    pa.array([pair[1] for pair in target_pairs])
).to_pylist()

# 构造哈希过滤表达式
hash_filter_expr = pc.hash_combine(
    pc.field("first_name"),
    pc.field("last_name")
).isin(target_hashes)

# 执行过滤
filtered_result_hash = user_ds.to_table(filter=hash_filter_expr)
print(filtered_result_hash.to_pandas())

这种方法性能更优，但要注意：哈希存在极低的冲突概率，对于字符串类型的姓名字段，冲突概率可忽略不计。

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