PyArrow表操作：将Parquet嵌套数据中的浮点数组列展开为独立列

嘿，我懂你为啥想从Polars转用PyArrow——追求更原生的Arrow性能嘛！针对你这种嵌套Parquet数据里的浮点数组列（也就是signals）展开成独立列的需求，我给你整理了一套实操步骤，结合你给出的Schema来讲解：

首先先明确你给出的原始数据Schema，方便后续操作对齐：

import pyarrow as pa

document_fields = pa.struct(
    [
        pa.field("url", pa.utf8()),
        pa.field("title", pa.utf8(), nullable=True),
        pa.field("content", pa.utf8(), nullable=True),
        pa.field("relevance", pa.int32()),
        pa.field("signals", pa.list_(pa.float32()), nullable=True),
        pa.field("click_proba", pa.float64()),
        pa.field("embedding", pa.binary()),
        pa.field("title_embedding", pa.binary()),
        pa.field("content_embedding", pa.binary()),
        pa.field("n_seen", pa.int32())
    ]
)

# 补全你没写完的顶层Schema，假设数据是包裹在`documents`结构体下的
PYARROW_SERP_SCHEMA = pa.schema([pa.field("documents", document_fields)])

接下来是具体的展开操作步骤：

1. 读取Parquet文件到PyArrow表 #

用PyArrow原生的Parquet读取工具加载数据，指定Schema可以确保数据类型完全符合预期：

import pyarrow.parquet as pq

# 替换成你的Parquet文件路径
table = pq.read_table("your_serp_data.parquet", schema=PYARROW_SERP_SCHEMA)

2. 提取目标数组列 #

从嵌套的documents结构体中取出signals列，这是我们要处理的核心列：

signals_col = table["documents"]["signals"]

3. 确定数组的最大长度（针对可变长度数组） #

如果你的signals数组每条数据长度不一致，先统计最大长度，确保我们能覆盖所有元素：

# 过滤掉null值后计算最大长度
max_signal_length = max(
    len(arr) for arr in signals_col.to_pylist() 
    if arr is not None
)

如果你的signals是固定长度数组，直接用已知的长度即可，跳过这一步

4. 展开数组为独立列 #

遍历数组索引，用PyArrow的list_element()方法提取每个位置的元素，生成新的列，同时保留原有的其他列：

new_columns = []

# 先把除了signals之外的原有列都加进来
for field in document_fields:
    if field.name != "signals":
        # 提取并重命名列（确保列名正确）
        new_columns.append(table["documents"][field.name].rename(field.name))

# 逐个提取signals的元素作为新列，命名为signal_0、signal_1...
for idx in range(max_signal_length):
    # list_element(idx)会自动处理数组长度不足的情况，填充null
    signal_col = signals_col.list_element(idx).rename(f"signal_{idx}")
    new_columns.append(signal_col)

5. 生成新表并保存 #

把所有列组合成新的PyArrow表，之后可以直接保存为Parquet文件或者继续后续处理：

# 创建新表
flattened_table = pa.Table.from_arrays(
    new_columns, 
    names=[col.name for col in new_columns]
)

# 保存为新的Parquet文件
pq.write_table(flattened_table, "flattened_serp_data.parquet")

一些注意事项 #

• 性能优势：PyArrow的这些操作都是矢量化执行的，完全利用Arrow的内存优势，比循环处理Python对象快得多，符合你追求性能的需求
• 空值处理：如果某条数据的signals数组长度小于最大长度，list_element()会自动填充null，不用担心数据丢失
• 固定长度优化：如果你的signals是固定长度数组，直接用固定的索引范围（比如range(5)），不用计算最大长度，效率会更高

备注：内容来源于stack exchange，提问作者Durand