[Python][ML] Accept file paths as input of the RDataLoader

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py

@@ -16,17 +16,21 @@
 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
     import torch
 
     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,8 @@ 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],
+        *,
         batch_size: int = 0,
         batches_in_memory: int = 1,
         columns: list[str] | None = None,
@@ -155,8 +184,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]):
+                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 +239,10 @@ def __init__(
 
         from ROOT import RDF
 
-        if rdataframes is None:
-            rdataframes = []
+        rdataframes = [] if inputs is None else _open_files(inputs)
+        if not rdataframes:
+            raise ValueError("Detected empty input, at least one valid input is required.")
 
-        if not hasattr(rdataframes, "__iter__"):
-            rdataframes = [rdataframes]
         self.noded_rdfs = [RDF.AsRNode(rdf) for rdf in rdataframes]
         _raise_if_implicit_mt(self.noded_rdfs)
 
@@ -596,7 +625,8 @@ 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 +646,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 +691,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,

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()

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,

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,

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,

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,
 )