feat(dataset_tools): union forms via DatasetSpec addition (#1478)

[NJManganelli]

🤖 AI text below 🤖

Note

Targets scikit-hep:master but builds on #1579 (RNTuple + switchable backends). Until #1579 merges, this PR's diff also shows #1579's commits; once it lands, GitHub collapses this to just the union-form changes. Review after #1579. Two further PRs stack on this one (metadata extraction, mutable-steps).

Implements the union-form-on-addition feature requested in #1478.

What this adds

When two DatasetSpecs that both carry a saved form are added, the result's form is now the union of their forms (every field appearing in either), instead of the form being discarded. This gives a second route to a dataset union form — combining already-preprocessed specs — without reopening any file. Adding a form-bearing spec to a form-less one raises, since the union could not describe the form-less operand's files.

• DatasetSpec.__add__ delegates to DatasetSpec.union_with(other, sort_fields=False); sort_fields=True recursively sorts record fields so the serialized form is byte-stable regardless of add order (awkward Form equality is field-order-insensitive either way).
• DatasetSpec.canonicalize_form() re-serializes an existing saved form with sorted fields (stable hashing/reproducibility).
• New coffea/dataset_tools/forms.py hosts the union helper (moved out of preprocess.py, which now imports it) plus sort_form_fields, prune_form_fields, and hex field-bitset encode/decode.
• Experimental: each concrete file spec gains experimental_field_bitset, a compact hex bitset over the dataset union form's top-level fields (bit i set ⇒ field i present in that file), populated during pydantic preprocessing. filter_files/limit_files prune the union form back to the surviving files' fields when every remaining file has a bitset (kept as a superset otherwise). Excluded from file-spec equality and from the legacy dict conversions, and prefixed experimental_ deliberately — the encoding may change.

The union merges flat-tuple-like differences: the HLT-trigger-bit / GenModel model-point case #1478 targets, where a field present in only some files becomes an option type (IndexedOptionArray(bool)) so it stays readable (as None) for the files that lack it.

Try it (against the bundled GenModel test files)

Two real CMS NanoAOD files with disjoint GenModel_TChiZH_* model-point flag subsets are added in this PR (file A has GenModel_TChiZH_700_1; file B does not). Highest-level API only — DataGroupSpec → preprocess → spec addition:

from coffea.dataset_tools import DataGroupSpec, preprocess
from coffea.dataset_tools.forms import decode_field_bitset

A = "tests/samples/nano_genmodel_with20_700_1_with0_1100_200_with0_950_400.root"
B = "tests/samples/nano_genmodel_without_700_1_with0_1100_200_with20_950_400.root"

da, _ = preprocess(DataGroupSpec({"s": {"files": {A: "Events"}}}), save_form=True, backend="iterative")
db, _ = preprocess(DataGroupSpec({"s": {"files": {B: "Events"}}}), save_form=True, backend="iterative")

combined = da["s"] + db["s"]            # unions the saved forms; no file reopened
uf = combined.form
print("union / A / B fields:", len(uf.fields), len(da["s"].form.fields), len(db["s"].form.fields))  # 897 804 887

# GenModel_TChiZH_700_1 (present only in file A) is an option type in the union
i = uf.fields.index("GenModel_TChiZH_700_1")
print(type(uf.contents[i]).__name__)   # IndexedOptionForm

# each file's experimental bitset decodes to its own branch set
bs = decode_field_bitset(combined.files[A].experimental_field_bitset, list(uf.fields))
print("700_1 in A:", "GenModel_TChiZH_700_1" in bs)   # True

# filtering back to file A prunes the union form to A's fields
print(set(combined.filter_files(filter_name=".*with20_700_1.*").form.fields) == set(da["s"].form.fields))  # True

Tests

tests/test_dataset_tools_forms.py covers the helpers (union order/associativity, sort, prune, bitset round-trip), __add__/union_with (matches joint preprocessing, equal-form short-circuit, one-sided/no-form cases, order independence), canonicalize_form, filter/limit pruning (including the no-bitset superset fallback and JSON round-trip), and a GenModel regression on the bundled files. At the form level (ROOT and parquet): the union is a superset of both, a flag present in only one file becomes IndexedOptionArray(bool), bitsets decode to each file's branch set, and filtering prunes the form. At the value level (dask read path, where the union form injects the option-typed flag absent from a file): reading file A through the union form gives the same GenModel_TChiZH_700_1 selection as reading A through its own form (20 events — the union does not falsify events passing a model-point mask), and for each model point the combined A+B count equals the per-file sum and the known truth (700_1: 20 from A; 950_400: 20 from B; 1100_200: 0). A companion test pins that selecting on the raw option-type mask keeps the injected None rows and over-selects (40 vs 20), so fill_none(flag, False) is required.

Important

Before converting from draft to review:

• Add the GenModel test files demonstrating GenModel functionality (and the prior union-form / schema bug) and use them in the example + tests/test_dataset_tools_forms.py, alongside the synthetic HLT unit tests. (done — nano_genmodel_* ROOT + parquet fixtures)

---

Developed with Claude Fable 5. Prototype — API may change.