feat: Record source performance counters in uproot.dask report (#1156)

* Record source performance counters in uproot.dask report

* style: pre-commit fixes

* Fix lint

* style: pre-commit fixes

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: nsmith-

src/uproot/_dask.py

@@ -6,6 +6,8 @@
 import time
 from collections.abc import Callable, Iterable, Mapping
 
+from uproot.source.chunk import SourcePerformanceCounters
+
 try:
     from typing import TYPE_CHECKING, Final
 
@@ -972,7 +974,9 @@ def allowed_exceptions(self):
     def return_report(self) -> bool:
         return bool(self.allow_read_errors_with_report)
 
-    def read_tree(self, tree: HasBranches, start: int, stop: int) -> AwkArray:
+    def read_tree(
+        self, tree: HasBranches, start: int, stop: int
+    ) -> tuple[AwkArray, SourcePerformanceCounters]:
         assert start <= stop
 
         from awkward._nplikes.numpy import Numpy
@@ -1017,13 +1021,15 @@ def read_tree(self, tree: HasBranches, start: int, stop: int) -> AwkArray:
                     dtype=dtype,
                 )
 
-        return awkward.from_buffers(
+        out = awkward.from_buffers(
             self.expected_form,
             stop - start,
             container,
             behavior=self.form_mapping_info.behavior,
             buffer_key=self.form_mapping_info.buffer_key,
         )
+        assert tree.source  # we must be reading something here
+        return out, tree.source.performance_counters
 
     def mock(self) -> AwkArray:
         awkward = uproot.extras.awkward()
@@ -1114,6 +1120,7 @@ def _report_failure(exception, call_time, *args, **kwargs):
             {
                 "call_time": call_time,
                 "duration": None,
+                "performance_counters": None,
                 "args": [repr(a) for a in args],
                 "kwargs": [[k, repr(v)] for k, v in kwargs.items()],
                 "exception": type(exception).__name__,
@@ -1127,11 +1134,13 @@ def _report_failure(exception, call_time, *args, **kwargs):
 
 def _report_success(duration, *args, **kwargs):
     awkward = uproot.extras.awkward()
+    counters = kwargs.pop("counters")
     return awkward.Array(
         [
             {
                 "call_time": None,
                 "duration": duration,
+                "performance_counters": counters.asdict(),
                 "args": [repr(a) for a in args],
                 "kwargs": [[k, repr(v)] for k, v in kwargs.items()],
                 "exception": None,
@@ -1195,14 +1204,17 @@ def __call__(self, i_start_stop):
         if self.return_report:
             call_time = time.time_ns()
             try:
-                result, duration = with_duration(self._call_impl)(i, start, stop)
+                (result, counters), duration = with_duration(self._call_impl)(
+                    i, start, stop
+                )
                 return (
                     result,
                     _report_success(
                         duration,
                         self.ttrees[i],
                         start,
                         stop,
+                        counters=counters,
                     ),
                 )
             except self.allowed_exceptions as err:
@@ -1217,7 +1229,8 @@ def __call__(self, i_start_stop):
                     ),
                 )
 
-        return self._call_impl(i, start, stop)
+        result, _ = self._call_impl(i, start, stop)
+        return result
 
     def _call_impl(self, i, start, stop):
         return self.read_tree(
@@ -1305,7 +1318,7 @@ def __call__(self, blockwise_args):
         if self.return_report:
             call_time = time.time_ns()
             try:
-                result, duration = with_duration(self._call_impl)(
+                (result, counters), duration = with_duration(self._call_impl)(
                     file_path, object_path, i_step_or_start, n_steps_or_stop, is_chunk
                 )
                 return (
@@ -1317,6 +1330,7 @@ def __call__(self, blockwise_args):
                         i_step_or_start,
                         n_steps_or_stop,
                         is_chunk,
+                        counters=counters,
                     ),
                 )
             except self.allowed_exceptions as err:
@@ -1333,9 +1347,10 @@ def __call__(self, blockwise_args):
                     ),
                 )
 
-        return self._call_impl(
+        result, _ = self._call_impl(
             file_path, object_path, i_step_or_start, n_steps_or_stop, is_chunk
         )
+        return result
 
     def project_keys(self: T, keys: frozenset[str]) -> T:
         return _UprootOpenAndRead(

src/uproot/behaviors/TBranch.py

@@ -23,6 +23,7 @@
 import uproot
 import uproot.interpretation.grouped
 import uproot.language.python
+import uproot.source.chunk
 from uproot._util import no_filter
 
 np_uint8 = numpy.dtype("u1")
@@ -1664,6 +1665,18 @@ def __iter__(self):
     def __len__(self):
         return len(self.branches)
 
+    @property
+    def source(self) -> uproot.source.chunk.Source | None:
+        """Returns the associated source of data for this container, if it exists
+
+        Returns: uproot.source.chunk.Source or None
+        """
+        if isinstance(self, uproot.model.Model) and isinstance(
+            self._file, uproot.reading.ReadOnlyFile
+        ):
+            return self._file.source
+        return None
+
 
 _branch_clean_name = re.compile(r"(.*\.)*([^\.\[\]]*)(\[.*\])*")
 _branch_clean_parent_name = re.compile(r"(.*\.)*([^\.\[\]]*)\.([^\.\[\]]*)(\[.*\])*")

src/uproot/source/chunk.py

@@ -12,6 +12,7 @@
 
 from __future__ import annotations
 
+import dataclasses
 import numbers
 import queue
 
@@ -41,6 +42,18 @@ def file_path(self) -> str:
         return self._file_path
 
 
+@dataclasses.dataclass
+class SourcePerformanceCounters:
+    """Container for performance counters"""
+
+    num_requested_bytes: int
+    num_requests: int
+    num_requested_chunks: int
+
+    def asdict(self) -> dict[str, int]:
+        return dataclasses.asdict(self)
+
+
 class Source:
     """
     Abstract class for physically reading and writing data from a file, which
@@ -138,6 +151,14 @@ def num_requested_bytes(self) -> int:
         """
         return self._num_requested_bytes
 
+    @property
+    def performance_counters(self) -> SourcePerformanceCounters:
+        return SourcePerformanceCounters(
+            self._num_requested_bytes,
+            self._num_requests,
+            self._num_requested_chunks,
+        )
+
     def close(self):
         """
         Manually closes the file(s) and stops any running threads.