From 1cc8de8a59ddd418d818e605db2a2d0eb1f4249d Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Tue, 13 Oct 2026 19:08:24 +0200 Subject: [PATCH 1/4] [Python][RDF] Enable pure Python callables in RDataFrame A Python callable that is not decorated with `ROOT.Numba.Declare` is now passed to `Define` and `Filter` as it is and runs in the Python interpreter, instead of being implicitly JIT compiled with Numba. Its type annotations give the function pointer type used to explicitly instantiate `Define` and `Filter`. Assisted-by: Claude Opus 5.5 --- .../python/ROOT/_pythonization/_rdf_pyz.py | 79 ++++++++++++------- 1 file changed, 51 insertions(+), 28 deletions(-) diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rdf_pyz.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rdf_pyz.py index 1801474b712ab..e2c79b5db3959 100755 --- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rdf_pyz.py +++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rdf_pyz.py @@ -7,10 +7,9 @@ # For the licensing terms see $ROOTSYS/LICENSE. # # For the list of contributors see $ROOTSYS/README/CREDITS. # ################################################################################ +import inspect import re -import textwrap import typing -import warnings from .._numbadeclare import _NumbaDeclareDecorator @@ -257,30 +256,36 @@ def _resolve_overload_based_on_cols(func, rdf, cols): return func.__overload__(signature) -class _WarnOnce: - called = False - - @classmethod - def warn(cls): - if not cls.called: - msg = """ - RDataFrame is implicitly calling numba to JIT compile your Python functions. This behaviour will be in the - future only opt-in. You can silence this warning and achieve the same behaviour by explicitly calling the - ROOT.Numba.Declare decorator on your Python function and pass a string expression to RDataFrame that calls - the same function, for example: - - ``` - @ROOT.Numba.Declare([input_type_1, input_type_2], ret_type) - def foo(x, y): - return x + y - - df.Define("sum", "Numba::foo(x, y)") - ``` - - See more details at https://root.cern/doc/master/classROOT_1_1RDataFrame.html#python - """ - warnings.warn(textwrap.dedent(msg), FutureWarning) - cls.called = True +# The C++ wrapper of a Python callable only holds a weak reference to it, but the event loop calls it later, +# so we keep the callables alive here +_python_callables = set() + +_PYTHON_TO_CPP_TYPES = {bool: "bool", int: "int", float: "double", str: "std::string"} + + +def _cpp_type_name(func, annotation): + if isinstance(annotation, str): # already a C++ type name + return annotation + if annotation in _PYTHON_TO_CPP_TYPES: + return _PYTHON_TO_CPP_TYPES[annotation] + if hasattr(annotation, "__cpp_name__"): # ROOT class + return annotation.__cpp_name__ + raise TypeError(f"Cannot convert the type annotation {annotation!r} of '{func.__name__}' to a C++ type.") + + +def _function_pointer_type(func): + """ + Returns the C++ function pointer type of an annotated Python callable, e.g. "double (*)(int, double)". + """ + signature = inspect.signature(func) + params = [p.annotation for p in signature.parameters.values()] + if signature.return_annotation is inspect.Signature.empty or any(p is inspect.Parameter.empty for p in params): + raise TypeError( + f"The Python callable '{func.__name__}' needs type annotations for all its parameters and for its " + "return value, e.g. 'def f(x: float) -> float'." + ) + ret = _cpp_type_name(func, signature.return_annotation) + return f"{ret} (*)({', '.join(_cpp_type_name(func, p) for p in params)})" def _PyFilter(rdf, callable_or_str, *args, extra_args={}): @@ -359,7 +364,16 @@ def x_more_than_y(x): if isinstance(func, ROOT._cppyy.types.Instance): return rdf._OriginalFilter(func, *args) - _WarnOnce.warn() + if not hasattr(func, "__cpp_wrapper__"): # not decorated with ROOT.Numba.Declare + if extra_args: + raise TypeError( + "extra_args is only supported for callables decorated with ROOT.Numba.Declare. " + "A Python callable can reference the values directly, e.g. from the enclosing scope." + ) + _python_callables.add(func) + col_list = col_list or list(inspect.signature(func).parameters) # map the parameters to columns by name + return rdf._OriginalFilter[_function_pointer_type(func)](func, col_list, filter_name) + jitter = FunctionJitter(rdf) func.__annotations__["return"] = ( "bool" # return type for Filters is bool # Note: You can keep double and Filter still works. @@ -423,7 +437,16 @@ def x_scaled(x): if isinstance(func, ROOT._cppyy.types.Instance): return rdf._OriginalDefine(col_name, func, cols) - _WarnOnce.warn() + if not hasattr(func, "__cpp_wrapper__"): # not decorated with ROOT.Numba.Declare + if extra_args: + raise TypeError( + "extra_args is only supported for callables decorated with ROOT.Numba.Declare. " + "A Python callable can reference the values directly, e.g. from the enclosing scope." + ) + _python_callables.add(func) + cols = cols or list(inspect.signature(func).parameters) # map the parameters to columns by name + return rdf._OriginalDefine[_function_pointer_type(func)](col_name, func, cols) + jitter = FunctionJitter(rdf) func_call = jitter.jit_function(func, cols, extra_args) return rdf._OriginalDefine(col_name, "Numba::" + func_call) From fa38dcd6a5856dbe327420548d929fc5e2af973f Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Tue, 13 Oct 2026 19:08:24 +0200 Subject: [PATCH 2/4] [Python][RDF] Update the Define and Filter docstrings for Python callables Assisted-by: Claude Opus 5.5 --- .../python/ROOT/_pythonization/_rdf_pyz.py | 36 +++++++++++++------ 1 file changed, 26 insertions(+), 10 deletions(-) diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rdf_pyz.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rdf_pyz.py index e2c79b5db3959..bce12a5e73cd7 100755 --- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rdf_pyz.py +++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rdf_pyz.py @@ -294,30 +294,38 @@ def _PyFilter(rdf, callable_or_str, *args, extra_args={}): Arguments: callable_or_str: The condition according to which the RDF is to be filtered. It can be either a python callable or a c style string. + A python callable runs in the Python interpreter and needs type annotations for its parameters and + return value, unless it is decorated with ROOT.Numba.Declare to be JIT compiled with Numba. *args: Can be at most two of them. a. List of columns that the callable will receive as argument. If not provided then it tries maps the name of the parameter to a column name of the RDF. b. The name of the Filter as a string. - extra_args: non-columnar arguments to be passed to the callable. + extra_args: non-columnar arguments to be passed to a callable decorated with ROOT.Numba.Declare. Returns: RDataFrame: result of transforming the rdf argument by applying the requested filter operation. Examples: - 1. rdf.Filter(lambda x,y: x>y) + 1. def x_more_than_y(x: float, y: float) -> bool: + return x>y + rdf.Filter(x_more_than_y) Filter using a python callable. - 2. rdf.Filter(lambda x,y: x>y, ["y", "x"]) + 2. rdf.Filter(x_more_than_y, ["y", "x"]) Filter using a python callable with column list passed explicitly. When the column list is passed, the function maps the starting input parameters to those columns instead. Here the function param x is mapped to column y, and the second function param y is mapped to column x. This is equivalent to the filter "y>x". - 3. rdf.Filter(lambda col1, const1: col1>const1, ["x"], extra_args = {"const1":0.5}) + 3. @ROOT.Numba.Declare(["double", "double"], "bool") + def col1_more_than_const1(col1, const1): + return col1>const1 + rdf.Filter(col1_more_than_const1, ["x"], extra_args = {"const1":0.5}) Filter using a python callable with column list passed explicitly and extra_args. Here the function parameter col1 is mapped to col "x" of the RDF, and the const1 parameter is mapped to the float value 0.5 This is equivalent to the filter x>0.5" 4. y = 0.5 + @ROOT.Numba.Declare(["double"], "bool") def x_more_than_y(x): - x>y + return x>y rdf.Filter(x_more_than_y) Here the value of y is captured from scope. Thus this is equivalent to the filter "x>0.5". Note: Any modifications to the value of y will not be reflected in the function as y would be treated as a compile time constant during the first jit. @@ -390,24 +398,32 @@ def _PyDefine(rdf, col_name, callable_or_str, cols=[], extra_args={}): 1. col_name: The name of the new column to be defined 2. callable_or_str: The definition of the contents of the new columns It can be either a python callable or a c style string. + A python callable runs in the Python interpreter and needs type annotations for its parameters and + return value, unless it is decorated with ROOT.Numba.Declare to be JIT compiled with Numba. 3. cols: list of columns that the callable will receive as argument. If not provided then it tries maps the name of the parameter to a column name of the RDF. - 4. extra_args: non-columnar arguments to be passed to the callable. + 4. extra_args: non-columnar arguments to be passed to a callable decorated with ROOT.Numba.Declare. Returns: RDataFrame: rdf with new column defined Examples: - 1. rdf.Define("x", lambda: np.random.rand()) - Define using a python lambda - 2. rdf.Define("x2", lambda x: x*x, ["y"]) + 1. def random() -> float: + return np.random.rand() + rdf.Define("x", random) + Define using a python callable + 2. def square(x: float) -> float: + return x*x + rdf.Define("x2", square, ["y"]) Maps the function argument to column y. 3. scale = 0.5 + @ROOT.Numba.Declare(["double"], "double") def x_scaled(x): return x*scale rdf.Define("X_scaled", x_scaled) The scale is captured from scope and treated as a compile time constant during first jit. Note: Any changes to scale will not be reflected in the jitted function as jitting is a one time operation. - 4. def x_scaled(x, u): + 4. @ROOT.Numba.Declare(["double", "double"], "double") + def x_scaled(x, u): return x*u rdf.Define("x_scaled", x_scaled, extra_args = {"u":0.5}) From bc96ef0a966f84cd74803ee992769f2ca9c89a93 Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Tue, 13 Oct 2026 19:08:24 +0200 Subject: [PATCH 3/4] [Python][RDF] Do not declare callables decorated with ROOT.Numba.Declare twice `FunctionJitter` declared again the callables already declared with `ROOT.Numba.Declare`, which failed with a redefinition error, so passing a decorated callable with `extra_args` to `Define` or `Filter` did not work. Assisted-by: Claude Opus 5.5 --- .../pythonizations/python/ROOT/_pythonization/_rdf_pyz.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rdf_pyz.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rdf_pyz.py index bce12a5e73cd7..ef7367548e3c2 100755 --- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rdf_pyz.py +++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rdf_pyz.py @@ -225,7 +225,8 @@ def jit_function(self, func, cols_list, extra_args): return self.func_call self.get_function_params_args_call(func, cols_list, extra_args) - _NumbaDeclareDecorator(self.func_sign, self.return_type)(self.func) + if not hasattr(self.func, "__cpp_wrapper__"): # not already declared with ROOT.Numba.Declare + _NumbaDeclareDecorator(self.func_sign, self.return_type)(self.func) FunctionJitter.function_cache[self.func.__name__] = (self.func_call, self.func_sign) return self.func_call From 403521900c63a9448892015db1c4826525bd5fb7 Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Tue, 13 Oct 2026 19:08:24 +0200 Subject: [PATCH 4/4] [Python][RDF] Update the Define and Filter tests for Python callables Assisted-by: Claude Opus 5.5 --- .../pythonizations/test/numbadeclare.py | 32 ++++-- .../pythonizations/test/rdf_define_pyz.py | 108 +++++++++++++++++- .../pythonizations/test/rdf_filter_pyz.py | 36 +++++- 3 files changed, 152 insertions(+), 24 deletions(-) diff --git a/bindings/pyroot/pythonizations/test/numbadeclare.py b/bindings/pyroot/pythonizations/test/numbadeclare.py index a1153b19b52af..c8bc5684356e3 100644 --- a/bindings/pyroot/pythonizations/test/numbadeclare.py +++ b/bindings/pyroot/pythonizations/test/numbadeclare.py @@ -654,15 +654,17 @@ def test_fund_types(self): with self.subTest("function"): + @ROOT.Numba.Declare(["unsigned long"]) def is_even(x): return x % 2 == 0 - df = df.Define("is_even_x_1", is_even, ["x"]) + df = df.Define("is_even_x_1", "Numba::is_even(x)") results = df.Take["bool"]("is_even_x_1").GetValue()[0] self.assertEqual(results, True) with self.subTest("lambda"): - df = df.Define("is_even_x_2", lambda x: x % 2 == 0, ["x"]) + ROOT.Numba.Declare(["unsigned long"], name="is_even_lambda")(lambda x: x % 2 == 0) + df = df.Define("is_even_x_2", "Numba::is_even_lambda(x)") results = df.Take["bool"]("is_even_x_2").GetValue()[0] self.assertEqual(results, True) @@ -674,15 +676,17 @@ def test_rvec(self): with self.subTest("function"): + @ROOT.Numba.Declare(["RVec"]) def square_rvec(v): return v * v - df = df.Define("square_rvec_1", square_rvec, ["x"]) + df = df.Define("square_rvec_1", "Numba::square_rvec(x)") results = df.Take["RVec"]("square_rvec_1").GetValue()[0] self.assertTrue(np.array_equal(results, np.array([1, 4, 9]))) with self.subTest("lambda"): - df = df.Define("square_rvec_2", lambda v: v * v, ["x"]) + ROOT.Numba.Declare(["RVec"], name="square_rvec_lambda")(lambda v: v * v) + df = df.Define("square_rvec_2", "Numba::square_rvec_lambda(x)") results = df.Take["RVec"]("square_rvec_2").GetValue()[0] self.assertTrue(np.array_equal(results, np.array([1, 4, 9]))) @@ -694,15 +698,17 @@ def test_std_vec(self): with self.subTest("function"): + @ROOT.Numba.Declare(["std::vector"]) def square_std_vec(v): return v * v - df = df.Define("square_std_vec_1", square_std_vec, ["x"]) + df = df.Define("square_std_vec_1", "Numba::square_std_vec(x)") results = df.Take["RVec"]("square_std_vec_1").GetValue()[0] self.assertTrue(np.array_equal(results, np.array([1, 4, 9]))) with self.subTest("lambda"): - df = df.Define("square_std_vec_2", lambda v: v * v, ["x"]) + ROOT.Numba.Declare(["std::vector"], name="square_std_vec_lambda")(lambda v: v * v) + df = df.Define("square_std_vec_2", "Numba::square_std_vec_lambda(x)") results = df.Take["RVec"]("square_std_vec_2").GetValue()[0] self.assertTrue(np.array_equal(results, np.array([1, 4, 9]))) @@ -714,15 +720,17 @@ def test_std_array(self): with self.subTest("function"): + @ROOT.Numba.Declare(["std::array"]) def square_std_arr(v): return v * v - df = df.Define("square_std_arr_1", square_std_arr, ["x"]) + df = df.Define("square_std_arr_1", "Numba::square_std_arr(x)") results = df.Take["RVec"]("square_std_arr_1").GetValue()[0] self.assertTrue(np.array_equal(results, np.array([1, 4, 9]))) with self.subTest("lambda"): - df = df.Define("square_std_arr_2", lambda v: v * v, ["x"]) + ROOT.Numba.Declare(["std::array"], name="square_std_arr_lambda")(lambda v: v * v) + df = df.Define("square_std_arr_2", "Numba::square_std_arr_lambda(x)") results = df.Take["RVec"]("square_std_arr_2").GetValue()[0] self.assertTrue(np.array_equal(results, np.array([1, 4, 9]))) @@ -732,11 +740,11 @@ def test_missing_signature_raises(self): and no explicit signature is provided in the decorator. """ - def f(x): - return x.M() - with self.assertRaises(Exception): - ROOT.RDataFrame(4).Define("v", "ROOT::Math::PtEtaPhiMVector(1, 2, 3, 4)").Define("m", f, ["v"]) + + @ROOT.Numba.Declare(["ROOT::Math::PtEtaPhiMVector"]) + def f(x): + return x.M() if __name__ == "__main__": diff --git a/bindings/pyroot/pythonizations/test/rdf_define_pyz.py b/bindings/pyroot/pythonizations/test/rdf_define_pyz.py index 0a45b47df71dc..c4f8bf2925292 100644 --- a/bindings/pyroot/pythonizations/test/rdf_define_pyz.py +++ b/bindings/pyroot/pythonizations/test/rdf_define_pyz.py @@ -7,6 +7,11 @@ # numba is not used directly, but tests can crash when ROOT is built with # builtin_llvm=OFF and numba is not imported at the beginning +# C++ type names, used as type annotations of the Python callables passed to Define +DOUBLE = "double" +ULONG64 = "ULong64_t" +RVEC_ULONG64 = "ROOT::VecOps::RVec" + class PyDefine(unittest.TestCase): """ @@ -22,7 +27,11 @@ def test_with_dtypes(self): for type in numba_declare_dtypes: col_name = "col_" + type.replace(" ", "") rdf = rdf.Define(col_name, f"({type}) rdfentry_") - rdf = rdf.Define(col_name + "_arr", lambda col: np.array([col, col]), [col_name]) + + def to_array(col: type) -> f"ROOT::VecOps::RVec<{type}>": + return [col, col] + + rdf = rdf.Define(col_name + "_arr", to_array, [col_name]) arr = np.arange(0, 10) if type == "bool": arr = np.array(arr, dtype="bool") @@ -35,14 +44,22 @@ def test_with_dtypes(self): def test_define_overload1(self): rdf = ROOT.RDataFrame(10).Define("x", "rdfentry_") - rdf = rdf.Define("x2", lambda y: y * y, ["x"]) + + def square(y: ULONG64) -> ULONG64: + return y * y + + rdf = rdf.Define("x2", square, ["x"]) arr = np.arange(0, 10) flag = np.array_equal(rdf.AsNumpy()["x2"], arr * arr) self.assertTrue(flag) def test_define_overload2(self): rdf = ROOT.RDataFrame(10).Define("x", "rdfentry_") - rdf = rdf.Define("x2", lambda x: x * x) + + def square(x: ULONG64) -> ULONG64: + return x * x + + rdf = rdf.Define("x2", square) arr = np.arange(0, 10) flag = np.array_equal(rdf.AsNumpy()["x2"], arr * arr) self.assertTrue(flag) @@ -50,6 +67,7 @@ def test_define_overload2(self): def test_define_extra_args(self): rdf = ROOT.RDataFrame(10).Define("x", "rdfentry_") + @ROOT.Numba.Declare(["unsigned long", "double"], "double") def x_y(x, y): return x * y @@ -58,11 +76,23 @@ def x_y(x, y): flag = np.array_equal(rdf.AsNumpy()["x_y"], arr * 0.5) self.assertTrue(flag) + def test_python_callable_extra_args(self): + """ + Test that extra_args are rejected for a Python callable not decorated with ROOT.Numba.Declare. + """ + rdf = ROOT.RDataFrame(10).Define("x", "rdfentry_") + + def x_y(x: ULONG64, y: DOUBLE) -> DOUBLE: + return x * y + + with self.assertRaisesRegex(TypeError, "extra_args"): + rdf.Define("x_y", x_y, ["x"], extra_args={"y": 0.5}) + def test_capture_from_scope(self): rdf = ROOT.RDataFrame(10).Define("x", "rdfentry_") y = 0.5 - def x_times_y(x): + def x_times_y(x: ULONG64) -> DOUBLE: return x * y rdf = rdf.Define("x_y", x_times_y) @@ -72,9 +102,13 @@ def x_times_y(x): def test_arrays(self): rdf = ROOT.RDataFrame(5).Define("x", "rdfentry_") - rdf = rdf.Define("x_arr", lambda x: np.array([x, x])) - def norm(x_arr): + def to_array(x: ULONG64) -> RVEC_ULONG64: + return np.array([x, x]) + + rdf = rdf.Define("x_arr", to_array) + + def norm(x_arr: RVEC_ULONG64) -> DOUBLE: return np.sqrt(x_arr[0] ** 2 + x_arr[1] ** 2) rdf = rdf.Define("mag", norm) @@ -238,6 +272,68 @@ def test_cpp_free_function_template(self): for x, y in zip(rdf.Take["ULong64_t"]("rdfentry_"), rdf.Take["ULong64_t"]("new_col")): self.assertEqual(x, y) + def test_python_callables_different_signatures(self): + """ + Test that each Python callable is instantiated with the signature of its own annotations, + not with the one of a previous Python callable. + """ + + def times_two(x: float) -> float: + return 2 * x + + def times_three(i: int) -> int: + return 3 * i + + def is_positive(x: float, i: int) -> bool: + return x > 0 and i > 0 + + def constant() -> float: + return 7.0 + + rdf = ( + ROOT + .RDataFrame(3) + .Define("x", "(double) rdfentry_") + .Define("i", "(int) rdfentry_") + .Define("x2", times_two, ["x"]) + .Define("i3", times_three, ["i"]) + .Define("pos", is_positive, ["x", "i"]) + .Define("c", constant) + ) + res = rdf.AsNumpy(["x2", "i3", "pos", "c"]) + + self.assertEqual(rdf.GetColumnType("i3"), "int") + self.assertEqual(rdf.GetColumnType("pos"), "bool") + self.assertTrue(np.array_equal(res["x2"], [0.0, 2.0, 4.0])) + self.assertTrue(np.array_equal(res["i3"], [0, 3, 6])) + self.assertTrue(np.array_equal(res["pos"], [False, True, True])) + self.assertTrue(np.array_equal(res["c"], [7.0, 7.0, 7.0])) + + def test_python_callable_root_class(self): + """ + Test that a ROOT class can be used as type annotation of a Python callable. + """ + + def get_mass(v: ROOT.Math.PtEtaPhiMVector) -> float: + return v.M() + + rdf = ( + ROOT + .RDataFrame(5) + .Define("v", "ROOT::Math::PtEtaPhiMVector(1., 0., 0., (double) rdfentry_)") + .Define("mass", get_mass, ["v"]) + ) + + self.assertTrue(np.allclose(rdf.AsNumpy(["mass"])["mass"], np.arange(5))) + + def test_python_callable_missing_annotations(self): + """ + Test that a Python callable without type annotations is rejected. + """ + rdf = ROOT.RDataFrame(3).Define("x", "(double) rdfentry_") + with self.assertRaisesRegex(TypeError, "needs type annotations"): + rdf.Define("y", lambda x: 2 * x, ["x"]) + if __name__ == "__main__": unittest.main() diff --git a/bindings/pyroot/pythonizations/test/rdf_filter_pyz.py b/bindings/pyroot/pythonizations/test/rdf_filter_pyz.py index 66106a12f1dfc..edfcc181d3f37 100755 --- a/bindings/pyroot/pythonizations/test/rdf_filter_pyz.py +++ b/bindings/pyroot/pythonizations/test/rdf_filter_pyz.py @@ -9,6 +9,9 @@ # numba is not used directly, but tests can crash when ROOT is built with # builtin_llvm=OFF and numba is not imported at the beginning +# C++ type name, used as type annotation of the Python callables passed to Filter +DOUBLE = "double" + class PyFilter(unittest.TestCase): """ @@ -27,7 +30,11 @@ def test_with_dtypes(self): x = rdf.Mean(col_name).GetValue() if col_name == "Bool_t": x = True - filtered = rdf.Filter(func, extra_args={"x": x}) + + def filter_x(col: rdf.GetColumnType(col_name)) -> bool: + return func(col, x) # x is referenced from the enclosing scope + + filtered = rdf.Filter(filter_x, [col_name]) res_root = filtered.AsNumpy()[col_name] if not isinstance(x, bool): filtered2 = rdf.Filter(f"{col_name} > {x}") @@ -49,7 +56,7 @@ def test_filter_overload1_a(self): """ rdf = ROOT.RDataFrame(5).Define("x", "(double) rdfentry_") - def x_greater_than_2(x): + def x_greater_than_2(x: DOUBLE) -> bool: return x > 2 fil1 = rdf.Filter(x_greater_than_2, ["x"], "x is more than 2") @@ -61,7 +68,11 @@ def test_filter_overload1_b(self): Filter(callable, col_list) """ rdf = ROOT.RDataFrame(5).Define("x", "(double) rdfentry_") - fil1 = rdf.Filter(lambda x: x > 2, ["x"]) + + def x_greater_than_2(x: DOUBLE) -> bool: + return x > 2 + + fil1 = rdf.Filter(x_greater_than_2, ["x"]) self.assertTrue(np.array_equal(fil1.AsNumpy()["x"], np.array([3, 4]))) def test_filter_overload1_c(self): @@ -71,7 +82,7 @@ def test_filter_overload1_c(self): """ rdf = ROOT.RDataFrame(5).Define("x", "(double) rdfentry_") - def x_greater_than_2(x): + def x_greater_than_2(x: DOUBLE) -> bool: return x > 2 fil1 = rdf.Filter(x_greater_than_2) @@ -85,7 +96,7 @@ def test_filter_overload3(self): """ rdf = ROOT.RDataFrame(5).Define("x", "(double) rdfentry_") - def x_greater_than_2(x): + def x_greater_than_2(x: DOUBLE) -> bool: return x > 2 fil1 = rdf.Filter(x_greater_than_2, "x is greater than 2") @@ -95,12 +106,25 @@ def test_capture_from_scope(self): rdf = ROOT.RDataFrame(5).Define("x", "(double) rdfentry_") y = 2 - def x_greater_than_y(x): + def x_greater_than_y(x: DOUBLE) -> bool: return x > y fil1 = rdf.Filter(x_greater_than_y, "x is greater than 2") self.assertTrue(np.array_equal(fil1.AsNumpy()["x"], np.array([3, 4]))) + def test_numba_declare_extra_args(self): + """ + Test a callable decorated with ROOT.Numba.Declare with extra_args. + """ + rdf = ROOT.RDataFrame(5).Define("x", "(double) rdfentry_") + + @ROOT.Numba.Declare(["double", "double"], "bool") + def x_greater_than_y(x, y): + return x > y + + fil1 = rdf.Filter(x_greater_than_y, ["x"], extra_args={"y": 2}) + self.assertTrue(np.array_equal(fil1.AsNumpy()["x"], np.array([3, 4]))) + def test_cpp_functor(self): """ Test that a C++ functor can be passed as a callable argument of a