diff --git a/bindings/distrdf/python/DistRDF/Backends/Utils.py b/bindings/distrdf/python/DistRDF/Backends/Utils.py index 6aa8d0a060cc2..c8de8abf65356 100644 --- a/bindings/distrdf/python/DistRDF/Backends/Utils.py +++ b/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, + ) diff --git a/bindings/distrdf/python/DistRDF/ComputationGraphGenerator.py b/bindings/distrdf/python/DistRDF/ComputationGraphGenerator.py index 52ca5a184421c..903720defeefe 100644 --- a/bindings/distrdf/python/DistRDF/ComputationGraphGenerator.py +++ b/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 diff --git a/bindings/distrdf/python/DistRDF/Operation.py b/bindings/distrdf/python/DistRDF/Operation.py index 577fa711cb99d..6a2ff48f161c5 100644 --- a/bindings/distrdf/python/DistRDF/Operation.py +++ b/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): diff --git a/bindings/distrdf/python/DistRDF/Proxy.py b/bindings/distrdf/python/DistRDF/Proxy.py index c1db4f6369daa..ebf1e622045ca 100644 --- a/bindings/distrdf/python/DistRDF/Proxy.py +++ b/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) diff --git a/bindings/distrdf/python/DistRDF/PythonMergeables.py b/bindings/distrdf/python/DistRDF/PythonMergeables.py index 8b4e0a27421dd..1472c11183bc3 100644 --- a/bindings/distrdf/python/DistRDF/PythonMergeables.py +++ b/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: diff --git a/bindings/distrdf/test/backend/test_dist.py b/bindings/distrdf/test/backend/test_dist.py index 19aabd0833d40..9208197d19c78 100644 --- a/bindings/distrdf/test/backend/test_dist.py +++ b/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) + diff --git a/bindings/distrdf/test/test_operation.py b/bindings/distrdf/test/test_operation.py index a12b794feff6b..b296cc0f40c68 100644 --- a/bindings/distrdf/test/test_operation.py +++ b/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) +