Initial attempt at RNTuple support in framework and uproot-raw transformer

Author: ponyisi

Diff for: code_generator_raw_uproot/servicex/raw_uproot_code_generator/request_translator.py

@@ -153,7 +153,7 @@ def run_single_query(file_path, query):
                     if 'cut' in sanitized_args:
                         sanitized_args.pop('cut')
                     arr = t.arrays(language=lang, entry_stop=0, **sanitized_args)
-                    rv_arrays_trees[outtreename] = (None, {{_: arr[_].type for _ in arr.fields}})
+                    rv_arrays_trees[outtreename] = (None, arr.layout)
         else:
             histograms = query['copy_histograms']
             keys = fl.keys(filter_name=histograms, cycle=False)

Diff for: code_generator_raw_uproot/servicex/templates/transform_single_file.py

@@ -26,20 +26,34 @@ def transform_single_file(file_path: str, output_path: Path, output_format: str)
 
         ttime = time.time()
 
-        if output_format == 'root-file':
+        if output_format in ('root-file', 'root-rntuple'):
             import uproot
             etime = time.time()
             # opening the file with open() is a workaround for a bug handling multiple colons
             # in the filename in uproot 5.3.9
             with open(output_path, 'b+w') as wfile:
                 with uproot.recreate(wfile) as writer:
                     for k, v in awkward_array_dict.items():
-                        if v[0] is not None:
-                            writer[k] = {field: v[0][field] for field in
-                                         v[0].fields} if v[0].fields \
-                                else v[0]
-                        else:
-                            writer.mktree(k, v[1])
+                        if output_format == 'root-file':
+                            if v[0] is not None:
+                                writer[k] = {field: v[0][field] for field in
+                                             v[0].fields}
+                            else:
+                                writer.mktree(k, dict(zip(v[1].form.columns(),
+                                                          v[1].form.column_types())))
+                        else:  # RNTuple
+                            if v[0] is not None:
+                                # Work around a limitation in uproot 5.6.0
+                                # If a cut is specified, we'll get ListArrays which can't be
+                                # written via uproot. Convert them to ListOffsetArrays
+                                # Assume the ListArrays are only at top level
+                                warr = ak.zip({_: v[0][_].layout.to_ListOffsetArray64()
+                                               if isinstance(v[0][_].layout, ak.contents.ListArray)
+                                               else v[0][_]
+                                               for _ in v[0].fields}, depth_limit=1)
+                                writer.mkrntuple(k, warr)
+                            else:
+                                writer.mkrntuple(k, v[1].form)
                     for k, v in histograms.items():
                         writer[k] = v
             wtime = time.time()

Diff for: code_generator_raw_uproot/tests/test_src.py

@@ -50,7 +50,7 @@ def test_generate_code():
                              'filter_name': ['lbn']},
                             {'copy_histograms': 'CutBookkeeper*'}
                             ])
-        expected_hash = "6f8ac79962ef753d3e7fd6161cba6fe8"
+        expected_hash = "9d7f01cdc875549daf3eed81a502147a"
         result = translator.generate_code(query, tmpdirname)
 
         # is the generated code at least syntactically valid Python?

Diff for: code_generator_raw_uproot/transformer_capabilities.json

@@ -2,7 +2,7 @@
   "name": "Uproot transformer using native uproot arguments",
   "description": "Extracts data from flat ntuple style root files.",
   "limitations": "Would be good to note what isn't implemented",
-  "file-formats": ["parquet"],
+  "file-formats": ["parquet", "root-file", "root-rntuple"],
   "stats-parser": "RawUprootStats",
   "language": "python",
   "command": "/generated/transform_single_file.py"

Diff for: servicex_app/servicex_app/resources/transformation/submit.py

@@ -76,7 +76,8 @@ def make_api(cls, rabbitmq_adaptor, object_store,
             TransformRequest.VOLUME_DEST
         ])
         cls.parser.add_argument(
-            'result-format', choices=['arrow', 'parquet', 'root-file'], default='arrow'
+            'result-format', choices=['parquet', 'root-file',
+                                      'root-rntuple'], default='parquet'
         )
         return cls
 

Diff for: transformer_sidecar/src/transformer_sidecar/transformer_argument_parser.py

@@ -56,8 +56,8 @@ def __init__(self, description="ServiceX Transformer"):
                           default=None, help='Local directory to output results')
 
         self.add_argument('--result-format', dest='result_format', action='store',
-                          default='arrow', help='arrow, parquet, root-file',
-                          choices=['arrow', 'parquet', 'root-file'])
+                          default='parquet', help='parquet, root-file, root-rntuple',
+                          choices=['parquet', 'root-file', 'root-rntuple'])
 
         self.add_argument('--rabbit-uri', dest="rabbit_uri", action='store',
                           default='host.docker.internal')

Diff for: transformer_sidecar/tests/test_transformer_argument_parser.py

@@ -44,7 +44,7 @@ def test_parse(self, mocker):
                     "--path", "/foo/bar",
                     "--limit", "10",
                     '--result-destination', 'object-store',
-                    '--result-format', 'arrow',
+                    '--result-format', 'parquet',
                     '--rabbit-uri', "http://rabbit.org",
                     '--request-id', "123-45-678"
                     ]
@@ -53,7 +53,7 @@ def test_parse(self, mocker):
         assert args.path == '/foo/bar'
         assert args.limit == 10
         assert args.result_destination == 'object-store'
-        assert args.result_format == 'arrow'
+        assert args.result_format == 'parquet'
         assert args.rabbit_uri == "http://rabbit.org"
         assert args.request_id == "123-45-678"