From 7f78c70750659416415944e575615448aaeb6b69 Mon Sep 17 00:00:00 2001 From: Angus Hollands Date: Fri, 7 Apr 2023 18:40:36 +0100 Subject: [PATCH 1/4] fix: `ak.unzip` visits all contents --- src/awkward/operations/ak_unzip.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/awkward/operations/ak_unzip.py b/src/awkward/operations/ak_unzip.py index 7165fe032d..707b8c4e71 100644 --- 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: From 245cd912799864018550a140147c79abe7644061 Mon Sep 17 00:00:00 2001 From: Angus Hollands Date: Fri, 7 Apr 2023 18:42:26 +0100 Subject: [PATCH 2/4] fix: don't return layouts --- src/awkward/operations/ak_unzip.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/awkward/operations/ak_unzip.py b/src/awkward/operations/ak_unzip.py index 707b8c4e71..6106c978b9 100644 --- a/src/awkward/operations/ak_unzip.py +++ b/src/awkward/operations/ak_unzip.py @@ -48,7 +48,7 @@ def _impl(array, highlevel, behavior): def check_for_union(layout, **kwargs): if isinstance(layout, (ak.contents.RecordArray, ak.Record)): - return layout # don't descend into nested records + return # don't descend into nested records elif layout.is_union: for content in layout.contents: From d51d1520acbab1b53641a56a62d35e2794fa235c Mon Sep 17 00:00:00 2001 From: Angus Hollands Date: Fri, 7 Apr 2023 18:52:46 +0100 Subject: [PATCH 3/4] fix: need to return layouts! --- src/awkward/operations/ak_unzip.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/awkward/operations/ak_unzip.py b/src/awkward/operations/ak_unzip.py index 6106c978b9..707b8c4e71 100644 --- a/src/awkward/operations/ak_unzip.py +++ b/src/awkward/operations/ak_unzip.py @@ -48,7 +48,7 @@ def _impl(array, highlevel, behavior): def check_for_union(layout, **kwargs): if isinstance(layout, (ak.contents.RecordArray, ak.Record)): - return # don't descend into nested records + return layout # don't descend into nested records elif layout.is_union: for content in layout.contents: From 5f8ac2ebaf716886c4f5cc822fcee47e27f71158 Mon Sep 17 00:00:00 2001 From: Angus Hollands Date: Fri, 7 Apr 2023 19:47:10 +0100 Subject: [PATCH 4/4] test: add test --- tests/test_2373_unzip_touching.py | 124 ++++++++++++++++++++++++++++++ 1 file changed, 124 insertions(+) create mode 100644 tests/test_2373_unzip_touching.py diff --git a/tests/test_2373_unzip_touching.py b/tests/test_2373_unzip_touching.py new file mode 100644 index 0000000000..dc4d12a196 --- /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!"]