Repository: scikit-hep/awkward

Issue:
`ak.unzip` touches all contents whilst typetracing
### Version of Awkward Array

HEAD, after #2370

### Description and code to reproduce

Even after applying #2370 (which fixes a similar bug), the following code:

```python
import json
import awkward as ak


def delta_r2(a, b):
    return (a.eta - b.eta) ** 2 + (a.phi - b.phi) ** 2


fromjson = {
    "class": "RecordArray",
    "fields": ["muon", "jet"],
    "contents": [
        {
            "class": "ListOffsetArray",
            "offsets": "i64",
            "content": {
                "class": "RecordArray",
                "fields": ["pt", "eta", "phi", "crossref"],
                "contents": [
                    {
                        "class": "NumpyArray",
                        "primitive": "int64",
                        "inner_shape": [],
                        "parameters": {},
                        "form_key": "muon_pt!",
                    },
                    {
                        "class": "NumpyArray",
                        "primitive": "int64",
                        "inner_shape": [],
                        "parameters": {},
                        "form_key": "muon_eta!",
                    },
                    {
                        "class": "NumpyArray",
                        "primitive": "int64",
                        "inner_shape": [],
                        "parameters": {},
                        "form_key": "muon_phi!",
                    },
                    {
                        "class": "ListOffsetArray",
                        "offsets": "i64",
                        "content": {
                            "class": "NumpyArray",
                            "primitive": "int64",
                            "inner_shape": [],
                            "parameters": {},
                            "form_key": "muon_crossref_content!",
                        },
                        "parameters": {},
                        "form_key": "muon_crossref_index!",
                    },
                ],
                "parameters": {},
                "form_key": "muon_record!",
            },
            "parameters": {},
            "form_key": "muon_list!",
        },
        {
            "class": "ListOffsetArray",
            "offsets": "i64",
            "content": {
                "class": "RecordArray",
                "fields": ["pt", "eta", "phi", "crossref", "thing1"],
                "contents": [
                    {
                        "class": "NumpyArray",
                        "primitive": "int64",
                        "inner_shape": [],
                        "parameters": {},
                        "form_key": "jet_pt!",
                    },
                    {
                        "class": "NumpyArray",
                        "primitive": "int64",
                        "inner_shape": [],
                        "parameters": {},
                        "form_key": "jet_eta!",
                    },
                    {
                        "class": "NumpyArray",
                        "primitive": "int64",
                        "inner_shape": [],
                        "parameters": {},
                        "form_key": "jet_phi!",
                    },
                    {
                        "class": "ListOffsetArray",
                        "offsets": "i64",
                        "content": {
                            "class": "NumpyArray",
                            "primitive": "int64",
                            "inner_shape": [],
                            "parameters": {},
                            "form_key": "jet_crossref_content!",
                        },
                        "parameters": {},
                        "form_key": "jet_crossref_index!",
                    },
                    {
                        "class": "NumpyArray",
                        "primitive": "int64",
                        "inner_shape": [],
                        "parameters": {},
                        "form_key": "jet_thing1!",
                    },
                ],
                "parameters": {},
                "form_key": "jet_record!",
            },
            "parameters": {},
            "form_key": "jet_list!",
        },
    ],
    "parameters": {},
    "form_key": "outer!",
}

form = ak.forms.from_json(json.dumps(fromjson))

ttlayout, report = ak._nplikes.typetracer.typetracer_with_report(form)

ttarray = ak.Array(ttlayout)

a, b = ak.unzip(ak.cartesian([ttarray.muon, ttarray.jet], axis=1, nested=True))

print("ab>>>", report.data_touched, "\n")

mval = delta_r2(a, b)

print("dr>>>>", report.data_touched, "\n")

mmin = ak.argmin(mval, axis=2)

print("mmin>>", report.data_touched, "\n")

ak.firsts(b[mmin], axis=2).pt

print("pt>>>>", report.data_touched, "\n")
```

produces the following output:

```
ab>>> ['muon_list!', 'jet_list!', 'muon_crossref_index!', 'jet_crossref_index!'] 

dr>>>> ['muon_list!', 'jet_list!', 'muon_crossref_index!', 'jet_crossref_index!', 'muon_eta!', 'jet_eta!', 'muon_phi!', 'jet_phi!'] 

mmin>> ['muon_list!', 'jet_list!', 'muon_crossref_index!', 'jet_crossref_index!', 'muon_eta!', 'jet_eta!', 'muon_phi!', 'jet_phi!'] 

pt>>>> ['muon_list!', 'jet_list!', 'muon_crossref_index!', 'jet_crossref_index!', 'muon_eta!', 'jet_eta!', 'muon_phi!', 'jet_phi!', 'jet_pt!', 'jet_crossref_content!', 'jet_thing1!'] 
```

It's a little mysterious (but not a major pain point) that the "crossref indexes" are touched by the `ak.cartesian` step, which ought to treat the muons and jets as opaque objects. A little mysterious, and maybe related.

But the real problem is represented by the "jet thing" in the final output. The "jet thing" should never have been accessed by any of these operations: `delta_r2`, `ak.argmin`, the `b[mmin]` slice, or `ak.firsts`. It's also a stand-in for dozens of large fields, so it's a real pain point that needs to be addressed.


Discussion:
From @lgray:

> so in particular it's when there is a ListOffsetArray in the records that are being cartesian'd and then `[index]`'d

and

> and it particular I'd also say that tagging the ListOffsetArrays in the cartesian is also wrong? Unless there's some mandatory manipulation that needs to happy to the objects? That makes a lot of assumptions about what those ListOffsetArrays are doing that doesn't make much sense (at least on first thought).

@lgray, if you mean the appearance of `'muon_list!', 'jet_list!'` in `ab>>>`, I expect them to be there because `ak.cartesian` needs to know the structure of the lists at the `axis` value where it's enumerating the Cartesian product. In an imperative implementation, it would be doing `for` loops over those lists, so it definitely has to see the lists and their lengths.

But if you mean the appearance of `'muon_crossref_index!', 'jet_crossref_index!'`, yeah, that's what I called mysterious. It might be a clue.
> But if you mean the appearance of 'muon_crossref_index!', 'jet_crossref_index!', yeah, that's what I called mysterious. It might be a clue.

Yes, I mean this *exactly*. `muon_list!` and `jet_list!` are expected.
The cross refs make little sense that they should be there, they play no role in broadcasting.
I'll also mention that the cross refs can be a pain since they're doubly nested indices, so their on disk / in-memory representations can be pretty large (since they're N^2 in the size of the collection, for two large collections). It would be best to avoid loading them spuriously.

Resolving patch:
diff --git a/src/awkward/operations/ak_unzip.py b/src/awkward/operations/ak_unzip.py
--- a/src/awkward/operations/ak_unzip.py
+++ b/src/awkward/operations/ak_unzip.py
@@ -48,18 +48,15 @@ def _impl(array, highlevel, behavior):
 
     def check_for_union(layout, **kwargs):
         if isinstance(layout, (ak.contents.RecordArray, ak.Record)):
-            pass  # don't descend into nested records
+            return layout  # don't descend into nested records
 
-        elif isinstance(layout, ak.contents.UnionArray):
+        elif layout.is_union:
             for content in layout.contents:
                 if set(ak.operations.fields(content)) != set(fields):
                     raise ak._errors.wrap_error(
                         ValueError("union of different sets of fields, cannot ak.unzip")
                     )
 
-        elif hasattr(layout, "content"):
-            check_for_union(layout.content)
-
     ak._do.recursively_apply(layout, check_for_union, behavior, return_array=False)
 
     if len(fields) == 0:


Test patch:
diff --git a/tests/test_2373_unzip_touching.py b/tests/test_2373_unzip_touching.py
new file mode 100644
--- /dev/null
+++ b/tests/test_2373_unzip_touching.py
@@ -0,0 +1,124 @@
+import awkward as ak
+
+
+def test():
+    form = ak.forms.from_dict(
+        {
+            "class": "RecordArray",
+            "fields": ["muon", "jet"],
+            "contents": [
+                {
+                    "class": "ListOffsetArray",
+                    "offsets": "i64",
+                    "content": {
+                        "class": "RecordArray",
+                        "fields": ["pt", "eta", "phi", "crossref"],
+                        "contents": [
+                            {
+                                "class": "NumpyArray",
+                                "primitive": "int64",
+                                "inner_shape": [],
+                                "parameters": {},
+                                "form_key": "muon_pt!",
+                            },
+                            {
+                                "class": "NumpyArray",
+                                "primitive": "int64",
+                                "inner_shape": [],
+                                "parameters": {},
+                                "form_key": "muon_eta!",
+                            },
+                            {
+                                "class": "NumpyArray",
+                                "primitive": "int64",
+                                "inner_shape": [],
+                                "parameters": {},
+                                "form_key": "muon_phi!",
+                            },
+                            {
+                                "class": "ListOffsetArray",
+                                "offsets": "i64",
+                                "content": {
+                                    "class": "NumpyArray",
+                                    "primitive": "int64",
+                                    "inner_shape": [],
+                                    "parameters": {},
+                                    "form_key": "muon_crossref_content!",
+                                },
+                                "parameters": {},
+                                "form_key": "muon_crossref_index!",
+                            },
+                        ],
+                        "parameters": {},
+                        "form_key": "muon_record!",
+                    },
+                    "parameters": {},
+                    "form_key": "muon_list!",
+                },
+                {
+                    "class": "ListOffsetArray",
+                    "offsets": "i64",
+                    "content": {
+                        "class": "RecordArray",
+                        "fields": ["pt", "eta", "phi", "crossref", "thing1"],
+                        "contents": [
+                            {
+                                "class": "NumpyArray",
+                                "primitive": "int64",
+                                "inner_shape": [],
+                                "parameters": {},
+                                "form_key": "jet_pt!",
+                            },
+                            {
+                                "class": "NumpyArray",
+                                "primitive": "int64",
+                                "inner_shape": [],
+                                "parameters": {},
+                                "form_key": "jet_eta!",
+                            },
+                            {
+                                "class": "NumpyArray",
+                                "primitive": "int64",
+                                "inner_shape": [],
+                                "parameters": {},
+                                "form_key": "jet_phi!",
+                            },
+                            {
+                                "class": "ListOffsetArray",
+                                "offsets": "i64",
+                                "content": {
+                                    "class": "NumpyArray",
+                                    "primitive": "int64",
+                                    "inner_shape": [],
+                                    "parameters": {},
+                                    "form_key": "jet_crossref_content!",
+                                },
+                                "parameters": {},
+                                "form_key": "jet_crossref_index!",
+                            },
+                            {
+                                "class": "NumpyArray",
+                                "primitive": "int64",
+                                "inner_shape": [],
+                                "parameters": {},
+                                "form_key": "jet_thing1!",
+                            },
+                        ],
+                        "parameters": {},
+                        "form_key": "jet_record!",
+                    },
+                    "parameters": {},
+                    "form_key": "jet_list!",
+                },
+            ],
+            "parameters": {},
+            "form_key": "outer!",
+        }
+    )
+
+    ttlayout, report = ak._nplikes.typetracer.typetracer_with_report(form)
+
+    ttarray = ak.Array(ttlayout)
+    pairs = ak.cartesian([ttarray.muon, ttarray.jet], axis=1, nested=True)
+    a, b = ak.unzip(pairs)
+    assert report.data_touched == ["muon_list!", "jet_list!"]