`axis=0` concatenation undertouches

[ikrommyd]

import uproot
import dask_awkward as dak

events = uproot.dask({"../coffea/tests/samples/nano_dy.root": "Events"})
lep = dak.concatenate([events.Electron_pt, events.Muon_pt], axis=0)
dak.necessary_columns(lep)

is going to report {'from-uproot-b2a53b89a3ce65faacb366f71b586347': frozenset({'Electron_pt'})}. It should have touched Muon_pt as well.
In the coffea world, the reproducer is

from coffea.nanoevents import NanoEventsFactory
import dask_awkward as dak

events = NanoEventsFactory.from_root({"../coffea/tests/samples/nano_dy.root": "Events"}, mode="dask").events()
lep = dak.concatenate([events.Electron.pt, events.Muon.pt], axis=0)
dak.necessary_columns(lep)

will report

{'from-uproot-327f1c37d0df7f06bac57dc64222280c': frozenset({'Electron_pt',
            'nElectron',
            'nMuon'})}

so it touches only the shape of the muons but not the data.
If you concatenate the whole objects and not just the pt:

from coffea.nanoevents import NanoEventsFactory
import dask_awkward as dak

events = NanoEventsFactory.from_root({"../coffea/tests/samples/nano_dy.root": "Events"}, mode="dask").events()
lep = dak.concatenate([events.Electron, events.Muon], axis=0)
dak.necessary_columns(lep)

will give you

{'from-uproot-327f1c37d0df7f06bac57dc64222280c': frozenset({'Electron_charge',
            'Electron_cleanmask',
            'Electron_convVeto',
            'Electron_cutBased',
            'Electron_cutBased_Fall17_V1',
            'Electron_cutBased_HEEP',
            'Electron_deltaEtaSC',
            'Electron_dr03EcalRecHitSumEt',
            'Electron_dr03HcalDepth1TowerSumEt',
            'Electron_dr03TkSumPt',
            'Electron_dr03TkSumPtHEEP',
            'Electron_dxy',
            'Electron_dxyErr',
            'Electron_dz',
            'Electron_dzErr',
            'Electron_eCorr',
            'Electron_eInvMinusPInv',
            'Electron_energyErr',
            'Electron_eta',
            'Electron_genPartFlav',
            'Electron_genPartIdx',
            'Electron_hoe',
            'Electron_ip3d',
            'Electron_isPFcand',
            'Electron_jetIdx',
            'Electron_jetPtRelv2',
            'Electron_jetRelIso',
            'Electron_lostHits',
            'Electron_mass',
            'Electron_miniPFRelIso_all',
            'Electron_miniPFRelIso_chg',
            'Electron_mvaFall17V1Iso',
            'Electron_mvaFall17V1Iso_WP80',
            'Electron_mvaFall17V1Iso_WP90',
            'Electron_mvaFall17V1Iso_WPL',
            'Electron_mvaFall17V1noIso',
            'Electron_mvaFall17V1noIso_WP80',
            'Electron_mvaFall17V1noIso_WP90',
            'Electron_mvaFall17V1noIso_WPL',
            'Electron_mvaFall17V2Iso',
            'Electron_mvaFall17V2Iso_WP80',
            'Electron_mvaFall17V2Iso_WP90',
            'Electron_mvaFall17V2Iso_WPL',
            'Electron_mvaFall17V2noIso',
            'Electron_mvaFall17V2noIso_WP80',
            'Electron_mvaFall17V2noIso_WP90',
            'Electron_mvaFall17V2noIso_WPL',
            'Electron_mvaTTH',
            'Electron_pdgId',
            'Electron_pfRelIso03_all',
            'Electron_pfRelIso03_chg',
            'Electron_phi',
            'Electron_photonIdx',
            'Electron_pt',
            'Electron_r9',
            'Electron_seedGain',
            'Electron_sieie',
            'Electron_sip3d',
            'Electron_tightCharge',
            'Electron_vidNestedWPBitmap',
            'Electron_vidNestedWPBitmapHEEP',
            'nElectron',
            'nGenPart',
            'nJet',
            'nMuon',
            'nPhoton'})}

So it touches all the electron branches properly and then it touches the shape of genpart, jet, muon and photon somehow? I guess probably due to the global indexes between all of those perhaps but it doesn't touch the muon data at all.

[pfackeldey]

Hm, when running the typetracer by hand and then checking the report everything looks alright:

import uproot
import dask_awkward as dak

events = uproot.dask({"../coffea/tests/samples/nano_dy.root": "Events"})

# create typetracer and an empty report
tt, report = ak.typetracer.typetracer_with_report(events.layout.form_with_key_path(), highlevel=True)
lep = ak.concatenate([tt.Electron_pt, tt.Muon_pt], axis=0)
print(report.data_touched)
# ["('Electron_pt',)",
# "('Electron_pt', None)",
# "('Muon_pt',)",
# "('Muon_pt', None)"]

So on awkward side everything seems to work correctly, I suspect there's a problem then with using this report to determine the columns to be loaded?

[pfackeldey]

@ikrommyd it looks like this under touching behavior was done on purpose: #305

I'm not sure what the thoughts were behind it - I don't understand (yet) how axis=0 concatenation should work at runtime with actual data if that is never touched (as shown by the error that you're seeing)?

Is there any crucial code that relies on axis=0 not touching the columns that should be concatenated?

Maybe @agoose77 or @lgray can comment on this?

edit: Ahhh, I think it's so you can do something like:

from coffea.nanoevents import NanoEventsFactory
import dask_awkward as dak

events = NanoEventsFactory.from_root({"../coffea/tests/samples/nano_dy.root": "Events"}, mode="dask").events()
lep = dak.concatenate([events.Electron, events.Muon], axis=0)
lep.pt.compute()

and not touch all fields, except for the "pt" fields. It's a 'structure-only' manipulation that just merges on a "partition level" in that sense. This does not seem to work however properly either.

edit 2: I'm just catching up with all of this, looks like @agoose77 started a fix already here: #457

[agoose77]

ak.concatenate shouldn't touch axis=0 unless the types change, in which case it needs to touch some (maybe all) partitions.
Right now I think it's broken, but the fix was a bit fiddly. It might be different now that things have stabilised.