Implemented iterate method

Author: ariostas

src/uproot/models/RNTuple.py

@@ -133,6 +133,55 @@ def _arrays(
     )[entry_start:entry_stop]
 
 
+def _num_entries_for(in_ntuple, target_num_bytes, filter_name):
+    # TODO: part of this is also done in _arrays, so we should refactor this
+    # TODO: there might be a better way to estimate the number of entries
+    entry_stop = in_ntuple.ntuple.num_entries
+
+    clusters = in_ntuple.ntuple.cluster_summaries
+    cluster_starts = numpy.array([c.num_first_entry for c in clusters])
+
+    start_cluster_idx = numpy.searchsorted(cluster_starts, 0, side="right") - 1
+    stop_cluster_idx = numpy.searchsorted(cluster_starts, entry_stop, side="right")
+
+    form = in_ntuple.to_akform().select_columns(
+        filter_name, prune_unions_and_records=False
+    )
+    target_cols = []
+    _recursive_find(form, target_cols)
+
+    total_bytes = 0
+    for key in target_cols:
+        if "column" in key and "union" not in key:
+            key_nr = int(key.split("-")[1])
+            for cluster in range(start_cluster_idx, stop_cluster_idx):
+                pages = in_ntuple.ntuple.page_list_envelopes.pagelinklist[cluster][
+                    key_nr
+                ]
+                total_bytes += sum(page.locator.num_bytes for page in pages)
+
+    total_entries = entry_stop
+    if total_bytes == 0:
+        num_entries = 0
+    else:
+        num_entries = int(round(target_num_bytes * total_entries / total_bytes))
+    if num_entries <= 0:
+        return 1
+    else:
+        return num_entries
+
+
+def _regularize_step_size(in_ntuple, step_size, filter_name):
+    if uproot._util.isint(step_size):
+        return step_size
+    target_num_bytes = uproot._util.memory_size(
+        step_size,
+        "number of entries or memory size string with units "
+        f"(such as '100 MB') required, not {step_size!r}",
+    )
+    return _num_entries_for(in_ntuple, target_num_bytes, filter_name)
+
+
 class Model_ROOT_3a3a_Experimental_3a3a_RNTuple(uproot.model.Model):
     """
     A versionless :doc:`uproot.model.Model` for ``ROOT::Experimental::RNTuple``.
@@ -742,6 +791,13 @@ def arrays(
             array_cache=array_cache,
         )
 
+    def iterate(self, filter_name="*", *args, step_size="100 MB", **kwargs):
+        step_size = _regularize_step_size(self, step_size, filter_name)
+        for start in range(0, self.num_entries, step_size):
+            yield self.arrays(
+                *args, entry_start=start, entry_stop=start + step_size, **kwargs
+            )
+
 
 # Supporting function and classes
 def _split_switch_bits(content):
@@ -1215,6 +1271,13 @@ def __array__(self, *args, **kwargs):
         else:
             return numpy.array(out, *args, **kwargs)
 
+    def iterate(self, filter_name="*", *args, step_size="100 MB", **kwargs):
+        step_size = _regularize_step_size(self, step_size, filter_name)
+        for start in range(0, self.ntuple.num_entries, step_size):
+            yield self.array(
+                *args, entry_start=start, entry_stop=start + step_size, **kwargs
+            )
+
 
 uproot.classes["ROOT::Experimental::RNTuple"] = (
     Model_ROOT_3a3a_Experimental_3a3a_RNTuple

tests/test_1250_rntuple_improvements.py

@@ -37,3 +37,49 @@ def test_array_methods():
         obj = f["RNT:CollectionTree"]
         jets = obj["AntiKt4TruthDressedWZJetsAux:"].arrays()
         assert len(jets[0].pt) == 5
+
+
+def test_iterate():
+    filename = skhep_testdata.data_path(
+        "Run2012BC_DoubleMuParked_Muons_rntuple_1000evts.root"
+    )
+    with uproot.open(filename) as f:
+        obj = f["Events"]
+        for i, arrays in enumerate(obj.iterate(step_size=100)):
+            assert len(arrays) == 100
+            if i == 0:
+                expected_pt = [10.763696670532227, 15.736522674560547]
+                expected_charge = [-1, -1]
+                assert arrays["Muon_pt"][0].tolist() == expected_pt
+                assert arrays["Muon_charge"][0].tolist() == expected_charge
+
+        for i, arrays in enumerate(obj.iterate(step_size="10 kB")):
+            if i == 0:
+                assert len(arrays) == 363
+                expected_pt = [10.763696670532227, 15.736522674560547]
+                expected_charge = [-1, -1]
+                assert arrays["Muon_pt"][0].tolist() == expected_pt
+                assert arrays["Muon_charge"][0].tolist() == expected_charge
+            elif i == 1:
+                assert len(arrays) == 363
+            elif i == 2:
+                assert len(arrays) == 274
+            else:
+                assert False
+
+        Muon_pt = obj["Muon_pt"]
+        for i, arrays in enumerate(Muon_pt.iterate(step_size=100)):
+            assert len(arrays) == 100
+            if i == 0:
+                expected_pt = [10.763696670532227, 15.736522674560547]
+                assert arrays[0].tolist() == expected_pt
+
+        for i, arrays in enumerate(Muon_pt.iterate(step_size="5 kB")):
+            if i == 0:
+                assert len(arrays) == 611
+                expected_pt = [10.763696670532227, 15.736522674560547]
+                assert arrays[0].tolist() == expected_pt
+            elif i == 1:
+                assert len(arrays) == 389
+            else:
+                assert False