Repository: scikit-hep/awkward

Issue:
`union[X, ?Y]` should be canonicalized to `union[?X, ?Y]`
Although records and unions have a canonical `contents` order that they preserve, it would be correct to say that two records or two unions are equal if they have the same set of `contents`, regardless of order (`zip(fields, contents)` in the case of records, which puts an ordering back in for tuples).

I was trying some examples with our new canonicalization rules to see if we'd have to distinguish between types like

```
A = union[option[X], Y]
```

and

```
B = union[X, option[Y]]
```

I found that an attempt to create

```
C = option[union[X, Y]]
```

forces you to go through `simplified` and makes

```
D = union[option[X], option[Y]]
```

but it's still possible to create an asymmetric case (`A` or `B`). How about if we squash that avenue as well? That is, an attempt to construct

```python
UnionArray(tags, index, [option_type_array, non_option_type_array])
```

would (with a future rule) fail and force you to use `UnionArray.simplified` (already required wherever the codebase makes `UnionArray`, and then `UnionArray.simplified` would wrap all the non-option type arrays in an `UnmaskedArray` if any one of the provided `contents` is option-type.

How about it? Then you'd have fewer cases to consider when checking equality of types: you wouldn't have to worry about whether `A` and `B` should be considered equivalent because they can't exist (they always be `D`).

This comment can become a new issue, if you like that idea.

_Originally posted by @jpivarski in https://github.com/scikit-hep/awkward/pull/2368#discussion_r1160806288_
            


Discussion:

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
@@ -117,6 +117,8 @@ def __init__(self, tags, index, contents, *, parameters=None):
 
         if len(contents) < 2:
             raise TypeError(f"{type(self).__name__} must have at least 2 'contents'")
+
+        n_seen_options = 0
         for content in contents:
             if not isinstance(content, Content):
                 raise TypeError(
@@ -130,6 +132,14 @@ def __init__(self, tags, index, contents, *, parameters=None):
                         type(self).__name__, type(content).__name__
                     )
                 )
+            elif content.is_option:
+                n_seen_options += 1
+
+        if n_seen_options not in {0, len(contents)}:
+            raise TypeError(
+                "{0} must either be comprised of entirely optional contents, or no optional contents; "
+                "try {0}.simplified instead".format(type(self).__name__)
+            )
 
         if (
             not (tags.length is unknown_length or index.length is unknown_length)
@@ -385,6 +395,17 @@ def simplified(
                 "FIXME: handle UnionArray with more than 127 contents"
             )
 
+        # If any contents are options, ensure all contents are options!
+        if any(c.is_option for c in contents):
+            contents = [
+                c
+                if c.is_option
+                else c.to_IndexedOptionArray64()
+                if c.is_indexed
+                else ak.contents.UnmaskedArray(c)
+                for c in contents
+            ]
+
         if len(contents) == 1:
             next = contents[0]._carry(index, True)
             return next.copy(parameters=parameters_union(next._parameters, parameters))


Test patch:
diff --git a/tests/test_1125_to_arrow_from_arrow.py b/tests/test_1125_to_arrow_from_arrow.py
--- a/tests/test_1125_to_arrow_from_arrow.py
+++ b/tests/test_1125_to_arrow_from_arrow.py
@@ -623,9 +623,11 @@ def test_unionarray(tmp_path, extensionarray):
         ak.index.Index8(np.array([0, 0, 1, 1, 1, 0, 1], dtype=np.int8)),
         ak.index.Index64(np.array([0, 1, 3, 2, 1, 2, 0], dtype=np.int64)),
         [
-            ak.contents.NumpyArray(
-                np.array([0.0, 1.1, 2.2]),
-                parameters={"which": "inner1"},
+            ak.contents.UnmaskedArray(
+                ak.contents.NumpyArray(
+                    np.array([0.0, 1.1, 2.2]),
+                    parameters={"which": "inner1"},
+                )
             ),
             ak.contents.ByteMaskedArray(
                 ak.index.Index8(np.array([False, False, True, False]).view(np.int8)),
@@ -674,8 +676,10 @@ def test_unionarray(tmp_path, extensionarray):
             ak.index.Index8(np.array([0, 0, 1, 1, 1, 0, 1], dtype=np.int8)),
             ak.index.Index64(np.array([0, 1, 3, 2, 1, 2, 0], dtype=np.int64)),
             [
-                ak.contents.NumpyArray(
-                    np.array([0.0, 1.1, 2.2]), parameters={"which": "inner1"}
+                ak.contents.UnmaskedArray(
+                    ak.contents.NumpyArray(
+                        np.array([0.0, 1.1, 2.2]), parameters={"which": "inner1"}
+                    )
                 ),
                 ak.contents.ByteMaskedArray(
                     ak.index.Index8(
diff --git a/tests/test_1294_to_and_from_parquet.py b/tests/test_1294_to_and_from_parquet.py
--- a/tests/test_1294_to_and_from_parquet.py
+++ b/tests/test_1294_to_and_from_parquet.py
@@ -740,9 +740,11 @@ def test_unionarray(tmp_path, through, extensionarray):
         ak.index.Index8(np.array([0, 0, 1, 1, 1, 0, 1], dtype=np.int8)),
         ak.index.Index64(np.array([0, 1, 3, 2, 1, 2, 0], dtype=np.int64)),
         [
-            ak.contents.NumpyArray(
-                np.array([0.0, 1.1, 2.2]),
-                parameters={"which": "inner1"},
+            ak.contents.UnmaskedArray(
+                ak.contents.NumpyArray(
+                    np.array([0.0, 1.1, 2.2]),
+                    parameters={"which": "inner1"},
+                )
             ),
             ak.contents.ByteMaskedArray(
                 ak.index.Index8(np.array([False, False, True, False]).view(np.int8)),
@@ -801,8 +803,10 @@ def test_unionarray(tmp_path, through, extensionarray):
             ak.index.Index8(np.array([0, 0, 1, 1, 1, 0, 1], dtype=np.int8)),
             ak.index.Index64(np.array([0, 1, 3, 2, 1, 2, 0], dtype=np.int64)),
             [
-                ak.contents.NumpyArray(
-                    np.array([0.0, 1.1, 2.2]), parameters={"which": "inner1"}
+                ak.contents.UnmaskedArray(
+                    ak.contents.NumpyArray(
+                        np.array([0.0, 1.1, 2.2]), parameters={"which": "inner1"}
+                    )
                 ),
                 ak.contents.ByteMaskedArray(
                     ak.index.Index8(
diff --git a/tests/test_1928_replace_simplify_method_with_classmethod_constructor.py b/tests/test_1928_replace_simplify_method_with_classmethod_constructor.py
--- a/tests/test_1928_replace_simplify_method_with_classmethod_constructor.py
+++ b/tests/test_1928_replace_simplify_method_with_classmethod_constructor.py
@@ -1,6 +1,7 @@
 # BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
 
 import numpy as np
+import pytest
 
 import awkward as ak
 
@@ -40,7 +41,21 @@ def test_indexedoption_of_union():
 
 
 def test_indexedoption_of_union_of_option_1():
-    unionarray = ak.contents.UnionArray(
+    with pytest.raises(
+        TypeError, match=r" must either be comprised of entirely optional contents"
+    ):
+        ak.contents.UnionArray(
+            ak.index.Index8(np.array([0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1], dtype=np.int8)),
+            ak.index.Index64(
+                np.array([0, 1, 0, 2, 1, 2, 3, 3, 4, 5, 4], dtype=np.int64)
+            ),
+            [
+                ak.from_iter([0.0, 1.1, 2.2, 3.3, None, 5.5], highlevel=False),
+                ak.from_iter(["zero", "one", "two", "three", "four"], highlevel=False),
+            ],
+        )
+
+    unionarray = ak.contents.UnionArray.simplified(
         ak.index.Index8(np.array([0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1], dtype=np.int8)),
         ak.index.Index64(np.array([0, 1, 0, 2, 1, 2, 3, 3, 4, 5, 4], dtype=np.int64)),
         [
@@ -66,7 +81,21 @@ def test_indexedoption_of_union_of_option_1():
 
 
 def test_indexedoption_of_union_of_option_2():
-    unionarray = ak.contents.UnionArray(
+    with pytest.raises(
+        TypeError, match=r"must either be comprised of entirely optional contents"
+    ):
+        ak.contents.UnionArray(
+            ak.index.Index8(np.array([0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1], dtype=np.int8)),
+            ak.index.Index64(
+                np.array([0, 1, 0, 2, 1, 2, 3, 3, 4, 5, 4], dtype=np.int64)
+            ),
+            [
+                ak.from_iter([0.0, 1.1, 2.2, 3.3, 4.4, 5.5], highlevel=False),
+                ak.from_iter(["zero", None, "two", "three", "four"], highlevel=False),
+            ],
+        )
+
+    unionarray = ak.contents.UnionArray.simplified(
         ak.index.Index8(np.array([0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1], dtype=np.int8)),
         ak.index.Index64(np.array([0, 1, 0, 2, 1, 2, 3, 3, 4, 5, 4], dtype=np.int64)),
         [
diff --git a/tests/test_2236_merge_union_of_records_option.py b/tests/test_2236_merge_union_of_records_option.py
--- a/tests/test_2236_merge_union_of_records_option.py
+++ b/tests/test_2236_merge_union_of_records_option.py
@@ -58,9 +58,11 @@ def test_option():
     assert z.type == ak.types.ArrayType(
         ak.types.UnionType(
             [
-                ak.types.RecordType(
-                    [ak.types.NumpyType("int64"), ak.types.NumpyType("int64")],
-                    ["a", "b"],
+                ak.types.OptionType(
+                    ak.types.RecordType(
+                        [ak.types.NumpyType("int64"), ak.types.NumpyType("int64")],
+                        ["a", "b"],
+                    )
                 ),
                 ak.types.OptionType(
                     ak.types.RecordType(
@@ -104,9 +106,11 @@ def test_option_unmasked():
     assert z.type == ak.types.ArrayType(
         ak.types.UnionType(
             [
-                ak.types.RecordType(
-                    [ak.types.NumpyType("int64"), ak.types.NumpyType("int64")],
-                    ["a", "b"],
+                ak.types.OptionType(
+                    ak.types.RecordType(
+                        [ak.types.NumpyType("int64"), ak.types.NumpyType("int64")],
+                        ["a", "b"],
+                    )
                 ),
                 ak.types.OptionType(
                     ak.types.RecordType(