Repository: scikit-hep/awkward

Issue:
IndexedArrays can be absorbed into a UnionArray's index (unless categorical)
### Version of Awkward Array

HEAD

### Description and code to reproduce

Currently, UnionArrays are allowed to contain IndexedArrays:

```python
array = ak.Array(
    ak.contents.UnionArray.simplified(
        ak.index.Index8([0, 0, 0, 1, 1, 1, 1, 1]),
        ak.index.Index64([0, 1, 2, 0, 1, 2, 3, 4]),
        [
            ak.from_iter(["one", "two", "three"], highlevel=False),
            ak.contents.IndexedArray(
                ak.index.Index64([4, 3, 2, 1, 0]),
                ak.from_iter([1.1, 2.2, 3.3, 4.4, 5.5], highlevel=False),
            )
        ],
    )
)
```

has layout

```python
>>> array.layout
<UnionArray len='8'>
    <tags><Index dtype='int8' len='8'>[0 0 0 1 1 1 1 1]</Index></tags>
    <index><Index dtype='int64' len='8'>[0 1 2 0 1 2 3 4]</Index></index>
    <content index='0'>
        <ListOffsetArray len='3'>
            <parameter name='__array__'>'string'</parameter>
            <offsets><Index dtype='int64' len='4'>
                [ 0  3  6 11]
            </Index></offsets>
            <content><NumpyArray dtype='uint8' len='11'>
                <parameter name='__array__'>'char'</parameter>
                [111 110 101 116 119 111 116 104 114 101 101]
            </NumpyArray></content>
        </ListOffsetArray>
    </content>
    <content index='1'>
        <IndexedArray len='5'>
            <index><Index dtype='int64' len='5'>[4 3 2 1 0]</Index></index>
            <content><NumpyArray dtype='float64' len='5'>[1.1 2.2 3.3 4.4 5.5]</NumpyArray></content>
        </IndexedArray>
    </content>
</UnionArray>
```

But because of

  * https://github.com/scikit-hep/awkward/blob/0.1.15/docs/theory/arrays-are-functions.pdf
  * page 22‒23 https://indico.cern.ch/event/833895/contributions/3577882/
  * page 25 https://www.thestrangeloop.com/2019/jagged-ragged-awkward-arrays.html, starting at 18:24 in https://youtu.be/2NxWpU7NArk?t=1104

the application of the `index[tags == tag]` array in UnionArray is associative with the application of the `index` in the IndexedArray. The latter can be merged into the former with no loss of information, and then we're using less memory and performing less indirection in computations. (Is it enough to matter for any applications? Unclear, because we like to avoid UnionArrays in general, but it is a net win.)

This issue would be resolved by doing a final pass at the end of `UnionArray.simplified` to check for any IndexedArrays in `contents`, not IndexedOptionArrays, that are not categorical. (IndexedOptionArrays can't be merged into `UnionArray.index` because they have `-1` for missing values, which `UnionArray.index` does not, and we don't want to apply this to categorical IndexedArrays because then we'd have to put the `__array__: "categorical"` parameter on the UnionArray, which downstream code is not expecting, and anyway it would lose information because if some of the `UnionArray.contents` were categorical and others not, we'd lose knowledge of which is which.)

If the `outindex` (`index` of the UnionArray that will be returned) is a new array, created inside the `UnionArray.simplified` function, it can be modified in place, something like

```python
for tag, content in enumerate(outcontents):
    if isinstance(content, IndexedArray) and content.parameter("__array__") != "categorical":
        # only want to affect the part of the UnionArray.index for this tag
        selection = (outtags == tag)
        # function-composition of this part of the UnionArray.index with the IndexedArray.index
        outindex[selection] = content.index[outindex[selection]]
        # now we don't have an IndexedArray anymore
        outcontents[tag] = content.content
```

If not, then `outindex` can be copied before the first modification.

Then `UnionArray.__init__` can forbid non-categorical IndexedArrays as `contents`. Even fewer combinations to worry about: yay!

We should consider the case of IndexedArray-of-RecordArray, which is created by `_carry` with `allow_lazy=True`. This would eliminate the IndexedArray that lazily carries the RecordArray, but the RecordArray would still be lazily carried (that is, the `_carry` would still not be propagated to all of the RecordArray's `contents`) because that laziness is in the `UnionArray.index` now. When/if the RecordArray ever gets projected out of this UnionArray, if that happens with `allow_lazy=True`, then it's _still_ lazy (not any worse for having made this optimization). But if it gets projected out with `allow_lazy=False`, then that specific case could have a performance degradation due to this change. That seems hyperspecialized, though.


Discussion:
@jpivarski I was looking at tackling and noticed the following:
- `UnionArray.simplified` accepts a `merge` parameters that can be set to `False`
- `UnionArray._validity_error` ensures that no two contents are mergeable

I believe we should therefore deprecate `UnionArray.simplified(..., merge=)`, as `merge=True` is allowed, but `merge=False` will only work for non-mergeable contents.

Resolving patch:
diff --git a/src/awkward/contents/content.py b/src/awkward/contents/content.py
--- a/src/awkward/contents/content.py
+++ b/src/awkward/contents/content.py
@@ -1276,7 +1276,7 @@ def is_equal_to(
         return (
             self.__class__ is other.__class__
             and len(self) == len(other)
-            and type_parameters_equal(self.parameters, other.parameters)
+            and type_parameters_equal(self._parameters, other._parameters)
             and self._is_equal_to(other, index_dtype, numpyarray)
         )
 
diff --git a/src/awkward/contents/indexedoptionarray.py b/src/awkward/contents/indexedoptionarray.py
--- a/src/awkward/contents/indexedoptionarray.py
+++ b/src/awkward/contents/indexedoptionarray.py
@@ -835,7 +835,6 @@ def _fill_none(self, value: Content) -> Content:
             index,
             contents,
             parameters=self._parameters,
-            merge=True,
             mergebool=True,
         )
 
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
@@ -8,7 +8,7 @@
 
 import awkward as ak
 from awkward._backends.backend import Backend
-from awkward._errors import AxisError
+from awkward._errors import AxisError, deprecate
 from awkward._layout import maybe_posaxis
 from awkward._nplikes.jax import Jax
 from awkward._nplikes.numpy import Numpy
@@ -135,6 +135,13 @@ def __init__(self, tags, index, contents, *, parameters=None):
                 )
             elif content.is_option:
                 n_seen_options += 1
+            elif content.is_indexed and content.parameter("__array__") != "categorical":
+                raise TypeError(
+                    (
+                        "{0} cannot contain non-categorical indexed-types in its 'contents' ({1}); "
+                        "try {0}.simplified instead"
+                    ).format(type(self).__name__, type(content).__name__)
+                )
 
         if n_seen_options not in {0, len(contents)}:
             raise TypeError(
@@ -221,9 +228,17 @@ def simplified(
         contents,
         *,
         parameters=None,
-        merge=True,
+        merge=UNSET,
         mergebool=False,
     ):
+        if merge is UNSET:
+            merge = True
+        else:
+            deprecate(
+                "The `merge` argument to UnionArray.simplified(...) is deprecated; "
+                "it is no longer possible to construct UnionArrays with mergeable contents.",
+                version="2.4.0",
+            )
         self_index = index
         self_tags = tags
         self_contents = contents
@@ -396,6 +411,27 @@ def simplified(
                 "FIXME: handle UnionArray with more than 127 contents"
             )
 
+        # If any contents are non-categorical index types, we can merge them into the union
+        # This is safe, because any remaining index types at this point in the routine are not considered
+        # mergeable with the other contents. This means none of the other contents are option or index types,
+        # nor are any contents mergeable with these index types' contents. We already know that the
+        # other contents cannot be unions. Thus, it's safe to simply erase the outer indexed type.
+        for tag, content in enumerate(contents):
+            if (
+                isinstance(content, ak.contents.IndexedArray)
+                and content.parameter("__array__") != "categorical"
+            ):
+                # only want to affect the part of the UnionArray.index for this tag
+                selection = tags.data == tag
+                # function-composition of this part of the UnionArray.index with the IndexedArray.index
+                index.data[selection] = content.index.data[index.data[selection]]
+                # now we don't have an IndexedArray anymore, but we want to preserve its parameters
+                contents[tag] = content.content.copy(
+                    parameters=parameters_union(
+                        content.content.parameters, content.parameters
+                    )
+                )
+
         # If any contents are options, ensure all contents are options!
         if any(c.is_option for c in contents):
             contents = [
@@ -1057,18 +1093,15 @@ def _mergemany(self, others: Sequence[Content]) -> Content:
         if len(others) == 0:
             return self
 
+        index_nplike = self._backend.index_nplike
         head, tail = self._merging_strategy(others)
 
         total_length = 0
         for array in head:
             total_length += array.length
 
-        nexttags = ak.index.Index8.empty(
-            total_length, nplike=self._backend.index_nplike
-        )
-        nextindex = ak.index.Index64.empty(
-            total_length, nplike=self._backend.index_nplike
-        )
+        nexttags = ak.index.Index8.empty(total_length, nplike=index_nplike)
+        nextindex = ak.index.Index64.empty(total_length, nplike=index_nplike)
 
         nextcontents = []
         length_so_far = 0
@@ -1084,8 +1117,8 @@ def _mergemany(self, others: Sequence[Content]) -> Content:
                 union_index = array.index
                 union_contents = array.contents
                 assert (
-                    nexttags.nplike is self._backend.index_nplike
-                    and union_tags.nplike is self._backend.index_nplike
+                    nexttags.nplike is index_nplike
+                    and union_tags.nplike is index_nplike
                 )
                 self._backend.maybe_kernel_error(
                     self._backend[
@@ -1101,8 +1134,8 @@ def _mergemany(self, others: Sequence[Content]) -> Content:
                     )
                 )
                 assert (
-                    nextindex.nplike is self._backend.index_nplike
-                    and union_index.nplike is self._backend.index_nplike
+                    nextindex.nplike is index_nplike
+                    and union_index.nplike is index_nplike
                 )
                 self._backend.maybe_kernel_error(
                     self._backend[
@@ -1120,8 +1153,41 @@ def _mergemany(self, others: Sequence[Content]) -> Content:
 
                 nextcontents.extend(union_contents)
 
+            # This pathway is covered by `.simplified`, but we can avoid an extra array operation
+            # by performing this now
+            elif (
+                isinstance(array, ak.contents.IndexedArray)
+                and array.parameter("__array__") != "categorical"
+            ):
+                assert nexttags.nplike is index_nplike
+                self._backend.maybe_kernel_error(
+                    self._backend[
+                        "awkward_UnionArray_filltags_const",
+                        nexttags.dtype.type,
+                    ](
+                        nexttags.data,
+                        length_so_far,
+                        array.length,
+                        len(nextcontents),
+                    )
+                )
+                nextindex.data[
+                    index_nplike.shape_item_as_index(
+                        length_so_far
+                    ) : index_nplike.shape_item_as_index(length_so_far + array.length)
+                ] = array.index.data
+
+                length_so_far += array.length
+
+                nextcontents.append(
+                    array.content.copy(
+                        parameters=parameters_union(
+                            array.content._parameters, array._parameters
+                        )
+                    )
+                )
             else:
-                assert nexttags.nplike is self._backend.index_nplike
+                assert nexttags.nplike is index_nplike
                 self._backend.maybe_kernel_error(
                     self._backend[
                         "awkward_UnionArray_filltags_const",
@@ -1134,7 +1200,7 @@ def _mergemany(self, others: Sequence[Content]) -> Content:
                     )
                 )
 
-                assert nextindex.nplike is self._backend.index_nplike
+                assert nextindex.nplike is index_nplike
                 self._backend.maybe_kernel_error(
                     self._backend[
                         "awkward_UnionArray_fillindex_count", nextindex.dtype.type
@@ -1221,7 +1287,6 @@ def _is_unique(self, negaxis, starts, parents, outlength):
             self._index,
             self._contents,
             parameters=self._parameters,
-            merge=True,
             mergebool=True,
         )
         if isinstance(simplified, ak.contents.UnionArray):
@@ -1235,7 +1300,6 @@ def _unique(self, negaxis, starts, parents, outlength):
             self._index,
             self._contents,
             parameters=self._parameters,
-            merge=True,
             mergebool=True,
         )
         if isinstance(simplified, ak.contents.UnionArray):


Test patch:
diff --git a/tests/test_2192_union_absorb_indexed.py b/tests/test_2192_union_absorb_indexed.py
new file mode 100644
--- /dev/null
+++ b/tests/test_2192_union_absorb_indexed.py
@@ -0,0 +1,176 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+import numpy as np
+import pytest  # noqa: F401
+
+import awkward as ak
+
+
+def test_simplified_indexed():
+    numbers = ak.contents.IndexedArray(
+        ak.index.Index64([0, 2, 3]),
+        ak.contents.NumpyArray(np.arange(10, dtype=np.int64)),
+    )
+    records = ak.contents.RecordArray(
+        [ak.contents.NumpyArray(np.array([4.0, 3.0, 1.0], dtype=np.int64))], ["x"]
+    )
+    result = ak.concatenate((numbers, records), highlevel=False)
+    assert result.is_equal_to(
+        ak.contents.UnionArray(
+            ak.index.Index8([0, 0, 0, 1, 1, 1]),
+            ak.index.Index64([0, 2, 3, 0, 1, 2]),
+            [
+                ak.contents.NumpyArray(np.arange(10, dtype=np.int64)),
+                ak.contents.RecordArray(
+                    [ak.contents.NumpyArray(np.array([4.0, 3.0, 1.0], dtype=np.int64))],
+                    ["x"],
+                ),
+            ],
+        )
+    )
+
+
+def test_simplified_indexed_categorical():
+    numbers = ak.contents.IndexedArray(
+        ak.index.Index64([0, 2, 3]),
+        ak.contents.NumpyArray(np.arange(10, dtype=np.int64)),
+        parameters={"__array__": "categorical"},
+    )
+    records = ak.contents.RecordArray(
+        [ak.contents.NumpyArray(np.array([4.0, 3.0, 1.0], dtype=np.int64))], ["x"]
+    )
+    result = ak.concatenate((numbers, records), highlevel=False)
+    assert result.is_equal_to(
+        ak.contents.UnionArray(
+            ak.index.Index8([0, 0, 0, 1, 1, 1]),
+            ak.index.Index64([0, 1, 2, 0, 1, 2]),
+            [
+                ak.contents.IndexedArray(
+                    ak.index.Index64([0, 2, 3]),
+                    ak.contents.NumpyArray(np.arange(10, dtype=np.int64)),
+                    parameters={"__array__": "categorical"},
+                ),
+                ak.contents.RecordArray(
+                    [ak.contents.NumpyArray(np.array([4.0, 3.0, 1.0], dtype=np.int64))],
+                    ["x"],
+                ),
+            ],
+        )
+    )
+
+
+def test_merge_indexed():
+    records = ak.contents.IndexedArray(
+        ak.index.Index64([0, 2, 3]),
+        ak.contents.RecordArray(
+            [
+                ak.contents.NumpyArray(
+                    np.array([4.0, 3.0, 1.0, 9.0, 8.0, 7.0], dtype=np.int64)
+                )
+            ],
+            ["x"],
+            parameters={"inner": "bar", "drop": "this"},
+        ),
+        parameters={"outer": "foo", "ignore": "me"},
+    )
+    union = ak.contents.UnionArray(
+        ak.index.Index8([0, 0, 0, 1, 1, 1]),
+        ak.index.Index64([0, 1, 2, 0, 1, 2]),
+        [
+            ak.contents.NumpyArray(np.arange(10, dtype=np.int64)),
+            ak.contents.RecordArray(
+                [ak.contents.NumpyArray(np.array([4.0, 3.0, 1.0], dtype=np.int64))],
+                ["x"],
+                parameters={"outer": "foo", "inner": "bar"},
+            ),
+        ],
+    )
+    result = ak.concatenate((union, records), highlevel=False)
+
+    assert result.is_equal_to(
+        ak.contents.UnionArray(
+            ak.index.Index8([0, 0, 0, 1, 1, 1, 1, 1, 1]),
+            ak.index.Index64([0, 1, 2, 0, 1, 2, 3, 5, 6]),
+            [
+                ak.contents.NumpyArray(np.arange(10, dtype=np.int64)),
+                ak.contents.RecordArray(
+                    [
+                        ak.contents.NumpyArray(
+                            np.array(
+                                [4.0, 3.0, 1.0, 4.0, 3.0, 1.0, 9.0, 8.0, 7.0],
+                                dtype=np.int64,
+                            )
+                        )
+                    ],
+                    ["x"],
+                ),
+            ],
+        )
+    )
+    # This test might be a bit strict; any code that views `layout.parameters` will change this result to `{}`
+    assert result.contents[0]._parameters is None
+    assert result.contents[1]._parameters == {"outer": "foo", "inner": "bar"}
+
+
+def test_merge_indexed_categorical():
+    records = ak.contents.IndexedArray(
+        ak.index.Index64([0, 2, 3]),
+        ak.contents.RecordArray(
+            [
+                ak.contents.NumpyArray(
+                    np.array([4.0, 3.0, 1.0, 9.0, 8.0, 7.0], dtype=np.int64)
+                )
+            ],
+            ["x"],
+            parameters={"inner": "bar", "drop": "this"},
+        ),
+        parameters={"outer": "foo", "ignore": "me", "__array__": "categorical"},
+    )
+    union = ak.contents.UnionArray(
+        ak.index.Index8([0, 0, 0, 1, 1, 1]),
+        ak.index.Index64([0, 1, 2, 0, 1, 2]),
+        [
+            ak.contents.NumpyArray(np.arange(10, dtype=np.int64)),
+            ak.contents.IndexedArray(
+                ak.index.Index64([0, 1, 2]),
+                ak.contents.RecordArray(
+                    [ak.contents.NumpyArray(np.array([4.0, 3.0, 1.0], dtype=np.int64))],
+                    ["x"],
+                    parameters={"inner": "bar"},
+                ),
+                parameters={"outer": "foo", "__array__": "categorical"},
+            ),
+        ],
+    )
+    result = ak.concatenate((union, records), highlevel=False)
+    assert result.is_equal_to(
+        ak.contents.UnionArray(
+            ak.index.Index8([0, 0, 0, 1, 1, 1, 1, 1, 1]),
+            ak.index.Index64([0, 1, 2, 0, 1, 2, 3, 4, 5]),
+            [
+                ak.contents.NumpyArray(np.arange(10, dtype=np.int64)),
+                ak.contents.IndexedArray(
+                    ak.index.Index64([0, 1, 2, 3, 5, 6]),
+                    ak.contents.RecordArray(
+                        [
+                            ak.contents.NumpyArray(
+                                np.array(
+                                    [4.0, 3.0, 1.0, 4.0, 3.0, 1.0, 9.0, 8.0, 7.0],
+                                    dtype=np.int64,
+                                )
+                            )
+                        ],
+                        ["x"],
+                    ),
+                    parameters={"__array__": "categorical"},
+                ),
+            ],
+        )
+    )
+    # This test might be a bit strict; any code that views `layout.parameters` will change this result to `{}`
+    assert result.contents[0]._parameters is None
+    assert result.contents[1]._parameters == {
+        "__array__": "categorical",
+        "outer": "foo",
+    }
+    assert result.contents[1].content._parameters == {"inner": "bar"}