[Python][RDF] Add tests for Numba.Declare with std::vector/std::array

Author: siliataider

bindings/pyroot/pythonizations/test/numbadeclare.py

@@ -113,6 +113,31 @@ def pass_temporary(v):
 
         self.assertTrue(np.array_equal(rvecf, np.array([4.])))
 
+    def test_rdataframe_std_vector(self):
+        """
+        Test function call as part of RDataFrame
+        """
+        @ROOT.Numba.Declare(["std::vector<int>"], "std::vector<int>")
+        def square_vec(x):
+            return x * x
+        df = ROOT.RDataFrame(4).Define("x", "std::vector{1, 2, 3}").Define("x_sq", "Numba::square_vec(x)")
+        df.Display().Print()
+        self.assertEqual(df.Sum("x").GetValue(), 24)
+        self.assertEqual(df.Sum("x_sq").GetValue(), 56)
+
+    def test_rdataframe_std_array(self):
+        """
+        Test function call as part of RDataFrame with std::array
+        """
+        @ROOT.Numba.Declare(["std::array<int, 3>"], "std::array<int, 3>")
+        def square_array(x):
+            return x * x
+
+        df = ROOT.RDataFrame(4).Define("x", "std::array{1, 2, 3}").Define("x_sq", "Numba::square_array(x)")
+        df.Display().Print()
+        self.assertEqual(df.Sum("x").GetValue(), 24)
+        self.assertEqual(df.Sum("x_sq").GetValue(), 56)
+
     # Test wrappings
     def test_wrapper_in_void(self):
         """