[DistRDF] Support return_filenames and merge options in Snapshot (fixes #17136)

bindings/distrdf/python/DistRDF/Backends/Utils.py

@@ -248,7 +248,13 @@ def _(resultptr):
     of the dataset and the path to the partial snapshot. We can directly return
     the object, no extra work needed.
     """
-    return SnapshotResult(resultptr.treename, resultptr.filenames)
+    return SnapshotResult(
+        resultptr.treename,
+        resultptr.filenames,
+        return_filenames=resultptr.return_filenames,
+        merge=resultptr.merge,
+        target_file=resultptr.target_file,
+    )
 
 
 @singledispatch
@@ -333,4 +339,11 @@ def _(snap, range_id: int):
     # Actually clone the RDF C++ Snapshot node with the new file name
     resptr = ROOT.Internal.RDF.CloneResultAndAction(snap._resultptr, path_with_range)
 
-    return SnapshotResult(snap.treename, [path_with_range], resptr)
+    return SnapshotResult(
+        snap.treename,
+        [path_with_range],
+        resptr,
+        return_filenames=snap.return_filenames,
+        merge=snap.merge,
+        target_file=snap.target_file,
+    )

bindings/distrdf/python/DistRDF/ComputationGraphGenerator.py

@@ -50,7 +50,16 @@ def _(operation: Union[Action, InstantAction], promise: Any, results: list) -> N
 
 @append_node_to_results.register
 def _(operation: Snapshot, promise: Any, results: list) -> None:
-    results.append(SnapshotResult(operation.args[0], [operation.args[1]], promise))
+    results.append(
+        SnapshotResult(
+            operation.args[0],
+            [operation.args[1]],
+            promise,
+            return_filenames=getattr(operation, "return_filenames", False),
+            merge=getattr(operation, "merge", False),
+            target_file=getattr(operation, "target_file", None),
+        )
+    )
 
 
 @singledispatch
@@ -103,9 +112,9 @@ def _(operation: Snapshot, range_id: int) -> Snapshot:
         # Snapshot(std::string_view treename, std::string_view filename, std::string_view columnNameRegexp = "")
         op_modified.args.append("")  # Append empty regex
 
-    if len(op_modified.args) == 4:
-        # An RSnapshotOptions instance was passed as fourth argument
-        # Make it lazy and keep the other options
+    if len(op_modified.args) == 3 and hasattr(op_modified.args[2], "fLazy"):
+        op_modified.args[2].fLazy = True
+    elif len(op_modified.args) == 4 and hasattr(op_modified.args[3], "fLazy"):
         op_modified.args[3].fLazy = True
     else:
         # We already appended an empty regex for the 2 mandatory arguments overload

bindings/distrdf/python/DistRDF/Operation.py

@@ -95,7 +95,12 @@ class AsNumpy(InstantAction):
 class Snapshot(InstantAction):
     """A 'Snapshot' instant action attached to a distributed RDataFrame graph node."""
 
-    pass
+    def __init__(self, name: str, *args, **kwargs):
+        self.return_filenames = kwargs.pop("return_filenames", False)
+        self.merge = kwargs.pop("merge", False)
+        self.lazy = kwargs.pop("lazy", False)
+        self.target_file = str(args[1]) if len(args) > 1 and args[1] is not None else None
+        super().__init__(name, *args, **kwargs)
 
 
 class Transformation(Operation):

bindings/distrdf/python/DistRDF/Proxy.py

@@ -295,17 +295,24 @@ def _(operation: Operation.InstantAction, node: Node) -> Any:
 
 @get_proxy_for.register
 def _(operation: Operation.Snapshot, node: Node) -> Union[ResultPtrProxy, Any]:
-    if len(operation.args) == 4:
-        # An RSnapshotOptions instance was passed as fourth argument
-
-        if str(operation.args[3].fMode).lower() == "update":
-            # UPDATE mode is not supported since it would required synchronised access to the same file from the
+    opts = None
+    if len(operation.args) == 4 and hasattr(operation.args[3], "fLazy"):
+        opts = operation.args[3]
+    elif len(operation.args) == 3 and hasattr(operation.args[2], "fLazy"):
+        opts = operation.args[2]
+
+    if opts is not None:
+        if str(opts.fMode).lower() == "update":
+            # UPDATE mode is not supported since it would require synchronised access to the same file from the
             # distributed workers
             raise ValueError("Opening a file in UPDATE mode is not supported in distributed Snapshot.")
 
-        if operation.args[3].fLazy:
+        if opts.fLazy:
             return get_proxy_for.dispatch(Operation.Action)(operation, node)
 
+    if getattr(operation, "lazy", False):
+        return get_proxy_for.dispatch(Operation.Action)(operation, node)
+
     return get_proxy_for.dispatch(Operation.InstantAction)(operation, node)
 
 

bindings/distrdf/python/DistRDF/PythonMergeables.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, List, Union
+from typing import TYPE_CHECKING, List, Optional, Union
 
+import os
 import ROOT
 from ROOT._pythonization._rdataframe import AsNumpyResult
 
@@ -15,13 +16,24 @@ class SnapshotResult(object):
     merge it with other objects of this type.
     """
 
-    def __init__(self, treename: str, filenames: List[str], resultptr: ROOT.RDF.RResultPtr = None) -> None:
+    def __init__(
+        self,
+        treename: str,
+        filenames: List[str],
+        resultptr: ROOT.RDF.RResultPtr = None,
+        return_filenames: bool = False,
+        merge: bool = False,
+        target_file: Optional[str] = None,
+    ) -> None:
         self.treename = treename
         self.filenames = filenames
         # Transient attribute, it will be discarded before the end of the mapper
         # function (in `Utils.get_mergeablevalue`) so that we don't incur in
         # serialization of the RResultPtr
         self._resultptr = resultptr
+        self.return_filenames = return_filenames
+        self.merge = merge
+        self.target_file = target_file
 
     def Merge(self, other: SnapshotResult) -> None:
         """
@@ -46,7 +58,44 @@ def GetValue(self, backend: BaseBackend):
         computation was created via a Spark backend, then this function will
         return another distributed RDataFrame build from a Spark backend
         instance. And so on for all other DistRDF backends.
+
+        If `merge` is True, the partial snapshot files are merged into the
+        target file using `ROOT.TFileMerger`.
+
+        If `return_filenames` is True, the list of snapshot file paths is
+        returned instead of a distributed RDataFrame.
         """
+        if self.merge and self.target_file and self.filenames:
+            merger = ROOT.TFileMerger(False)
+            if not merger.OutputFile(self.target_file):
+                raise RuntimeError(f"Could not open output file '{self.target_file}' for merging.")
+            for filename in self.filenames:
+                if not merger.AddFile(filename):
+                    raise RuntimeError(f"Could not add file '{filename}' to merger.")
+            if not merger.Merge():
+                raise RuntimeError(f"Failed to merge snapshot files into '{self.target_file}'.")
+            if hasattr(merger, "CloseOutputFile"):
+                merger.CloseOutputFile()
+
+            for filename in self.filenames:
+                if filename != self.target_file and os.path.exists(filename):
+                    try:
+                        os.remove(filename)
+                    except OSError:
+                        pass
+
+            self.filenames = [self.target_file]
+
+            if self.return_filenames:
+                return list(self.filenames)
+
+            snapshot_chain = ROOT.TChain(self.treename)
+            snapshot_chain.Add(self.target_file)
+            return backend.make_dataframe(snapshot_chain)
+
+        if self.return_filenames:
+            return list(self.filenames)
+
         snapshot_chain = ROOT.TChain(self.treename)
         # Add partial snapshot files to the chain
         for filename in self.filenames:

bindings/distrdf/test/backend/test_dist.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 
 from ROOT._distrdf import DataFrame, HeadNode
@@ -58,7 +59,7 @@ def distribute_unique_paths(self, includes_list):
 
         def make_dataframe(self, *args, **kwargs):
             """Dummy make_dataframe"""
-            pass
+            return args[0]
 
         def optimize_npartitions(self):
             pass
@@ -80,3 +81,67 @@ def test_count_result_invariance(self):
             headnode = HeadNode.get_headnode(backend, npartitions, treename, filenames)
             rdf = DataFrame.RDataFrame(headnode)
             self.assertEqual(rdf.Count().GetValue(), 100)
+
+    def test_snapshot_return_filenames(self):
+        """Test Snapshot with return_filenames=True returns list of partial files."""
+        treename = "entries"
+        filenames = ["1cluster_20entries.root"] * 2
+        backend = DistRDataFrameInvariants.TestBackend()
+        headnode = HeadNode.get_headnode(backend, 2, treename, filenames)
+        rdf = DataFrame.RDataFrame(headnode)
+        output_file = "test_snapshot_return_filenames_output.root"
+        try:
+            files = rdf.Snapshot(treename, output_file, return_filenames=True)
+            self.assertIsInstance(files, list)
+            self.assertEqual(len(files), 2)
+            for f in files:
+                self.assertTrue(os.path.exists(f))
+        finally:
+            for i in range(2):
+                p = f"test_snapshot_return_filenames_output_{i}.root"
+                if os.path.exists(p):
+                    os.remove(p)
+
+    def test_snapshot_merge(self):
+        """Test Snapshot with merge=True merges partial files into target file."""
+        treename = "entries"
+        filenames = ["1cluster_20entries.root"] * 2
+        backend = DistRDataFrameInvariants.TestBackend()
+        headnode = HeadNode.get_headnode(backend, 2, treename, filenames)
+        rdf = DataFrame.RDataFrame(headnode)
+        output_file = "test_snapshot_merge_output.root"
+        try:
+            rdf.Snapshot(treename, output_file, merge=True)
+            self.assertTrue(os.path.exists(output_file))
+            for i in range(2):
+                self.assertFalse(os.path.exists(f"test_snapshot_merge_output_{i}.root"))
+        finally:
+            if os.path.exists(output_file):
+                os.remove(output_file)
+            for i in range(2):
+                p = f"test_snapshot_merge_output_{i}.root"
+                if os.path.exists(p):
+                    os.remove(p)
+
+    def test_snapshot_merge_and_return_filenames(self):
+        """Test Snapshot with merge=True and return_filenames=True."""
+        treename = "entries"
+        filenames = ["1cluster_20entries.root"] * 2
+        backend = DistRDataFrameInvariants.TestBackend()
+        headnode = HeadNode.get_headnode(backend, 2, treename, filenames)
+        rdf = DataFrame.RDataFrame(headnode)
+        output_file = "test_snapshot_merge_return_output.root"
+        try:
+            files = rdf.Snapshot(treename, output_file, merge=True, return_filenames=True)
+            self.assertEqual(files, [output_file])
+            self.assertTrue(os.path.exists(output_file))
+            for i in range(2):
+                self.assertFalse(os.path.exists(f"test_snapshot_merge_return_output_{i}.root"))
+        finally:
+            if os.path.exists(output_file):
+                os.remove(output_file)
+            for i in range(2):
+                p = f"test_snapshot_merge_return_output_{i}.root"
+                if os.path.exists(p):
+                    os.remove(p)
+

bindings/distrdf/test/test_operation.py

@@ -122,3 +122,48 @@ def test_histonsparsed_without_model(self):
         """Creating a histogram without model raises ValueError."""
         with self.assertRaises(ValueError):
             _ = Operation.create_op("HistoNSparseD", ["a", "b", "c", "d"])
+
+
+class SnapshotTests(unittest.TestCase):
+    """
+    Test Snapshot operation construction with different options.
+    """
+
+    def test_snapshot_default_options(self):
+        """Snapshot creates with default return_filenames=False and merge=False."""
+        op = Operation.create_op("Snapshot", "tree", "output.root")
+        self.assertIsInstance(op, Operation.Snapshot)
+        self.assertEqual(op.target_file, "output.root")
+        self.assertFalse(op.return_filenames)
+        self.assertFalse(op.merge)
+        self.assertNotIn("return_filenames", op.kwargs)
+        self.assertNotIn("merge", op.kwargs)
+
+    def test_snapshot_with_return_filenames(self):
+        """Snapshot handles return_filenames keyword argument."""
+        op = Operation.create_op("Snapshot", "tree", "output.root", return_filenames=True)
+        self.assertIsInstance(op, Operation.Snapshot)
+        self.assertEqual(op.target_file, "output.root")
+        self.assertTrue(op.return_filenames)
+        self.assertFalse(op.merge)
+        self.assertNotIn("return_filenames", op.kwargs)
+
+    def test_snapshot_with_merge(self):
+        """Snapshot handles merge keyword argument."""
+        op = Operation.create_op("Snapshot", "tree", "output.root", merge=True)
+        self.assertIsInstance(op, Operation.Snapshot)
+        self.assertEqual(op.target_file, "output.root")
+        self.assertFalse(op.return_filenames)
+        self.assertTrue(op.merge)
+        self.assertNotIn("merge", op.kwargs)
+
+    def test_snapshot_with_both_options(self):
+        """Snapshot handles both return_filenames and merge keyword arguments."""
+        op = Operation.create_op("Snapshot", "tree", "output.root", return_filenames=True, merge=True)
+        self.assertIsInstance(op, Operation.Snapshot)
+        self.assertEqual(op.target_file, "output.root")
+        self.assertTrue(op.return_filenames)
+        self.assertTrue(op.merge)
+        self.assertNotIn("return_filenames", op.kwargs)
+        self.assertNotIn("merge", op.kwargs)
+