RNTuple support (#586)

ponyisi · web-flow · commit df8380293df5 · 2025-05-17T10:02:52.000-05:00
* RNTuple support
* Update tests and documentation
diff --git a/docs/transform_request.rst b/docs/transform_request.rst
@@ -32,7 +32,7 @@ The General Section
 ^^^^^^^^^^^^^^^^^^^
 The General section of the request includes the following fields:
 
-* (Optional) ``OutputFormat``: Can be ``root-ttree`` (default) or ``parquet`` (you can also use the enums ``servicex.OutputFormat.root_ttree`` and ``servicex.OutputFormat.parquet``)
+* (Optional) ``OutputFormat``: Can be ``root-ttree`` (default), ``root-rntuple``, or ``parquet`` (you can also use the enums ``servicex.OutputFormat.root_ttree``, ``servicex.OutputFormat.root_rntuple``, and ``servicex.OutputFormat.parquet``)
 * (Optional) ``Delivery``: Can be ``LocalCache`` (default) to download the results to the system on which you are running the client, or ``URLs`` to provide HTTPS URLs to the output on the ServiceX storage (you can also use the enums ``servicex.Delivery.LocalCache`` and ``servicex.Delivery.URLs``). The output files on the ServiceX storage will periodically get cleaned, so if you need to keep the results for long term use it is recommended that you download the output to local cache, but for transient use the URLs will be faster.
 
 In general, if you are running on your laptop away from the ServiceX site and are working with a small amount of
diff --git a/docs/transformer_matrix.md b/docs/transformer_matrix.md
@@ -33,7 +33,7 @@ In brief, here are the currently supported combinations:
 
 ### Uproot-Raw
 
-The Uproot-Raw query language can be thought of as essentially calling the [`TTree.arrays()`](https://uproot.readthedocs.io/en/latest/uproot.behaviors.TTree.TTree.html#arrays) function of `uproot` with the possibility to specify several of the arguments, and returning the result. Multiple queries can be bundled into one request. It is also possible to copy objects from the input file.
+The Uproot-Raw query language can be thought of as essentially calling the [`TTree.arrays()`](https://uproot.readthedocs.io/en/latest/uproot.behaviors.TTree.TTree.html#arrays) function of `uproot` with the possibility to specify several of the arguments, and returning the result. (Tree here should be understood to also include `RNtuple` objects; nothing needs to be changed to read those.) Multiple queries can be bundled into one request. It is also possible to copy objects from the input file.
 
 Let's look at the structure of an Uproot-Raw query.
 
@@ -56,7 +56,7 @@ Each dictionary either has a `treename` key (indicating that it is a query on a
 
 * **Query dictionaries**: these dictionaries contain a `treename` key, which specifies the tree(s) to be queried. The keys are:
   * `treename`: either a string, a list of strings, or a dictionary. This selects the names of the tree(s) to which the query will be applied. In the case that a dictionary is passed, the keys will be used to choose the input trees, and the values will be used as the name of the tree that results from the query - this allows the user to run multiple queries on the same tree, saving the output to a different tree each time.
-  * `expressions`, `cut`, `filter_name`, `aliases`: have the same meaning as for [`TTree.arrays()`](https://uproot.readthedocs.io/en/latest/uproot.behaviors.TTree.TTree.html#arrays) in `uproot`, except that functions aren't permitted (but *glob*s and _regular expressions_, which are special kinds of strings, are).
+  * `expressions`, `cut`, `filter_name`, `aliases`: have the same meaning as for [`TTree.arrays()`](https://uproot.readthedocs.io/en/latest/uproot.behaviors.TTree.TTree.html#arrays)/[`RNTuple.arrays()`](https://uproot.readthedocs.io/en/stable/uproot.models.RNTuple.Model_ROOT_3a3a_RNTuple.html#arrays) in `uproot`, except that functions aren't permitted (but *glob*s and _regular expressions_, which are special kinds of strings, are).
 
   Other keys will be ignored.
 
diff --git a/servicex/databinder_models.py b/servicex/databinder_models.py
@@ -198,6 +198,11 @@ class OutputFormatEnum(str, Enum):
         a ROOT TTree https://root.cern.ch/doc/master/classTTree.html
         """
 
+        root_rntuple = "root-rntuple"
+        """
+        Save the output as an RNtuple https://root.cern/doc/master/classROOT_1_1RNTuple.html
+        """
+
         def to_ResultFormat(self) -> ResultFormat:
             """This method is used to convert the OutputFormatEnum enum to the ResultFormat enum,
             which is what is actually used for the TransformRequest. This allows us to use
@@ -207,6 +212,8 @@ def to_ResultFormat(self) -> ResultFormat:
                 return ResultFormat.parquet
             elif self == self.root_ttree:
                 return ResultFormat.root_ttree
+            elif self == self.root_rntuple:
+                return ResultFormat.root_rntuple
             else:  # pragma: no cover
                 raise RuntimeError(f"Bad OutputFormatEnum {self}")
 
diff --git a/servicex/models.py b/servicex/models.py
@@ -79,6 +79,7 @@ class ResultFormat(str, Enum):
 
     parquet = "parquet"
     root_ttree = "root-file"
+    root_rntuple = "root-rntuple"
 
 
 class Status(str, Enum):
diff --git a/tests/test_databinder.py b/tests/test_databinder.py
@@ -97,10 +97,14 @@ def test_output_format():
     spec = basic_spec()
     spec["General"] = {"OutputFormat": "root-ttree"}
     ServiceXSpec.model_validate(spec)
+    spec["General"] = {"OutputFormat": "root-rntuple"}
+    ServiceXSpec.model_validate(spec)
     spec["General"] = {"OutputFormat": "parquet"}
     ServiceXSpec.model_validate(spec)
     spec["General"] = {"OutputFormat": OutputFormat.root_ttree}
     ServiceXSpec.model_validate(spec)
+    spec["General"] = {"OutputFormat": OutputFormat.root_rntuple}
+    ServiceXSpec.model_validate(spec)
     spec["General"] = {"OutputFormat": OutputFormat.parquet}
     ServiceXSpec.model_validate(spec)
     with pytest.raises(ValidationError):
@@ -846,6 +850,35 @@ def test_uproot_raw_query_parquet(transformed_result, codegen_list, with_event_l
         deliver(spec, config_path="tests/example_config.yaml")
 
 
+def test_uproot_raw_query_rntuple(transformed_result, codegen_list, with_event_loop):
+    from servicex import deliver
+    from servicex.query import UprootRaw  # type: ignore
+
+    spec = ServiceXSpec.model_validate(
+        {
+            "General": {"OutputFormat": "root-rntuple"},
+            "Sample": [
+                {
+                    "Name": "sampleA",
+                    "RucioDID": "user.ivukotic:user.ivukotic.single_top_tW__nominal",
+                    "Query": UprootRaw([{"treename": "nominal"}]),
+                }
+            ],
+        }
+    )
+    with (
+        patch(
+            "servicex.dataset_group.DatasetGroup.as_files",
+            return_value=[transformed_result],
+        ),
+        patch(
+            "servicex.servicex_client.ServiceXClient.get_code_generators",
+            return_value=codegen_list,
+        ),
+    ):
+        deliver(spec, config_path="tests/example_config.yaml")
+
+
 def test_generic_query(codegen_list):
     from servicex.servicex_client import ServiceXClient
 
