[Python][RDF] Extend FunctionJitter with newest supported numba jitting types

Author: siliataider

bindings/pyroot/pythonizations/python/ROOT/_numbadeclare.py

@@ -46,12 +46,12 @@ def _NumbaDeclareDecorator(input_types, return_type=None, name=None):
             "match_pattern": r"(?:ROOT::)?(?:VecOps::)?RVec\w+|(?:ROOT::)?(?:VecOps::)?RVec<[\w\s]+>",
             "cpp_name": ["ROOT::RVec", "ROOT::VecOps::RVec"],
         },
-        "std::vector": {
-            "match_pattern": r"std::vector<[\w\s]+>",
+        "vector": {
+            "match_pattern": r"(?:std::)?vector<[\w\s]+>",
             "cpp_name": ["std::vector"],
         },
-        "std::array": {
-            "match_pattern": r"std::array<[\w\s,<>]+>",
+        "array": {
+            "match_pattern": r"(?:std::)?array<[\w\s,<>]+>",
             "cpp_name": ["std::array"],
         },
     }
@@ -233,7 +233,6 @@ def inner(func, input_types=input_types, return_type=return_type, name=name):
         """
         Inner decorator without arguments, see outer decorator for documentation
         """
-
         # Jit the given Python callable with numba
         try:
             nb_return_type, nb_input_types = get_numba_signature(input_types, return_type)
@@ -255,6 +254,13 @@ def inner(func, input_types=input_types, return_type=return_type, name=name):
                         "See https://cppyy.readthedocs.io/en/latest/numba.html#numba-support"
                     )
                 nbjit = nb.jit(nopython=True, inline="always")(func)
+                # In this case, the user has to explictly provide the return type, cannot be inferred
+                if return_type is None:
+                    raise RuntimeError(
+                        "Failed to infer the return type for the provided function. "
+                        "Please specify the signature explicitly in the decorator, e.g.: "
+                        "@ROOT.NumbaDeclare(['double'], 'double')"
+                    )
             except:  # noqa E722
                 raise Exception("Failed to jit Python callable {} with numba.jit".format(func))
         func.numba_func = nbjit

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rdf_pyz.py

@@ -81,17 +81,14 @@ def find_type(self, x):
                 t = self.rdf.GetColumnType(x)
                 if t in TREE_TO_NUMBA:  # The column is a fundamental type from tree
                     return TREE_TO_NUMBA[t]
-                elif "<" in t:  # The column type is a RVec<type>
-                    if ">>" in t:  # It is a RVec<RVec<T>>
-                        raise TypeError(
-                            f"Only columns with 'RVec<T>' where T is is a fundamental type are supported, not '{t}'."
-                        )
-                    g = re.match("(.*)<(.*)>", t).groups(0)
-                    if g[1] in TREE_TO_NUMBA:
-                        return "RVec<" + TREE_TO_NUMBA[g[1]] + ">"
-                    # There are data type that leak into here. Not sure from where. But need to implement something here such that this condition is never met.
-                    return "RVec<" + str(g[1]) + ">"
 
+                match = re.match(r"([\w:]+)<(.+)>", t)
+                if match:
+                    container_type, inner_type = match.groups()
+                    container_type = container_type.strip()
+                    inner_type = inner_type.strip()
+                    inner_mapped = TREE_TO_NUMBA.get(inner_type, inner_type)
+                    return f"{container_type}<{inner_mapped}>"
                 else:
                     return t
             else: