Repository: scikit-hep/awkward

Issue:
Can't concatenate more than 3 types into Union
### Version of Awkward Array

2.0.7

### Description and code to reproduce

When attempting to concatenate more than 3 awkward arrays with different RecordTypes, the concatenation fails because it refuses to create the respective UnionType. Not sure how this relates to #2235, because essentially, I don't want the union, but an OptionType instead. 

```python
>>> ak.concatenate(
    [
        ak.Array(x)
        for x in [
            [{"a": 3}],
            [{"c": 3}],
            [{"d": 3}],
            [{"e": 3}]
        ]
    ]
)
TypeError: while calling

    ak.concatenate(
        arrays = [<Array [{a: 3}] type='1 * {a: int64}'>, <Array [{c: 3}] typ...
        axis = 0
        mergebool = True
        highlevel = True
        behavior = None
    )

Error details: UnionArray cannot contain union-types in its 'contents' (UnionArray); try UnionArray.simplified instead
```
<details>
<summary>Stacktrace</summary>

```pytb
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
~/tmp/ipykernel_993961/1534422542.py in <cell line: 1>()
----> 1 ak.concatenate(
      2     [
      3         ak.Array(x)
      4         for x in [
      5             [{"a": 3}],

~/anaconda3/envs/scirpy_dev/lib/python3.9/site-packages/awkward/operations/ak_concatenate.py in concatenate(arrays, axis, mergebool, highlevel, behavior)
     43         },
     44     ):
---> 45         return _impl(arrays, axis, mergebool, highlevel, behavior)
     46 
     47 

~/anaconda3/envs/scirpy_dev/lib/python3.9/site-packages/awkward/operations/ak_concatenate.py in _impl(arrays, axis, mergebool, highlevel, behavior)
    109             else:
    110                 collapsed = ak._do.mergemany(batch)
--> 111                 batch = [ak._do.merge_as_union(collapsed, x)]
    112 
    113         out = ak._do.mergemany(batch)

~/anaconda3/envs/scirpy_dev/lib/python3.9/site-packages/awkward/_do.py in merge_as_union(one, two)
    307     )
    308 
--> 309     return ak.contents.UnionArray(tags, index, contents, parameters=None)
    310 
    311 

~/anaconda3/envs/scirpy_dev/lib/python3.9/site-packages/awkward/contents/unionarray.py in __init__(self, tags, index, contents, parameters)
     74                 )
     75             if content.is_union:
---> 76                 raise ak._errors.wrap_error(
     77                     TypeError(
     78                         "{0} cannot contain union-types in its 'contents' ({1}); try {0}.simplified instead".format(

TypeError: while calling

    ak.concatenate(
        arrays = [<Array [{a: 3}] type='1 * {a: int64}'>, <Array [{c: 3}] typ...
        axis = 0
        mergebool = True
        highlevel = True
        behavior = None
    )

Error details: UnionArray cannot contain union-types in its 'contents' (UnionArray); try UnionArray.simplified instead
```

</details>


With only three elements, this works: 
```python
>>> ak.concatenate(
    [
        ak.Array(x)
        for x in [
            [{"a": 3}],
            [{"c": 3}],
            [{"d": 3}],
        ]
    ]
)
[{a: 3},
 {c: 3},
 {d: 3}]
----------------
type: 3 * union[
    {
        a: int64
    },
    {
        c: int64
    },
    {
        d: int64
    }
]
```

**Context:**
I've been playing with the new awkward in AnnData Feature with real data and found this when concatenating a list of 32 AnnData objects, each of which contains a list-of-lists-of-records awkward array. 

CC @ivirshup


Discussion:
This error indicates that we've assumed something will be in canonical form, but the initialiser of the layout node finds that this is not the case. For performance reasons (and predicting output types), we don't want to invoke `.simplified` (the canonical constructor) unless we have to. So, this is a bug in Awkward! :tada: 


Resolving patch:
diff --git a/src/awkward/_do.py b/src/awkward/_do.py
--- a/src/awkward/_do.py
+++ b/src/awkward/_do.py
@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 import copy
-from collections.abc import Mapping, MutableMapping
+from collections.abc import Mapping, MutableMapping, Sequence
 from numbers import Integral
 
 import awkward as ak
@@ -270,42 +270,31 @@ def mergeable(one: Content, two: Content, mergebool: bool = True) -> bool:
     return one._mergeable_next(two, mergebool=mergebool)
 
 
-def merge_as_union(one: Content, two: Content) -> ak.contents.UnionArray:
-    mylength = one.length
-    theirlength = two.length
-    tags = ak.index.Index8.empty(
-        mylength + theirlength,
-        one.backend.index_nplike,
-    )
-    index = ak.index.Index64.empty(
-        mylength + theirlength,
-        one.backend.index_nplike,
-    )
-    contents = [one, two]
-    assert tags.nplike is one.backend.index_nplike
-    one._handle_error(
-        one.backend["awkward_UnionArray_filltags_const", tags.dtype.type](
-            tags.data, 0, mylength, 0
-        )
-    )
-    assert index.nplike is one.backend.index_nplike
-    one._handle_error(
-        one.backend["awkward_UnionArray_fillindex_count", index.dtype.type](
-            index.data, 0, mylength
-        )
-    )
-    one._handle_error(
-        one.backend["awkward_UnionArray_filltags_const", tags.dtype.type](
-            tags.data, mylength, theirlength, 1
+def merge_as_union(
+    contents: Sequence[Content], parameters=None
+) -> ak.contents.UnionArray:
+    length = sum([c.length for c in contents])
+    first = contents[0]
+    tags = ak.index.Index8.empty(length, first.backend.index_nplike)
+    index = ak.index.Index64.empty(length, first.backend.index_nplike)
+
+    offset = 0
+    for i, content in enumerate(contents):
+        content._handle_error(
+            content.backend["awkward_UnionArray_filltags_const", tags.dtype.type](
+                tags.data, offset, content.length, i
+            )
         )
-    )
-    one._handle_error(
-        one.backend["awkward_UnionArray_fillindex_count", index.dtype.type](
-            index.data, mylength, theirlength
+        content._handle_error(
+            content.backend["awkward_UnionArray_fillindex_count", index.dtype.type](
+                index.data, offset, content.length
+            )
         )
-    )
+        offset += content.length
 
-    return ak.contents.UnionArray(tags, index, contents, parameters=None)
+    return ak.contents.UnionArray.simplified(
+        tags, index, contents, parameters=parameters
+    )
 
 
 def mergemany(contents: list[Content]) -> Content:
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
@@ -138,20 +138,6 @@ def __init__(self, tags, index, contents, *, parameters=None):
                     )
                 )
 
-        for content in contents[1:]:
-            if contents[0]._mergeable_next(content, mergebool=False):
-                raise ak._errors.wrap_error(
-                    TypeError(
-                        "{0} cannot contain mergeable 'contents' ({1} of {2} and {3} of {4}); try {0}.simplified instead".format(
-                            type(self).__name__,
-                            type(contents[0]).__name__,
-                            repr(str(contents[0].form.type)),
-                            type(content).__name__,
-                            repr(str(content.form.type)),
-                        )
-                    )
-                )
-
         if (
             not (tags.length is unknown_length or index.length is unknown_length)
             and tags.length > index.length
@@ -1174,14 +1160,9 @@ def _mergemany(self, others):
             parameters=parameters,
         )
 
-        if len(tail) == 0:
-            return next
-
-        reversed = tail[0]._reverse_merge(next)
-        if len(tail) == 1:
-            return reversed
-        else:
-            return reversed._mergemany(tail[1:])
+        # `tail` should always be empty for UnionArray's merging strategy
+        assert len(tail) == 0
+        return next
 
     def _fill_none(self, value: Content) -> Content:
         contents = []
@@ -1339,48 +1320,54 @@ def _reduce_next(
         )
 
     def _validity_error(self, path):
-        for i in range(len(self.contents)):
-            if self._backend.nplike.known_data and self.index.length < self.tags.length:
-                return f"at {path} ({type(self)!r}): len(index) < len(tags)"
+        if self._backend.nplike.known_data and self.index.length < self.tags.length:
+            return f"at {path} ({type(self)!r}): len(index) < len(tags)"
 
-            lencontents = self._backend.index_nplike.empty(
-                len(self.contents), dtype=np.int64
-            )
-            if self._backend.nplike.known_data:
-                for i in range(len(self.contents)):
-                    lencontents[i] = self.contents[i].length
+        lencontents = self._backend.index_nplike.empty(
+            len(self.contents), dtype=np.int64
+        )
+        if self._backend.nplike.known_data:
+            for j, _content_j in enumerate(self._contents):
+                lencontents[j] = _content_j.length
+
+        error = self._backend[
+            "awkward_UnionArray_validity",
+            self.tags.dtype.type,
+            self.index.dtype.type,
+            np.int64,
+        ](
+            self.tags.data,
+            self.index.data,
+            self.tags.length,
+            len(self.contents),
+            lencontents,
+        )
 
-            error = self._backend[
-                "awkward_UnionArray_validity",
-                self.tags.dtype.type,
-                self.index.dtype.type,
-                np.int64,
-            ](
-                self.tags.data,
-                self.index.data,
-                self.tags.length,
-                len(self.contents),
-                lencontents,
+        if error.str is not None:
+            if error.filename is None:
+                filename = ""
+            else:
+                filename = " (in compiled code: " + error.filename.decode(
+                    errors="surrogateescape"
+                ).lstrip("\n").lstrip("(")
+            message = error.str.decode(errors="surrogateescape")
+            return 'at {} ("{}"): {} at i={}{}'.format(
+                path, type(self), message, error.id, filename
             )
 
-            if error.str is not None:
-                if error.filename is None:
-                    filename = ""
-                else:
-                    filename = " (in compiled code: " + error.filename.decode(
-                        errors="surrogateescape"
-                    ).lstrip("\n").lstrip("(")
-                message = error.str.decode(errors="surrogateescape")
-                return 'at {} ("{}"): {} at i={}{}'.format(
-                    path, type(self), message, error.id, filename
-                )
+        # Check triangular pairs (i, j) are not mergeable
+        for i, content_i in enumerate(self._contents):
+            for j in range(0, i):
+                content_j = self._contents[j]
+                if ak._do.mergeable(content_i, content_j, mergebool=False):
+                    return f"at {path}: content({i}) is mergeable with content({j})"
 
-            for i in range(len(self.contents)):
-                sub = self.contents[i]._validity_error(path + f".content({i})")
-                if sub != "":
-                    return sub
+        for i, content in enumerate(self._contents):
+            sub = content._validity_error(path + f".content({i})")
+            if sub != "":
+                return sub
 
-            return ""
+        return ""
 
     def _nbytes_part(self):
         result = self.tags._nbytes_part() + self.index._nbytes_part()
diff --git a/src/awkward/operations/ak_concatenate.py b/src/awkward/operations/ak_concatenate.py
--- a/src/awkward/operations/ak_concatenate.py
+++ b/src/awkward/operations/ak_concatenate.py
@@ -102,15 +102,19 @@ def _impl(arrays, axis, mergebool, highlevel, behavior):
             x if isinstance(x, ak.contents.Content) else ak.operations.to_layout([x])
             for x in content_or_others
         ]
-        batch = [content_or_others[0]]
+        batches = [[content_or_others[0]]]
         for x in content_or_others[1:]:
+            batch = batches[-1]
             if ak._do.mergeable(batch[-1], x, mergebool=mergebool):
                 batch.append(x)
             else:
-                collapsed = ak._do.mergemany(batch)
-                batch = [ak._do.merge_as_union(collapsed, x)]
+                batches.append([x])
 
-        out = ak._do.mergemany(batch)
+        contents = [ak._do.mergemany(b) for b in batches]
+        if len(contents) > 1:
+            out = ak._do.merge_as_union(contents)
+        else:
+            out = contents[0]
 
         if isinstance(out, ak.contents.UnionArray):
             out = type(out).simplified(


Test patch:
diff --git a/tests/test_2240_simplify_merge_as_union.py b/tests/test_2240_simplify_merge_as_union.py
new file mode 100644
--- /dev/null
+++ b/tests/test_2240_simplify_merge_as_union.py
@@ -0,0 +1,37 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+
+import numpy as np
+
+import awkward as ak
+
+
+def test_many():
+    result = ak.concatenate(
+        [ak.Array(x) for x in [[{"a": 3}], [{"c": 3}], [{"d": 3}], [{"e": 3}]]]
+    )
+    assert result.tolist() == [{"a": 3}, {"c": 3}, {"d": 3}, {"e": 3}]
+
+
+def test_validity_error_simple():
+    layout = ak.contents.UnionArray(
+        ak.index.Index8(np.array([0, 0, 0, 0, 1, 1, 1, 1, 1], dtype=np.int8)),
+        ak.index.Index64(np.array([0, 1, 2, 3, 0, 1, 2, 3, 4], dtype=np.int64)),
+        [ak.to_layout([1, 2, 3, 4]), ak.to_layout([5, 6, 7, 8, 9])],
+    )
+    assert layout.to_list() == [1, 2, 3, 4, 5, 6, 7, 8, 9]
+    assert "content(1) is mergeable with content(0)" in ak.validity_error(layout)
+
+
+def test_validity_error_complex():
+    layout = ak.contents.UnionArray(
+        ak.index.Index8(np.array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2], dtype=np.int8)),
+        ak.index.Index64(np.array([0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2], dtype=np.int64)),
+        [
+            ak.to_layout([1, 2, 3, 4]),
+            ak.to_layout(["a", "b", "c", "d"]),
+            ak.to_layout([5, 6, 7]),
+        ],
+    )
+    assert layout.to_list() == [1, 2, 3, 4, "a", "b", "c", "d", 5, 6, 7]
+    assert "content(2) is mergeable with content(0)" in ak.validity_error(layout)