Repository: scikit-hep/awkward

Issue:
Broadcasting of unions overtouches data
### Version of Awkward Array

master

### Description and code to reproduce

```python3
from coffea.nanoevents import NanoEventsFactory, NanoAODSchema
import awkward as ak
import dask_awkward as dak

if __name__ == "__main__":
    NanoAODSchema.error_missing_event_ids = False

    events = NanoEventsFactory.from_root(
        {"tests/samples/nano_dy.root": {"object_path": "Events", "steps": [[0,40]]}},
        metadata={"dataset": "nano_dy"},
        schemaclass=NanoAODSchema,
    ).events()

    events["leptons"] = ak.concatenate(
        [events.Electron, events.Muon],
        axis=1,
    )

    sum_p4 = events.leptons + events.leptons

    print(dak.necessary_columns(sum_p4.pt))

    print(sum_p4.pt.compute())
```

While this does function now and does not yield incorrect data it reads a very large number of fields when it should need 10 (the counts of electrons, muons, and the four-momentum components of each). The current performance is rather suboptimal.

The code presently outputs:
```
(coffea-dev) lgray@Lindseys-MacBook-Pro coffea % python -i concat_issue_nano.py
/Users/lgray/coffea-dev/coffea/src/coffea/nanoevents/schemas/nanoaod.py:243: RuntimeWarning: Missing cross-reference index for FatJet_genJetAK8Idx => GenJetAK8
  warnings.warn(
{'from-uproot-ad7806f24d5481e9a913897472bb24fa': frozenset({'Electron_cutBased', 'Muon_miniPFRelIso_chg', 'Electron_charge', 'Muon_highPtId', 'Electron_cutBased_HEEP', 'Muon_isGlobal', 'Muon_nStations', 'Muon_softMvaId', 'Electron_mvaFall17V1noIso_WP90', 'nGenPart', 'Electron_vidNestedWPBitmap', 'Muon_tkRelIso', 'Electron_pdgId', 'Electron_energyErr', 'Muon_pfIsoId', 'Electron_mvaFall17V1Iso_WPL', 'Muon_dxyErr', 'Electron_mvaFall17V1Iso_WP90', 'Muon_miniIsoId', 'Electron_r9', 'Electron_dr03TkSumPt', 'Electron_dr03HcalDepth1TowerSumEt', 'Electron_mvaFall17V2noIso', 'Electron_mvaFall17V2Iso_WP80', 'Electron_isPFcand', 'Electron_cutBased_Fall17_V1', 'Muon_dz', 'Muon_isTracker', 'Electron_hoe', 'Electron_mvaFall17V1noIso_WP80', 'nFsrPhoton', 'Electron_mvaTTH', 'Electron_photonIdx', 'Electron_dxy', 'Muon_nTrackerLayers', 'Muon_segmentComp', 'Muon_genPartFlav', 'Muon_tightCharge', 'Electron_dz', 'Muon_dzErr', 'Muon_tunepRelPt', 'Muon_dxy', 'Electron_genPartFlav', 'Muon_mediumId', 'Muon_looseId', 'Electron_mvaFall17V1Iso_WP80', 'Electron_jetRelIso', 'Electron_lostHits', 'Electron_dr03EcalRecHitSumEt', 'Electron_vidNestedWPBitmapHEEP', 'Muon_mvaLowPt', 'Muon_mediumPromptId', 'Muon_softMva', 'nPhoton', 'Electron_miniPFRelIso_chg', 'Muon_fsrPhotonIdx', 'Muon_mvaTTH', 'Muon_pfRelIso04_all', 'Muon_mvaId', 'Electron_mvaFall17V2Iso', 'Muon_ip3d', 'Electron_phi', 'Muon_inTimeMuon', 'Electron_eInvMinusPInv', 'Muon_mass', 'Muon_jetPtRelv2', 'Muon_sip3d', 'Muon_tightId', 'Electron_eCorr', 'Electron_dxyErr', 'Electron_dr03TkSumPtHEEP', 'Muon_jetIdx', 'Electron_pfRelIso03_chg', 'Electron_mvaFall17V1noIso_WPL', 'Electron_mvaFall17V2Iso_WP90', 'Muon_pfRelIso03_chg', 'Electron_pfRelIso03_all', 'Muon_pfRelIso03_all', 'Electron_eta', 'Electron_convVeto', 'Electron_mvaFall17V2Iso_WPL', 'Muon_pt', 'Electron_mvaFall17V1noIso', 'nElectron', 'Electron_sip3d', 'Electron_mvaFall17V1Iso', 'Muon_triggerIdLoose', 'Electron_pt', 'Electron_mvaFall17V2noIso_WPL', 'Muon_multiIsoId', 'Electron_tightCharge', 'Muon_phi', 'Muon_ptErr', 'Muon_eta', 'Electron_miniPFRelIso_all', 'nMuon', 'Electron_deltaEtaSC', 'Muon_jetRelIso', 'Muon_isPFcand', 'Electron_mvaFall17V2noIso_WP90', 'nJet', 'Electron_jetPtRelv2', 'Electron_sieie', 'Electron_cleanmask', 'Muon_tkIsoId', 'Muon_cleanmask', 'Electron_mass', 'Electron_ip3d', 'Muon_genPartIdx', 'Electron_genPartIdx', 'Electron_seedGain', 'Muon_miniPFRelIso_all', 'Muon_charge', 'Electron_dzErr', 'Electron_mvaFall17V2noIso_WP80', 'Muon_pdgId', 'Muon_softId', 'Electron_jetIdx'})}
[[], [59.1], [120, 103], [21.3, 17.2], [154, ...], ..., [], [15.4], [], []]
```

Apologies for the coffea reproducer.


Discussion:
This was working as expected in a prior version of awkward/dask-awkward (query 8 here run for CHEP):
https://github.com/CoffeaTeam/coffea-benchmarks/blob/master/coffea-adl-benchmarks.ipynb

I know things have changed greatly, but I'm sure we can get back to that state. This issue may likely hang between awkward and dask-awkward. @douglasdavis FYI
What's happening here is that we broadcast through the union, which necessarily involves projecting out each content before recursing. This projection is eager, so we end up touching everything below-and-including each content of the union. Once we recurse through each content, we find the same ufunc overload for `np.add` (because leptons are leptons), and the resulting union over the result therefore simplifies away.

@jpivarski this concerns the choice for `allow_lazy=False` vs `allow_lazy=True` in `UnionArray.project`. This code predates my time on Awkward, and I know you usually have well-formed ideas about whether it's reasonable to create an index or not. There are two solutions: either we change `UnionArray.project` to use `allow_lazy=True` internally, *or* we add a new `UnionArray.project_as_indexed` that explicitly sets `allow_lazy=True`. I think the former is fine; it makes a lot of sense to me that we should avoid the carry; it's trivial to project an `IndexedArray`.

N.B.: presently, there is an implicit dependency upon the union formed by `ak.concatenate`.

Resolving patch:
diff --git a/src/awkward/contents/unionarray.py b/src/awkward/contents/unionarray.py
--- a/src/awkward/contents/unionarray.py
+++ b/src/awkward/contents/unionarray.py
@@ -708,7 +708,7 @@ def project(self, index):
         nextcarry = ak.index.Index64(
             tmpcarry.data[: lenout[0]], nplike=self._backend.index_nplike
         )
-        return self._contents[index]._carry(nextcarry, False)
+        return self._contents[index]._carry(nextcarry, True)
 
     @staticmethod
     def regular_index(
diff --git a/src/awkward/operations/ak_enforce_type.py b/src/awkward/operations/ak_enforce_type.py
--- a/src/awkward/operations/ak_enforce_type.py
+++ b/src/awkward/operations/ak_enforce_type.py
@@ -947,7 +947,7 @@ def _recurse_union_non_union(
 
         # Index over them
         index = ak.index.Index64(index_data)
-        return ak.contents.IndexedArray(index, next_content)
+        return ak.contents.IndexedArray.simplified(index, next_content)
     else:
         # Find the first content whose type equals the given type
         for tag, content in enumerate(layout.contents):  # noqa: B007


Test patch: