如何使用pyarrow.compute查询PyArrow表结构体/列表字段实现条件过滤

测试数据构造 #

先复现问题中的示例表：

import pyarrow as pa
import pyarrow.compute as pc

A = pa.array([["X", "Y"], ["Y", "Z"]])
B = pa.array([[1, 3], [5, 6]])
tbl = pa.table({"A": A, "B": B})

两个核心需求实现 #

直接使用pyarrow.compute内置的列表专用操作函数即可实现，不需要额外做复杂结构转换：

• 需求1：获取每行列A中元素'Y'的索引位置
  用pc.list_index逐行查找元素下标，不存在匹配值时返回null，示例中返回结果正好是[1, 0]：

  y_pos = pc.list_index(tbl["A"], value="Y")
  

• 需求2：按索引提取列B对应位置的值
  pc.list_element支持传入位置数组，逐行提取对应下标的元素，示例中返回结果为[3, 5]：

  y_val = pc.list_element(tbl["B"], y_pos)
  

实现等价SQL筛选逻辑 #

要实现SELECT * FROM table WHERE A.Y = 5的效果，直接用匹配到的Y值生成布尔掩码过滤表即可，最终会返回符合条件的第二行数据：

filter_mask = pc.equal(y_val, 5)
res = tbl.filter(filter_mask)

说明：之前尝试的pc.index_in、pc.equals直接作用于列表列时，是把整个列表作为单个值做比较，不会遍历列表内元素，因此无法得到预期结果；直接将A、B两列拼接为StructArray的思路不适用，因为StructArray是固定命名字段的结构，不是键值对映射结构，无法实现按A列元素值作为键取B列值的逻辑。

实现逆堆叠（unstack）得到关联结构 #

要得到按Y值关联拆分后的表，逐行过滤掉Y所在位置的剩余元素即可，最终输出结构和预期完全一致：

# 生成每行的下标序列
row_len = pc.list_value_length(tbl["A"])
max_len = max(row_len.as_py())
base_pos = pa.array([list(range(max_len)) for _ in range(len(tbl))])
row_pos = pc.list_slice(base_pos, length=row_len)
# 过滤掉Y所在下标，得到剩余元素位置
keep_pos = pc.apply_filters(
    row_pos,
    pc.not_equal(row_pos.flatten(), pc.repeat(y_pos, row_len))
)
# 提取剩余元素拼接为结果表
unstack_tbl = pa.table({
    "A": pc.list_element(tbl["A"], keep_pos),
    "B": pc.list_element(tbl["B"], keep_pos),
    "Y": y_val
})

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