RVecs as numpy arrays via numba in RDataFrame

numba_numpy_rvec.py

import numba as nb
from numba.types import CPointer, Array, void
import ROOT
import numpy as np
import ctypes

def c_friendly_signature(args, ret):
    """Take a normalized numba signature, return a C-friendly signature:
    array types are substituted with a pointer/size pair.
    """
    cargs = []
    for arg in args:
        if isinstance(arg, Array):
            # convert to T* and size_t input parameters
            cargs.extend((CPointer(arg.dtype), nb.i4))
        else:
            cargs.append(arg)
    if isinstance(ret, Array):
        # convert to T** and size_t* output parameters, return void
        cargs.extend((CPointer(CPointer(ret.dtype)), CPointer(nb.i4)))
        cret = void
    else:
        cret = ret
    return cargs, cret

# decorator: returns function that takes a python function
# and wraps it so that instead of arrays, CPointer and length are
# exposed, making it easy to jit it and call it from C++.
def cpp_callable_with_array(sig):
    def decorator(func):
        args, ret = nb.sigutils.normalize_signature(sig)
        cargs, cret = c_friendly_signature(args, ret)
        if cret == void:
            raise NotImplementedError("RVec as return type not supported yet")
        else:
            globals()["jitted_func"] = nb.jit(ret(*args))(func)

            # TODO all three parameter lists must be programmatically generated
            c_friendly_fun_str = """def c_friendly_func(vptr, size):
                arr = nb.carray(vptr, (size,))
                return jitted_func(arr)
            """
            exec(c_friendly_fun_str, globals(), locals())

            cf = nb.cfunc(cret(*cargs), inline="always")(locals()["c_friendly_func"])
            cf_addr = cf.address

            # TODO reinterpret_cast target type must be programmatically generated
            # TODO C++ function signature must be programmatically generated
            ROOT.gInterpreter.Declare("""namespace cpp_callable {{
                float sum(ROOT::RVec<float> &v) {{
                    return reinterpret_cast<float(*)(float*, int)>({})(v.data(), v.size());
                }}
            }}""".format(cf_addr))
        return func
    return decorator

def test_ret_float():
    @cpp_callable_with_array("f4(f4[:])")
    def sum(arr):
        return arr.sum()
    
    m = ROOT.RDataFrame(10).Define("v", "ROOT::RVec<float>{1,2,3}").Define("s", "cpp_callable::sum(v)").Mean("s");
    assert np.isclose(m.GetValue(), 6)

def main():
    print("test_ret_float...")
    test_ret_float()
    print("done")

if __name__ == "__main__":
    main()

When return type is an RVec, need something like

def original_python(arr):
  return arr[arr > 3]

def python_adapter(arrptr, arrsize, retptrptr, retsizeptr):
  arr = numba.carray(arrptr, arrsize)
  ret = original_python(arr)
  retptrptr[0] = ret.__array_interface__["data"][0]
  retsizeptr[0] = len(ret)
  ret.__array_interface__["data"][1] = False # set as not owning

RVec<float> cpp_wrapper(RVec<float> &v) {
  float** ret = nullptr;
  int *retsize = new int();
  python_adapter_cfunc(v.data(), v.size(), ret, retsize);
  RVec<float> ret_v(*ret, *ret + *retsize);
  delete retsize;
  delete *ret[];
  return ret_v;
}