refactor: split ak._do into meta and content

src/awkward/__init__.py

@@ -16,7 +16,8 @@
 import awkward.forms
 
 # internal
-import awkward._do
+import awkward._do.content
+import awkward._do.meta
 import awkward._slicing
 import awkward._broadcasting
 import awkward._reducers

src/awkward/_connect/jax/trees.py

@@ -9,6 +9,7 @@
 from awkward._backends.backend import Backend
 from awkward._backends.jax import JaxBackend
 from awkward._behavior import behavior_of
+from awkward._do.content import recursively_apply
 from awkward._layout import wrap_layout
 from awkward._nplikes.jax import Jax
 from awkward._nplikes.numpy import Numpy
@@ -30,9 +31,7 @@ def action(node, **kwargs):
         if isinstance(node, ak.contents.NumpyArray):
             data_ptrs.append(node.data)
 
-    ak._do.recursively_apply(
-        layout, action=action, return_array=False, numpy_to_regular=False
-    )
+    recursively_apply(layout, action=action, return_array=False, numpy_to_regular=False)
 
     return data_ptrs
 
@@ -54,7 +53,7 @@ def action(node, **kwargs):
                     buffer, parameters=node.parameters, backend=backend
                 )
 
-    return ak._do.recursively_apply(layout, action=action, numpy_to_regular=False)
+    return recursively_apply(layout, action=action, numpy_to_regular=False)
 
 
 T = TypeVar(

src/awkward/_connect/numpy.py

@@ -5,7 +5,7 @@
 import collections
 import functools
 import inspect
-from collections.abc import Iterable
+from collections.abc import Iterable, Iterator, Mapping
 from itertools import chain
 
 import numpy
@@ -21,11 +21,12 @@
     find_ufunc_generic,
 )
 from awkward._categorical import as_hashable
+from awkward._do.content import reduce as do_reduce
 from awkward._layout import wrap_layout
 from awkward._nplikes import to_nplike
 from awkward._parameters import parameters_intersect
 from awkward._regularize import is_non_string_like_iterable
-from awkward._typing import Any, Iterator, Mapping
+from awkward._typing import Any
 from awkward._util import Sentinel
 from awkward.contents.numpyarray import NumpyArray
 
@@ -316,11 +317,9 @@ def _array_ufunc_string_likes(
     nplike = left.backend.nplike
 
     # first condition: string lengths must be the same
-    left_counts_layout = ak._do.reduce(left, ak._reducers.Count(), axis=-1, mask=False)
+    left_counts_layout = do_reduce(left, ak._reducers.Count(), axis=-1, mask=False)
     assert left_counts_layout.is_numpy
-    right_counts_layout = ak._do.reduce(
-        right, ak._reducers.Count(), axis=-1, mask=False
-    )
+    right_counts_layout = do_reduce(right, ak._reducers.Count(), axis=-1, mask=False)
     assert right_counts_layout.is_numpy
 
     counts1 = nplike.asarray(left_counts_layout.data)

src/awkward/_do.py → src/awkward/_do/content.py

@@ -9,14 +9,16 @@
 import awkward as ak
 from awkward._backends.backend import Backend
 from awkward._nplikes.numpy_like import NumpyMetadata
-from awkward._typing import Any, AxisMaybeNone, Literal
-from awkward.contents.content import ActionType, Content
+from awkward._typing import TYPE_CHECKING, Any, AxisMaybeNone, Literal
 from awkward.errors import AxisError
 from awkward.forms import form
-from awkward.record import Record
 
 np = NumpyMetadata.instance()
 
+if TYPE_CHECKING:
+    from awkward.contents.content import ActionType, Content
+    from awkward.record import Record
+
 
 def recursively_apply(
     layout: Content | Record,
@@ -32,6 +34,9 @@ def recursively_apply(
     function_name: str | None = None,
     regular_to_jagged=False,
 ) -> Content | Record | None:
+    from awkward.contents.content import Content
+    from awkward.record import Record
+
     if isinstance(layout, Content):
         return layout._recursively_apply(
             action,
@@ -201,6 +206,8 @@ def remove_structure(
     allow_records: bool = False,
     list_to_regular: bool = False,
 ):
+    from awkward.record import Record
+
     if isinstance(layout, Record):
         return remove_structure(
             layout._array[layout._at : layout._at + 1],
@@ -246,10 +253,6 @@ def num(layout, axis):
     return layout._num(axis, 0)
 
 
-def mergeable(one: Content, two: Content, mergebool: bool = True) -> bool:
-    return one._mergeable_next(two, mergebool=mergebool)
-
-
 def mergemany(contents: list[Content]) -> Content:
     assert len(contents) != 0
     return contents[0]._mergemany(contents[1:])

src/awkward/_do/meta.py

@@ -0,0 +1,15 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
+
+from __future__ import annotations
+
+from awkward._nplikes.numpy_like import NumpyMetadata
+from awkward._typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from awkward._meta.meta import Meta
+
+np = NumpyMetadata.instance()
+
+
+def mergeable(one: Meta, two: Meta, mergebool: bool = True) -> bool:
+    return one._mergeable_next(two, mergebool=mergebool)

src/awkward/_meta/bitmaskedmeta.py

@@ -2,14 +2,13 @@
 
 from __future__ import annotations
 
-from awkward._meta.meta import Meta
-from awkward._typing import Generic, JSONSerializable, TypeVar
+from awkward._meta.meta import Meta, is_indexed, is_option
+from awkward._parameters import type_parameters_equal
+from awkward._typing import JSONSerializable
 
-T = TypeVar("T", bound=Meta)
 
-
-class BitMaskedMeta(Meta, Generic[T]):
-    _content: T
+class BitMaskedMeta(Meta):
+    _content: Meta
     is_option = True
 
     @property
@@ -53,5 +52,17 @@ def dimension_optiontype(self) -> bool:
         return True
 
     @property
-    def content(self) -> T:
+    def content(self) -> Meta:
         return self._content
+
+    def _mergeable_next(self, other: Meta, mergebool: bool) -> bool:
+        # Is the other content is an identity, or a union?
+        if other.is_identity_like or other.is_union:
+            return True
+        # We can only combine option/indexed types whose array-record parameters agree
+        elif is_option(other) or is_indexed(other):
+            return self._content._mergeable_next(
+                other.content, mergebool
+            ) and type_parameters_equal(self._parameters, other._parameters)
+        else:
+            return self._content._mergeable_next(other, mergebool)

src/awkward/_meta/bytemaskedmeta.py

@@ -2,14 +2,13 @@
 
 from __future__ import annotations
 
-from awkward._meta.meta import Meta
-from awkward._typing import Generic, JSONSerializable, TypeVar
+from awkward._meta.meta import Meta, is_indexed, is_option
+from awkward._parameters import type_parameters_equal
+from awkward._typing import JSONSerializable
 
-T = TypeVar("T", bound=Meta)
 
-
-class ByteMaskedMeta(Meta, Generic[T]):
-    _content: T
+class ByteMaskedMeta(Meta):
+    _content: Meta
     is_option = True
 
     @property
@@ -53,5 +52,17 @@ def dimension_optiontype(self) -> bool:
         return True
 
     @property
-    def content(self) -> T:
+    def content(self) -> Meta:
         return self._content
+
+    def _mergeable_next(self, other: Meta, mergebool: bool) -> bool:
+        # Is the other content is an identity, or a union?
+        if other.is_identity_like or other.is_union:
+            return True
+        # We can only combine option/indexed types whose array-record parameters agree
+        elif is_option(other) or is_indexed(other):
+            return self._content._mergeable_next(
+                other.content, mergebool
+            ) and type_parameters_equal(self._parameters, other._parameters)
+        else:
+            return self._content._mergeable_next(other, mergebool)

src/awkward/_meta/emptymeta.py

@@ -44,3 +44,6 @@ def is_tuple(self) -> bool:
     @property
     def dimension_optiontype(self) -> bool:
         return False
+
+    def _mergeable_next(self, other: Meta, mergebool: bool) -> bool:
+        return True

src/awkward/_meta/indexedmeta.py

@@ -2,16 +2,15 @@
 
 from __future__ import annotations
 
-from awkward._meta.meta import Meta
-from awkward._typing import Generic, JSONSerializable, TypeVar
+from awkward._meta.meta import Meta, is_indexed, is_option
+from awkward._parameters import type_parameters_equal
+from awkward._typing import JSONSerializable
 
-T = TypeVar("T", bound=Meta)
 
-
-class IndexedMeta(Meta, Generic[T]):
+class IndexedMeta(Meta):
     is_indexed = True
 
-    _content: T
+    _content: Meta
 
     def purelist_parameters(self, *keys: str) -> JSONSerializable:
         if self._parameters is not None:
@@ -54,5 +53,17 @@ def dimension_optiontype(self) -> bool:
         return False
 
     @property
-    def content(self) -> T:
+    def content(self) -> Meta:
         return self._content
+
+    def _mergeable_next(self, other: Meta, mergebool: bool) -> bool:
+        # Is the other content is an identity, or a union?
+        if other.is_identity_like or other.is_union:
+            return True
+        # We can only combine option/indexed types whose array-record parameters agree
+        elif is_option(other) or is_indexed(other):
+            return self._content._mergeable_next(
+                other.content, mergebool
+            ) and type_parameters_equal(self._parameters, other._parameters)
+        else:
+            return self._content._mergeable_next(other, mergebool)

src/awkward/_meta/indexedoptionmeta.py

@@ -2,17 +2,16 @@
 
 from __future__ import annotations
 
-from awkward._meta.meta import Meta
-from awkward._typing import Generic, JSONSerializable, TypeVar
+from awkward._meta.meta import Meta, is_indexed, is_option
+from awkward._parameters import type_parameters_equal
+from awkward._typing import JSONSerializable
 
-T = TypeVar("T", bound=Meta)
 
-
-class IndexedOptionMeta(Meta, Generic[T]):
+class IndexedOptionMeta(Meta):
     is_indexed = True
     is_option = True
 
-    _content: T
+    _content: Meta
 
     def purelist_parameters(self, *keys: str) -> JSONSerializable:
         if self._parameters is not None:
@@ -55,5 +54,17 @@ def dimension_optiontype(self) -> bool:
         return True
 
     @property
-    def content(self) -> T:
+    def content(self) -> Meta:
         return self._content
+
+    def _mergeable_next(self, other: Meta, mergebool: bool) -> bool:
+        # Is the other content is an identity, or a union?
+        if other.is_identity_like or other.is_union:
+            return True
+        # We can only combine option/indexed types whose array-record parameters agree
+        elif is_option(other) or is_indexed(other):
+            return self._content._mergeable_next(
+                other.content, mergebool
+            ) and type_parameters_equal(self._parameters, other._parameters)
+        else:
+            return self._content._mergeable_next(other, mergebool)

src/awkward/_meta/listmeta.py

@@ -2,16 +2,15 @@
 
 from __future__ import annotations
 
-from awkward._meta.meta import Meta
-from awkward._typing import Generic, JSONSerializable, TypeVar
+from awkward._meta.meta import Meta, is_indexed, is_list, is_numpy, is_option
+from awkward._parameters import type_parameters_equal
+from awkward._typing import JSONSerializable
 
-T = TypeVar("T", bound=Meta)
 
-
-class ListMeta(Meta, Generic[T]):
+class ListMeta(Meta):
     is_list = True
 
-    _content: T
+    _content: Meta
 
     def purelist_parameters(self, *keys: str) -> JSONSerializable:
         if self._parameters is not None:
@@ -65,5 +64,23 @@ def dimension_optiontype(self) -> bool:
         return False
 
     @property
-    def content(self) -> T:
+    def content(self) -> Meta:
         return self._content
+
+    def _mergeable_next(self, other: Meta, mergebool: bool) -> bool:
+        # Is the other content is an identity, or a union?
+        if other.is_identity_like or other.is_union:
+            return True
+        # Check against option contents
+        elif is_option(other) or is_indexed(other):
+            assert hasattr(other, "content")
+            return self._mergeable_next(other.content, mergebool)
+        # Otherwise, do the parameters match? If not, we can't merge.
+        elif not type_parameters_equal(self._parameters, other._parameters):
+            return False
+        elif is_list(other):
+            return self._content._mergeable_next(other.content, mergebool)
+        elif is_numpy(other) and len(other.inner_shape) > 0:
+            return self._mergeable_next(other._to_regular_primitive(), mergebool)
+        else:
+            return False