diff --git a/km3io/rootio.py b/km3io/rootio.py index 17442feb9d13bdd47fe92595b5d8957764011d15..e49fd6324ca17a703a45b1bb3886b4d68bfbadcd 100644 --- a/km3io/rootio.py +++ b/km3io/rootio.py @@ -279,7 +279,24 @@ class Branch: def __getattr__(self, attr): if attr not in self._aliases: raise AttributeError(f"No field named {attr}. Available fields: {self.fields}") - return unfold_indices(self._branch[self._aliases[attr]].array(), self._index_chain) + key = self._aliases[attr] + + if self._index_chain: + idx0 = self._index_chain[0] + if isinstance(idx0, (int, np.int32, np.int64)): + # optimise single-element and slice lookups + start = idx0 + stop = idx0 + 1 + arr = ak.flatten(self._branch[key].array(entry_start=start, entry_stop=stop)) + return unfold_indices(arr, self._index_chain[1:]) + if isinstance(idx0, slice): + if idx0.step is None or idx0.step == 1: + start = idx0.start + stop = idx0.stop + arr = self._branch[key].array(entry_start=start, entry_stop=stop) + return unfold_indices(arr, self._index_chain[1:]) + + return unfold_indices(self._branch[key].array(), self._index_chain) def __getitem__(self, key): return self.__class__(self._branch, self.fields, self._aliases, self._index_chain + [key])