如何保存含不同字段的拼接Awkward Array至Parquet文件？

处理电子和μ子数据时，筛选同味异号双轻子对（e+e-或μ+μ-），使用Awkward 1.10.3拼接双轻子数组后，保存到Parquet文件时触发ArrowNotImplementedError。原因是电子数组含MVAid字段、μ子数组含tightId字段，拼接后形成union类型，而Arrow暂不支持该类型的Parquet转换。

复现代码（简化版）：

import awkward as ak
import vector
vector.register_awkward()

elec = ak.Array([[{'pt' : 1, 'phi' : 2.2, 'eta' : 1.5, 'mass' : 1, 'MVAid' : True}, {'pt' : 1, 'phi' : 2.2, 'eta' : 1.5, 'mass' : 1, 'MVAid' : True}]])
mu = ak.Array([[{'pt' : 2, 'phi' : 2.2, 'eta' : 2.5, 'mass' : 2, 'tightId' : True}, {'pt' : 2, 'phi' : 2.2, 'eta' : 2.5, 'mass' : 2, 'tightId' : True}]])

electrons_4V = ak.Array(elec, with_name="Momentum4D")
ee_pairs = ak.combinations(electrons_4V, 2, fields=["LeadLepton", "SubleadLepton"])

muons_4V = ak.Array(mu, with_name="Momentum4D")
mm_pairs = ak.combinations(muons_4V, 2, fields=["LeadLepton", "SubleadLepton"])

dileptons = ak.concatenate([ee_pairs, mm_pairs], axis=1)
ak.to_parquet(dileptons, 'test.parquet')

核心错误信息：

ArrowNotImplementedError: Unhandled type for Arrow to Parquet schema conversion: dense_union<0: extension> not null=0, 1: extension> not null=1>

1. 统一轻子字段结构 #

给电子数组补充tightId字段，给μ子数组补充MVAid字段，使两者结构完全一致，拼接后不会生成union类型。

示例代码：

import awkward as ak
import vector
vector.register_awkward()

elec = ak.Array([[{'pt' : 1, 'phi' : 2.2, 'eta' : 1.5, 'mass' : 1, 'MVAid' : True}, {'pt' : 1, 'phi' : 2.2, 'eta' : 1.5, 'mass' : 1, 'MVAid' : True}]])
mu = ak.Array([[{'pt' : 2, 'phi' : 2.2, 'eta' : 2.5, 'mass' : 2, 'tightId' : True}, {'pt' : 2, 'phi' : 2.2, 'eta' : 2.5, 'mass' : 2, 'tightId' : True}]])

# 统一字段：给电子加tightId（设为None或默认值），给μ子加MVAid
elec['tightId'] = None
mu['MVAid'] = None

electrons_4V = ak.Array(elec, with_name="Momentum4D")
ee_pairs = ak.combinations(electrons_4V, 2, fields=["LeadLepton", "SubleadLepton"])

muons_4V = ak.Array(mu, with_name="Momentum4D")
mm_pairs = ak.combinations(muons_4V, 2, fields=["LeadLepton", "SubleadLepton"])

dileptons = ak.concatenate([ee_pairs, mm_pairs], axis=1)
ak.to_parquet(dileptons, 'test.parquet')  # 正常保存

2. 提取公共字段后拼接 #

只保留电子和μ子的公共字段（pt、phi、eta、mass），丢弃各自独有的字段，确保拼接后的数组结构一致。

示例代码：

import awkward as ak
import vector
vector.register_awkward()

elec = ak.Array([[{'pt' : 1, 'phi' : 2.2, 'eta' : 1.5, 'mass' : 1, 'MVAid' : True}, {'pt' : 1, 'phi' : 2.2, 'eta' : 1.5, 'mass' : 1, 'MVAid' : True}]])
mu = ak.Array([[{'pt' : 2, 'phi' : 2.2, 'eta' : 2.5, 'mass' : 2, 'tightId' : True}, {'pt' : 2, 'phi' : 2.2, 'eta' : 2.5, 'mass' : 2, 'tightId' : True}]])

# 提取公共字段
elec_common = elec[['pt', 'phi', 'eta', 'mass']]
mu_common = mu[['pt', 'phi', 'eta', 'mass']]

electrons_4V = ak.Array(elec_common, with_name="Momentum4D")
ee_pairs = ak.combinations(electrons_4V, 2, fields=["LeadLepton", "SubleadLepton"])

muons_4V = ak.Array(mu_common, with_name="Momentum4D")
mm_pairs = ak.combinations(muons_4V, 2, fields=["LeadLepton", "SubleadLepton"])

dileptons = ak.concatenate([ee_pairs, mm_pairs], axis=1)
ak.to_parquet(dileptons, 'test.parquet')  # 正常保存

3. 转换Union类型为Struct类型（带类型标记） #

给每个轻子添加类型标记字段（如lepton_type），将Union类型转换为Struct类型，使Parquet能够识别处理。

示例代码：

import awkward as ak
import vector
vector.register_awkward()

elec = ak.Array([[{'pt' : 1, 'phi' : 2.2, 'eta' : 1.5, 'mass' : 1, 'MVAid' : True}, {'pt' : 1, 'phi' : 2.2, 'eta' : 1.5, 'mass' : 1, 'MVAid' : True}]])
mu = ak.Array([[{'pt' : 2, 'phi' : 2.2, 'eta' : 2.5, 'mass' : 2, 'tightId' : True}, {'pt' : 2, 'phi' : 2.2, 'eta' : 2.5, 'mass' : 2, 'tightId' : True}]])

# 添加类型标记
elec['lepton_type'] = "electron"
mu['lepton_type'] = "muon"

# 转换为struct类型，缺失字段自动设为None
all_fields = set(ak.fields(elec)) | set(ak.fields(mu))
elec_struct = ak.zip({k: elec[k] if k in ak.fields(elec) else None for k in all_fields})
mu_struct = ak.zip({k: mu[k] if k in ak.fields(mu) else None for k in all_fields})

electrons_4V = ak.Array(elec_struct, with_name="Momentum4D")
ee_pairs = ak.combinations(electrons_4V, 2, fields=["LeadLepton", "SubleadLepton"])

muons_4V = ak.Array(mu_struct, with_name="Momentum4D")
mm_pairs = ak.combinations(muons_4V, 2, fields=["LeadLepton", "SubleadLepton"])

dileptons = ak.concatenate([ee_pairs, mm_pairs], axis=1)
ak.to_parquet(dileptons, 'test.parquet')  # 正常保存

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