使用Uproot写入TTree时如何解决'Basket data should be added to all branches'异常

异常原因分析 #

这个异常是uproot3调用extend方法时，发现传入的分支数据字典与已创建的TTree分支不匹配导致的，常见触发场景：

• 某次调用writeTree时，传入的brs分支列表和初始化TTree时的分支数量/名称不一致
• brs中存在重复分支定义，或部分分支的事件数据长度与其他分支不匹配
• 不同分支逻辑复用了brs变量，导致分支列表被意外修改

结合你的代码片段，核心问题大概率是：elif分支中动态扩展了brs，但else分支直接复用了这个被修改过的brs，导致前后写入的分支列表不一致。

具体修复方案 #

方案1：避免复用brs变量，每次写入重新构建分支列表 #

elif "fakeWeight" in ltree.weight.fields and "em" not in lep:
    print ("extending tree for fakes")
    # 重新构建分支列表，不依赖外部变量
    brs = brlist(lep, True)
    brs.extend([ ["weight", "fakeWeight"], ["weight", "leptonSF_iso"] ])
    utils.writeTree(outf, ltree, outtreename, brs)
else:
    # 每次都重新生成分支列表，确保一致性
    brs = brlist(lep, True)
    utils.writeTree(outf, ltree, outtreename, brs)

方案2：在writeTree中添加数据校验 #

修改utils.py中的writeTree函数，提前校验分支数据的一致性，避免触发uproot3的异常：

def writeTree(outf, ltree, treename, branches):
    # 先收集所有分支数据并检查长度
    branch_data = {}
    target_length = None
    for br in branches:
        data = ltree[br[0]][br[1]]
        branch_data[br[2]] = data
        # 校验所有分支数据长度一致
        if target_length is None:
            target_length = len(data)
        elif len(data) != target_length:
            raise ValueError(f"Branch {br[2]} has mismatched length: {len(data)} vs {target_length}")
    
    # 初始化TTree（如果不存在）
    if treename not in outf:
        outf.mktree(treename, {br[2]: ltree[br[0]][br[1]].dtype for br in branches})
    
    # 执行写入
    outf[treename].extend(branch_data)

方案3：确保分支列表完全一致 #

如果需要固定分支集合，提前定义完整的分支列表，避免动态修改：

# 提前定义基础分支+扩展分支的完整列表
full_brs = brlist(lep, True) + [ ["weight", "fakeWeight"], ["weight", "leptonSF_iso"] ]

elif "fakeWeight" in ltree.weight.fields and "em" not in lep:
    print ("extending tree for fakes")
    utils.writeTree(outf, ltree, outtreename, full_brs)
else:
    utils.writeTree(outf, ltree, outtreename, brlist(lep, True))

额外排查建议 #

• 打印每次写入的brs内容，对比前后分支的数量和名称，确认是否存在不一致
• 检查ltree.weight.fakeWeight和ltree.weight.leptonSF_iso的事件数，是否与其他分支的事件数完全相同
• 确认初始化TTree时的分支定义，与后续写入时的分支列表完全匹配

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