Merge branch 'main' into jpivarski/protect-ak-to_parquet-against-memo…

…ry-explosion

awkward-cpp/pyproject.toml

@@ -7,7 +7,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "awkward_cpp"
-version = "16"
+version = "17"
 dependencies = [
     "numpy>=1.17.0"
 ]

pyproject.toml

@@ -40,7 +40,7 @@ classifiers = [
     "Topic :: Utilities",
 ]
 dependencies = [
-    "awkward_cpp==16",
+    "awkward_cpp==17",
     "importlib_resources;python_version < \"3.9\"",
     "numpy>=1.17.0",
     "packaging",

src/awkward/_broadcasting.py

@@ -934,7 +934,13 @@ def broadcast_any_union():
         )
 
     def broadcast_any_indexed():
-        nextinputs = [x.project() if isinstance(x, IndexedArray) else x for x in inputs]
+        # The `apply` function may exit at the level of a `RecordArray`. We can avoid projection
+        # of the record array in such cases, in favour of a deferred carry. This can be done by
+        # "pushing" the `IndexedArray` _into_ the record (i.e., wrapping each `content`).
+        nextinputs = [
+            x._push_inside_record_or_project() if isinstance(x, IndexedArray) else x
+            for x in inputs
+        ]
         return apply_step(
             backend,
             nextinputs,

src/awkward/_errors.py

@@ -1,6 +1,7 @@
 # BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
 from __future__ import annotations
 
+import builtins
 import sys
 import threading
 import warnings
@@ -31,6 +32,11 @@ def __call__(self):
         return self.func(*self.args, **self.kwargs)
 
 
+class KeyError(builtins.KeyError):
+    def __str__(self):
+        return super(Exception, self).__str__()
+
+
 class ErrorContext:
     # Any other threads should get a completely independent _slate.
     _slate = threading.local()
@@ -87,6 +93,9 @@ def decorate_exception(self, cls: type[E], exception: E) -> E:
                     + "\n\nSee if this has been reported at https://github.com/scikit-hep/awkward/issues"
                 )
                 new_exception.__cause__ = exception
+            elif issubclass(cls, builtins.KeyError):
+                new_exception = KeyError(self.format_exception(exception))
+                new_exception.__cause__ = exception
             else:
                 new_exception = cls(self.format_exception(exception))
                 new_exception.__cause__ = exception

src/awkward/contents/indexedarray.py

@@ -1136,3 +1136,18 @@ def _is_equal_to(self, other, index_dtype, numpyarray):
         return self.index.is_equal_to(
             other.index, index_dtype, numpyarray
         ) and self.content.is_equal_to(other.content, index_dtype, numpyarray)
+
+    def _push_inside_record_or_project(self) -> Self | ak.contents.RecordArray:
+        if self.content.is_record:
+            return ak.contents.RecordArray(
+                contents=[
+                    ak.contents.IndexedArray.simplified(self._index, c)
+                    for c in self.content.contents
+                ],
+                fields=self.content._fields,
+                length=self.length,
+                backend=self._backend,
+                parameters=parameters_union(self.content._parameters, self._parameters),
+            )
+        else:
+            return self.project()