Skip to content

Commit

Permalink
fix: support lists of strings in ak.zip with `optiontype_outside_re…
Browse files Browse the repository at this point in the history
…cord=True` (#2623)

* fix: remove string special-casing

* test: add test for fix

* Also test IndexedOptionArray.

---------

Co-authored-by: Jim Pivarski <[email protected]>
  • Loading branch information
agoose77 and jpivarski authored Aug 9, 2023
1 parent 1cfea2f commit 6c52131
Show file tree
Hide file tree
Showing 2 changed files with 61 additions and 8 deletions.
9 changes: 1 addition & 8 deletions src/awkward/operations/ak_zip.py
Original file line number Diff line number Diff line change
Expand Up @@ -213,14 +213,7 @@ def _impl(
parameters["__record__"] = with_name

def action(inputs, depth, **ignore):
if depth_limit == depth or all(
x.purelist_depth == 1
or (
x.purelist_depth == 2
and x.purelist_parameter("__array__") in ("string", "bytestring")
)
for x in inputs
):
if depth_limit == depth or all(x.purelist_depth == 1 for x in inputs):
# If we want to zip after option types at this depth
if optiontype_outside_record and any(x.is_option for x in inputs):
return None
Expand Down
60 changes: 60 additions & 0 deletions tests/test_2623_optiontype_outside_record_strings.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE

import awkward as ak


def test_ByteMaskedArray():
record = ak.zip(
{
"x": ak.mask(["foo", "bar", "world"], [True, True, False]),
"y": ak.mask(["do", "re", "mi"], [False, True, True]),
},
optiontype_outside_record=True,
)
assert record.to_list() == [None, {"x": "bar", "y": "re"}, None]
assert record.type == ak.types.ArrayType(
ak.types.OptionType(
ak.types.RecordType(
[
ak.types.ListType(
ak.types.NumpyType("uint8", parameters={"__array__": "char"}),
parameters={"__array__": "string"},
),
ak.types.ListType(
ak.types.NumpyType("uint8", parameters={"__array__": "char"}),
parameters={"__array__": "string"},
),
],
["x", "y"],
)
),
3,
None,
)


def test_IndexedOptionArray():
record = ak.zip(
{"x": ["foo", "bar", None], "y": [None, "re", "mi"]},
optiontype_outside_record=True,
)
assert record.to_list() == [None, {"x": "bar", "y": "re"}, None]
assert record.type == ak.types.ArrayType(
ak.types.OptionType(
ak.types.RecordType(
[
ak.types.ListType(
ak.types.NumpyType("uint8", parameters={"__array__": "char"}),
parameters={"__array__": "string"},
),
ak.types.ListType(
ak.types.NumpyType("uint8", parameters={"__array__": "char"}),
parameters={"__array__": "string"},
),
],
["x", "y"],
)
),
3,
None,
)

0 comments on commit 6c52131

Please sign in to comment.