[Python][RDF] Add tests

Author: siliataider

bindings/pyroot/pythonizations/test/numbadeclare.py

@@ -633,5 +633,67 @@ def pass_reference(v):
         self.assertTrue(np.array_equal(rvecf, np.array([1.0, 4.0])))
 
 
+class NumbaDeclareInferred(unittest.TestCase):
+    """
+    Test decorator created with a reconstructed list of arguments using RDF column types,
+    and a return type inferred from the numba jitted function.
+    """
+
+    def test_fund_types(self):
+        """
+        Test fundamental types
+        """
+        def is_even(x):
+            return x % 2 == 0
+
+        df = ROOT.RDataFrame(4).Define("x", "rdfentry_").Define("is_even_x", is_even, ["x"])
+        results = df.Take["bool"]("is_even_x").GetValue()[0]
+        self.assertEqual(results, True)
+
+    def test_rvec(self):
+        """
+        Test RVec
+        """
+        def square_rvec(v):
+            return v*v
+
+        df = ROOT.RDataFrame(4).Define("x", "ROOT::VecOps::RVec<int>({1, 2, 3})").Define("square_rvec", square_rvec, ["x"])
+        results = df.Take["RVec<int>"]("square_rvec").GetValue()[0]
+        self.assertTrue(np.array_equal(results, np.array([1, 4, 9])))
+
+    def test_std_vec(self):
+        """
+        Test std::vector
+        """
+        def square_std_vec(v):
+            return v*v
+
+        df = ROOT.RDataFrame(4).Define("x", "std::vector<int>({1, 2, 3})").Define("square_std_vec", square_std_vec, ["x"])
+        results = df.Take["RVec<int>"]("square_std_vec").GetValue()[0]
+        self.assertTrue(np.array_equal(results, np.array([1, 4, 9])))
+
+    def test_std_array(self):
+        """
+        Test std::array
+        """
+        def square_std_arr(v):
+            return v*v
+
+        df = ROOT.RDataFrame(4).Define("x", "std::array<int, 3>({1, 2, 3})").Define("square_std_arr", square_std_arr, ["x"])
+        results = df.Take["RVec<int>"]("square_std_arr").GetValue()[0]
+        self.assertTrue(np.array_equal(results, np.array([1, 4, 9])))
+
+    def test_missing_signature_raises(self):
+        """
+        Ensure an Exception is raised when return type cannot be inferred
+        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"])
+
+
 if __name__ == "__main__":
     unittest.main()