From a9b570a3f12cba80ac7c7e15205dae68f8767d21 Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Tue, 20 Oct 2026 11:57:01 +0200 Subject: [PATCH 1/4] [Python][ML] Accept file paths as input of the RDataLoader --- .../ROOT/_pythonization/_ml_dataloader.py | 53 ++++++++++++++----- 1 file changed, 41 insertions(+), 12 deletions(-) diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py index bef5ea1683f02..8ff2df63585d5 100644 --- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py +++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py @@ -16,10 +16,11 @@ from __future__ import annotations import atexit +import os from typing import TYPE_CHECKING, Any, Callable, Tuple if TYPE_CHECKING: - import os + from collections.abc import Iterable import numpy as np import tensorflow as tf @@ -27,6 +28,9 @@ import ROOT + # a dataset read from files: the name of its TTree or RNTuple, and one or more file paths + Spec = Tuple[str, str | os.PathLike | Iterable[str | os.PathLike]] + def _normalize_target(target: str | list[str] | None) -> list[str]: """`target` can be a string, list of strings, or None. Normalize to a list of strings @@ -54,6 +58,30 @@ def _raise_if_implicit_mt(rdfs: list) -> None: ) +def _is_spec(item: Any) -> bool: + """Whether `item` is a `(dataset_name, files)` tuple.""" + return isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], str) + + +def _open_files(inputs: Any) -> list: + """Return the inputs as a list of RDataFrames, each `(dataset_name, files)` tuple being read into one RDataFrame.""" + import ROOT + + if _is_spec(inputs) or not hasattr(inputs, "__iter__"): + inputs = [inputs] + + rdataframes = [] + for item in inputs: + if isinstance(item, (str, os.PathLike)): + raise TypeError("File paths are given together with the dataset to read, as (dataset_name, files) tuples") + if _is_spec(item): + dataset_name, files = item + files = [files] if isinstance(files, (str, os.PathLike)) else list(files) + item = ROOT.RDataFrame(dataset_name, [os.fspath(f) for f in files]) + rdataframes.append(item) + return rdataframes + + class _RDataLoader: def _resolve_columns( self, @@ -135,7 +163,7 @@ def _resolve_columns( def __init__( self, - rdataframes: ROOT.RDF.RNode | list[ROOT.RDF.RNode] | None = None, + inputs: ROOT.RDF.RNode | Spec | Iterable[ROOT.RDF.RNode | Spec] | None = None, batch_size: int = 0, batches_in_memory: int = 1, columns: list[str] | None = None, @@ -155,8 +183,9 @@ def __init__( """Wrapper around the C++ DataLoader Args: - rdataframes (ROOT.RDF.RNode | list[ROOT.RDF.RNode] | None): - RDataFrame or list of RDataFrames to load from. + inputs (ROOT.RDF.RNode | Spec | Iterable[ROOT.RDF.RNode | Spec] | None): + RDataFrame or ``(dataset_name, files)`` tuple, or a list of them, to load from. + A tuple is used to construct a single RDataFrame to read ``dataset_name`` from one or more files. batch_size (int): Number of entries per batch returned by the generator. batches_in_memory (int): @@ -209,11 +238,10 @@ def __init__( from ROOT import RDF - if rdataframes is None: - rdataframes = [] + if inputs is None: + inputs = [] - if not hasattr(rdataframes, "__iter__"): - rdataframes = [rdataframes] + rdataframes = _open_files(inputs) self.noded_rdfs = [RDF.AsRNode(rdf) for rdf in rdataframes] _raise_if_implicit_mt(self.noded_rdfs) @@ -596,7 +624,7 @@ class RDataLoader: def __init__( self, - rdataframes: ROOT.RDF.RNode | list[ROOT.RDF.RNode], + inputs: ROOT.RDF.RNode | Spec | Iterable[ROOT.RDF.RNode | Spec], batch_size: int = 64, batches_in_memory: int = 10, columns: list[str] | None = None, @@ -616,8 +644,9 @@ def __init__( \ingroup Py_ML Args: - rdataframes: - RDataFrame or list of RDataFrames to load from. + inputs: + RDataFrame or (dataset_name, files) tuple, or a list of them, to load from. + A tuple is used to construct a single RDataFrame to read dataset_name from one or more files. batch_size: Number of entries per batch. batches_in_memory: @@ -660,7 +689,7 @@ def __init__( # created lazily on the first call to as_torch/as_numpy/as_tensorflow or # train_test_split. self._params = dict( - rdataframes=rdataframes, + inputs=inputs, batch_size=batch_size, batches_in_memory=batches_in_memory, columns=columns, From 661b5d4dabc737e96264114902688c6ad31b014b Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Wed, 21 Oct 2026 13:41:32 +0200 Subject: [PATCH 2/4] [Python][ML] Raise if no input is given to the RDataLoader --- .../python/ROOT/_pythonization/_ml_dataloader.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py index 8ff2df63585d5..8c0f7afd5e946 100644 --- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py +++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py @@ -163,7 +163,7 @@ def _resolve_columns( def __init__( self, - inputs: ROOT.RDF.RNode | Spec | Iterable[ROOT.RDF.RNode | Spec] | None = None, + inputs: ROOT.RDF.RNode | Spec | Iterable[ROOT.RDF.RNode | Spec], batch_size: int = 0, batches_in_memory: int = 1, columns: list[str] | None = None, @@ -183,7 +183,7 @@ def __init__( """Wrapper around the C++ DataLoader Args: - inputs (ROOT.RDF.RNode | Spec | Iterable[ROOT.RDF.RNode | Spec] | None): + inputs (ROOT.RDF.RNode | Spec | Iterable[ROOT.RDF.RNode | Spec]): RDataFrame or ``(dataset_name, files)`` tuple, or a list of them, to load from. A tuple is used to construct a single RDataFrame to read ``dataset_name`` from one or more files. batch_size (int): @@ -238,10 +238,10 @@ def __init__( from ROOT import RDF - if inputs is None: - inputs = [] + rdataframes = [] if inputs is None else _open_files(inputs) + if not rdataframes: + raise ValueError("Detected empty input, at least one valid input is required.") - rdataframes = _open_files(inputs) self.noded_rdfs = [RDF.AsRNode(rdf) for rdf in rdataframes] _raise_if_implicit_mt(self.noded_rdfs) From 575e419c2b6949d1b8ca21eff522fa3977cf5eae Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Wed, 21 Oct 2026 13:41:52 +0200 Subject: [PATCH 3/4] [Python][ML] Make all the RDataLoader arguments but the inputs keyword-only --- .../python/ROOT/_pythonization/_ml_dataloader.py | 2 ++ tutorials/machine_learning/ml_dataloader_NumPy.py | 2 +- tutorials/machine_learning/ml_dataloader_PyTorch.py | 2 +- tutorials/machine_learning/ml_dataloader_TensorFlow.py | 4 ++-- tutorials/machine_learning/ml_dataloader_filters_vectors.py | 2 +- 5 files changed, 7 insertions(+), 5 deletions(-) diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py index 8c0f7afd5e946..3590026c54a28 100644 --- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py +++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py @@ -164,6 +164,7 @@ def _resolve_columns( def __init__( self, inputs: ROOT.RDF.RNode | Spec | Iterable[ROOT.RDF.RNode | Spec], + *, batch_size: int = 0, batches_in_memory: int = 1, columns: list[str] | None = None, @@ -625,6 +626,7 @@ class RDataLoader: def __init__( self, inputs: ROOT.RDF.RNode | Spec | Iterable[ROOT.RDF.RNode | Spec], + *, batch_size: int = 64, batches_in_memory: int = 10, columns: list[str] | None = None, diff --git a/tutorials/machine_learning/ml_dataloader_NumPy.py b/tutorials/machine_learning/ml_dataloader_NumPy.py index 4751087b1ef1c..7f0d2e29e388c 100644 --- a/tutorials/machine_learning/ml_dataloader_NumPy.py +++ b/tutorials/machine_learning/ml_dataloader_NumPy.py @@ -23,7 +23,7 @@ dl = ROOT.Experimental.ML.RDataLoader( rdataframe, - batch_size, + batch_size=batch_size, target=target, shuffle=True, drop_remainder=True, diff --git a/tutorials/machine_learning/ml_dataloader_PyTorch.py b/tutorials/machine_learning/ml_dataloader_PyTorch.py index e401adbb4b834..fff5e21b7a079 100644 --- a/tutorials/machine_learning/ml_dataloader_PyTorch.py +++ b/tutorials/machine_learning/ml_dataloader_PyTorch.py @@ -24,7 +24,7 @@ # as PyTorch tensors. dl = ROOT.Experimental.ML.RDataLoader( rdataframe, - batch_size, + batch_size=batch_size, target=target, shuffle=True, drop_remainder=True, diff --git a/tutorials/machine_learning/ml_dataloader_TensorFlow.py b/tutorials/machine_learning/ml_dataloader_TensorFlow.py index 195f06a168079..d0056c6122d2b 100644 --- a/tutorials/machine_learning/ml_dataloader_TensorFlow.py +++ b/tutorials/machine_learning/ml_dataloader_TensorFlow.py @@ -22,8 +22,8 @@ # Returns two TF.Dataset for training and validation batches. dl = ROOT.Experimental.ML.RDataLoader( rdataframe, - batch_size, - approx_batches_in_memory, + batch_size=batch_size, + batches_in_memory=approx_batches_in_memory, target=target, shuffle=True, drop_remainder=True, diff --git a/tutorials/machine_learning/ml_dataloader_filters_vectors.py b/tutorials/machine_learning/ml_dataloader_filters_vectors.py index 1ee9d88caded5..12d9acef591db 100644 --- a/tutorials/machine_learning/ml_dataloader_filters_vectors.py +++ b/tutorials/machine_learning/ml_dataloader_filters_vectors.py @@ -29,7 +29,7 @@ dl = ROOT.Experimental.ML.RDataLoader( filteredrdf, - batch_size, + batch_size=batch_size, max_vec_sizes=max_vec_sizes, shuffle=False, ) From a9927b8ba8e7f0b513b8bfc7f7302d8f6c64896f Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Tue, 20 Oct 2026 11:57:34 +0200 Subject: [PATCH 4/4] [Python][ML] Add tests for the different inputs Assisted-by: Claude Opus 5.5 --- .../pythonizations/test/ml_dataloader.py | 126 ++++++++++++++++++ 1 file changed, 126 insertions(+) diff --git a/bindings/pyroot/pythonizations/test/ml_dataloader.py b/bindings/pyroot/pythonizations/test/ml_dataloader.py index c018ba7cdd50d..f60406361916e 100644 --- a/bindings/pyroot/pythonizations/test/ml_dataloader.py +++ b/bindings/pyroot/pythonizations/test/ml_dataloader.py @@ -8141,5 +8141,131 @@ def test01_implicit_mt_not_supported(self): ROOT.Experimental.ML.RDataLoader(df, batch_size=2).as_numpy() +class DataLoaderFromFiles(unittest.TestCase): + file_name1 = "dataloader_from_files1.root" + file_name2 = "dataloader_from_files2.root" + tree_name = "mytree" + + # Helpers + def create_file(self, file_name, num_of_entries=10, offset=0, output_format="ttree"): + opts = ROOT.RDF.RSnapshotOptions() + if output_format == "rntuple": + opts.fOutputFormat = ROOT.RDF.ESnapshotOutputFormat.kRNTuple + ( + ROOT + .RDataFrame(num_of_entries) + .Define("b1", f"(int) rdfentry_ + {offset}") + .Define("b2", "(double) b1 * b1") + .Snapshot(self.tree_name, file_name, ["b1", "b2"], opts) + ) + + def tearDown(self): + for file in (self.file_name1, self.file_name2): + if os.path.exists(file): + os.remove(file) + + def load_all(self, inputs, **kwargs): + """Concatenate the features and the target of all the batches of an unshuffled loader.""" + dl = ROOT.Experimental.ML.RDataLoader( + inputs, batch_size=3, batches_in_memory=2, target="b2", shuffle=False, drop_remainder=False, **kwargs + ) + batches = list(dl.as_numpy()) + x = np.concatenate([x for x, _ in batches]).flatten() + y = np.concatenate([y for _, y in batches]).flatten() + return x, y + + def test01_single_path(self): + self.create_file(self.file_name1) + + for load_eager in (False, True): + with self.subTest(load_eager=load_eager): + x, y = self.load_all((self.tree_name, self.file_name1), load_eager=load_eager) + np.testing.assert_array_equal(x, np.arange(10)) + np.testing.assert_array_equal(y, np.arange(10) ** 2) + + def test02_pathlike(self): + from pathlib import Path + + self.create_file(self.file_name1) + + x, y = self.load_all((self.tree_name, Path(self.file_name1))) + np.testing.assert_array_equal(x, np.arange(10)) + np.testing.assert_array_equal(y, np.arange(10) ** 2) + + def test03_list_of_paths(self): + from pathlib import Path + + self.create_file(self.file_name1, num_of_entries=9) + self.create_file(self.file_name2, num_of_entries=6, offset=9) + + # the files are read as a single dataset, in the given order + x, y = self.load_all((self.tree_name, [self.file_name1, Path(self.file_name2)])) + np.testing.assert_array_equal(x, np.arange(15)) + np.testing.assert_array_equal(y, np.arange(15) ** 2) + + def test04_glob(self): + self.create_file(self.file_name1) + self.create_file(self.file_name2, num_of_entries=5, offset=10) + + x, _ = self.load_all((self.tree_name, "dataloader_from_files*.root")) + np.testing.assert_array_equal(np.sort(x), np.arange(15)) + + def test05_rntuple(self): + self.create_file(self.file_name1, output_format="rntuple") + + x, y = self.load_all((self.tree_name, self.file_name1)) + np.testing.assert_array_equal(x, np.arange(10)) + np.testing.assert_array_equal(y, np.arange(10) ** 2) + + def test06_train_test_split(self): + self.create_file(self.file_name1) + + dl = ROOT.Experimental.ML.RDataLoader( + (self.tree_name, self.file_name1), + batch_size=2, + columns=["b1"], + shuffle=False, + drop_remainder=False, + ) + gen_train, gen_validation = dl.train_test_split(0.4) + + train_x = np.concatenate(list(gen_train.as_numpy())).flatten() + val_x = np.concatenate(list(gen_validation.as_numpy())).flatten() + np.testing.assert_array_equal(train_x, np.arange(6)) + np.testing.assert_array_equal(val_x, np.arange(6, 10)) + + def test07_invalid_inputs(self): + self.create_file(self.file_name1) + df = ROOT.RDataFrame(self.tree_name, self.file_name1) + + # file paths need the name of the dataset to read + with self.assertRaisesRegex(TypeError, "dataset_name, files"): + self.load_all(self.file_name1) + + with self.assertRaisesRegex(ValueError, "at least one valid input is required"): + self.load_all([]) + with self.assertRaisesRegex(ValueError, "at least one valid input is required"): + self.load_all(df for df in []) + + # all the arguments but the inputs are keyword-only + with self.assertRaises(TypeError): + ROOT.Experimental.ML.RDataLoader(df, 3) + + def test08_several_datasets(self): + self.create_file(self.file_name1, num_of_entries=9) + self.create_file(self.file_name2, num_of_entries=6, offset=9) + + # one RDataFrame per tuple, which can be mixed with RDataFrames + df = ROOT.RDataFrame(self.tree_name, self.file_name2) + for inputs in ( + [(self.tree_name, self.file_name1), (self.tree_name, self.file_name2)], + [(self.tree_name, self.file_name1), df], + ): + with self.subTest(inputs=inputs): + x, y = self.load_all(inputs) + np.testing.assert_array_equal(x, np.arange(15)) + np.testing.assert_array_equal(y, np.arange(15) ** 2) + + if __name__ == "__main__": unittest.main()