From ac97e2a64fb78a636550a8992f6b2a7a563c1b57 Mon Sep 17 00:00:00 2001
From: Angus Hollands <goosey15@gmail.com>
Date: Mon, 4 Mar 2024 19:52:56 +0000
Subject: [PATCH] fix: performance of ufunc resolution for non-nominal
 signatures (#3030)

---
 src/awkward/_behavior.py      | 45 +++++++++++++++++++----------------
 src/awkward/_connect/numpy.py | 22 +++++++++++------
 2 files changed, 39 insertions(+), 28 deletions(-)

diff --git a/src/awkward/_behavior.py b/src/awkward/_behavior.py
index e47051b65b..f994aad370 100644
--- a/src/awkward/_behavior.py
+++ b/src/awkward/_behavior.py
@@ -100,29 +100,32 @@ def find_ufunc_generic(
 
 
 def find_ufunc(behavior: Mapping | None, signature: tuple) -> UfuncLike | None:
-    if all(s is not None for s in signature):
-        behavior = overlay_behavior(behavior)
+    if any(s is None for s in signature):
+        return None
+
+    behavior = overlay_behavior(behavior)
 
-        # Special case all strings or hashable types.
+    # Try and fast-path the lookup
+    try:
+        return behavior[signature]
+    except KeyError:
+        # We didn't find an exact overload, and we won't find any!
         if all(isinstance(x, str) for x in signature):
-            return behavior.get(signature)
-        else:
-            for key, custom in behavior.items():
-                if (
-                    isinstance(key, tuple)
-                    and len(key) == len(signature)
-                    and key[0] == signature[0]
-                    and all(
-                        k == s
-                        or (
-                            isinstance(k, type)
-                            and isinstance(s, type)
-                            and issubclass(s, k)
-                        )
-                        for k, s in zip(key[1:], signature[1:])
-                    )
-                ):
-                    return custom
+            return None
+
+    # Fall back on linear search (first-wins)
+    for key, custom in behavior.items():
+        if (
+            isinstance(key, tuple)
+            and len(key) == len(signature)
+            and key[0] == signature[0]
+            and all(
+                k == s
+                or (isinstance(k, type) and isinstance(s, type) and issubclass(s, k))
+                for k, s in zip(key[1:], signature[1:])
+            )
+        ):
+            return custom
     return None
 
 
diff --git a/src/awkward/_connect/numpy.py b/src/awkward/_connect/numpy.py
index 4529677dfb..f17ee98b36 100644
--- a/src/awkward/_connect/numpy.py
+++ b/src/awkward/_connect/numpy.py
@@ -207,15 +207,18 @@ def _array_ufunc_adjust_apply(
         )
 
 
-def _array_ufunc_signature(ufunc, inputs):
-    signature = [ufunc]
+def _array_ufunc_signature(ufunc, inputs) -> tuple[Any, ...] | None:
+    signature = []
+    has_seen_nominal_type = False
     for x in inputs:
         if isinstance(x, ak.contents.Content):
             record_name, list_name = x.parameter("__record__"), x.parameter("__list__")
             if record_name is not None:
                 signature.append(record_name)
+                has_seen_nominal_type = True
             elif list_name is not None:
                 signature.append(list_name)
+                has_seen_nominal_type = True
             elif isinstance(x, NumpyArray):
                 signature.append(x.dtype.type)
             else:
@@ -223,7 +226,10 @@ def _array_ufunc_signature(ufunc, inputs):
         else:
             signature.append(type(x))
 
-    return tuple(signature)
+    if has_seen_nominal_type:
+        return (ufunc, *signature)
+    else:
+        return None
 
 
 def _array_ufunc_categorical(
@@ -364,10 +370,12 @@ def action(inputs, **ignore):
         assert len(contents) >= 1
 
         signature = _array_ufunc_signature(ufunc, inputs)
-        # Do we have a custom ufunc (an override of the given ufunc)?
-        custom = find_ufunc(behavior, signature)
-        if custom is not None:
-            return _array_ufunc_adjust(custom, inputs, kwargs, behavior)
+        # Should we allow ufunc overloads for this signature?
+        if signature is not None:
+            # Do we have a custom ufunc (an override of the given ufunc)?
+            custom = find_ufunc(behavior, signature)
+            if custom is not None:
+                return _array_ufunc_adjust(custom, inputs, kwargs, behavior)
 
         # Do we have any categoricals?
         if any(