From 1edca86fbe6bc57f5583f27e53e999b6532b5005 Mon Sep 17 00:00:00 2001
From: Vincenzo Eduardo Padulano <vincenzo.eduardo.padulano@cern.ch>
Date: Thu, 28 Mar 2024 15:43:20 +0100
Subject: [PATCH] [df] Enable cloning of AsNumpyResult

---
 .../python/ROOT/_pythonization/_rdataframe.py | 15 +++++++++++
 .../pythonizations/test/rdataframe_asnumpy.py | 25 +++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rdataframe.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rdataframe.py
index 1a466d88c4e6e..6ca1bcaacd757 100644
--- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rdataframe.py
+++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rdataframe.py
@@ -345,6 +345,21 @@ def __setstate__(self, state):
         self._py_arrays = state
 
 
+def _clone_asnumpyresult(res: AsNumpyResult) -> AsNumpyResult:
+    """
+    Clones the internal actions held by the input result and returns a new
+    result.
+    """
+    import ROOT
+    return AsNumpyResult(
+        {
+            col: ROOT.Internal.RDF.CloneResultAndAction(ptr)
+            for (col, ptr) in res._result_ptrs.items()
+        },
+        res._columns
+    )
+
+
 class HistoProfileWrapper(MethodTemplateWrapper):
     """
     Subclass of MethodTemplateWrapper that pythonizes HistoXD and ProfileXD
diff --git a/bindings/pyroot/pythonizations/test/rdataframe_asnumpy.py b/bindings/pyroot/pythonizations/test/rdataframe_asnumpy.py
index 0f13563df150c..8c0670b2753c2 100644
--- a/bindings/pyroot/pythonizations/test/rdataframe_asnumpy.py
+++ b/bindings/pyroot/pythonizations/test/rdataframe_asnumpy.py
@@ -3,6 +3,8 @@
 import numpy as np
 import pickle
 
+from ROOT._pythonization._rdataframe import _clone_asnumpyresult
+
 
 def make_tree(*dtypes):
     """
@@ -308,6 +310,29 @@ def test_memory_adoption_complex_types(self):
         pyarr[0][0] = 42
         self.assertTrue(cpparr[0][0] == pyarr[0][0])
 
+    def test_cloning(self):
+        """
+        Testing cloning of AsNumpy results
+        """
+        df = ROOT.RDataFrame(20).Define("x", "rdfentry_")
+        ranges = [(0, 5), (5, 10), (10, 15), (15, 20)]
+
+        # Get the result for the first range
+        (begin, end) = ranges.pop(0)
+        ROOT.Internal.RDF.ChangeEmptyEntryRange(
+            ROOT.RDF.AsRNode(df), (begin, end))
+        asnumpyres = df.AsNumpy(["x"], lazy=True)  # To return an AsNumpyResult
+        self.assertSequenceEqual(
+            asnumpyres.GetValue()["x"].tolist(), np.arange(begin, end).tolist())
+
+        # Clone the result for following ranges
+        for (begin, end) in ranges:
+            ROOT.Internal.RDF.ChangeEmptyEntryRange(
+                ROOT.RDF.AsRNode(df), (begin, end))
+            asnumpyres = _clone_asnumpyresult(asnumpyres)
+            self.assertSequenceEqual(
+                asnumpyres.GetValue()["x"].tolist(), np.arange(begin, end).tolist())
+
 
 if __name__ == '__main__':
     unittest.main()