Add counters for xarray_beam rechunk and combiners

PiperOrigin-RevId: 820964883

Author: shoyer

xarray_beam/_src/combiners.py

@@ -76,20 +76,28 @@ def create_accumulator(self):
     return (0, 0)
 
   def add_input(self, sum_count, element):
+    core.inc_counter(self.__class__, 'add-input-calls')
     (sum_, count) = sum_count
     if self.sum_and_count is not None:
+      core.inc_counter(self.__class__, 'add-input-in-bytes', element.nbytes)
       sum_increment, count_increment = self.sum_and_count(element)
     else:
       sum_increment, count_increment = element
+      nbytes = sum_increment.nbytes + count_increment.nbytes
+      core.inc_counter(self.__class__, 'add-input-bytes', nbytes)
     new_sum = sum_ + sum_increment
     new_count = count + count_increment
+    nbytes = new_sum.nbytes + new_count.nbytes
+    core.inc_counter(self.__class__, 'add-input-out-bytes', nbytes)
     return new_sum, new_count
 
   def merge_accumulators(self, accumulators):
+    core.inc_counter(self.__class__, 'merge-accumulators')
     sums, counts = zip(*accumulators)
     return sum(sums), sum(counts)
 
   def extract_output(self, sum_count):
+    core.inc_counter(self.__class__, 'extract-outputs')
     if self.finalize:
       (sum_, count) = sum_count
       return sum_ / count
@@ -317,12 +325,29 @@ def __post_init__(self):
         f' pre_aggregate={self.pre_aggregate}'
     )
 
+  @property
+  def _sum_and_count(self):
+    return _SumAndCount(self.dims, self.skipna, self.dtype)
+
+  def _pre_aggregate(
+      self, key: core.Key, chunk: xarray.Dataset
+  ) -> tuple[core.Key, tuple[xarray.Dataset, xarray.Dataset]]:
+    core.inc_counter(self.__class__, 'preaggregate-calls')
+    core.inc_counter(self.__class__, 'preaggregate-in-bytes', chunk.nbytes)
+    sum_increment, count_increment = self._sum_and_count(chunk)
+    out_bytes = sum_increment.nbytes + count_increment.nbytes
+    core.inc_counter(self.__class__, 'preaggregate-out-bytes', out_bytes)
+    return key, (sum_increment, count_increment)
+
   def _finalize_no_combiner(
       self, key: core.Key, sum_count: tuple[xarray.Dataset, xarray.Dataset]
   ) -> tuple[core.Key, xarray.Dataset]:
     key = key.with_offsets(**{d: None for d in self.dims if d in key.offsets})
     sum_, count = sum_count
-    return key, sum_ / count
+    chunk = sum_ / count
+    core.inc_counter(self.__class__, 'finalize-calls')
+    core.inc_counter(self.__class__, 'finalize-bytes', chunk.nbytes)
+    return key, chunk
 
   def _prepare_key(
       self, key: core.Key, chunk: xarray.Dataset
@@ -348,24 +373,20 @@ def _strip_fanout_bins(
     return key, value
 
   def expand(self, pcoll):
-    sum_and_count = _SumAndCount(self.dims, self.skipna, self.dtype)
-
     if not self.bins_per_stage:  # no combiner needed
-      pcoll |= 'Aggregate' >> beam.MapTuple(lambda k, v: (k, sum_and_count(v)))
+      pcoll |= 'Aggregate' >> beam.MapTuple(self._pre_aggregate)
       pcoll |= 'Finalize' >> beam.MapTuple(self._finalize_no_combiner)
       return pcoll
 
     if self.pre_aggregate:
-      pcoll |= 'PreAggregate' >> beam.MapTuple(
-          lambda k, v: (k, sum_and_count(v))
-      )
+      pcoll |= 'PreAggregate' >> beam.MapTuple(self._pre_aggregate)
     pcoll |= 'PrepareKey' >> beam.MapTuple(self._prepare_key)
     for i in range(len(self.bins_per_stage)):
       final_stage = i + 1 >= len(self.bins_per_stage)
       if self.pre_aggregate or i > 0:
         combine_fn = MeanCombineFn(None, finalize=final_stage)
       else:
-        combine_fn = MeanCombineFn(sum_and_count, finalize=final_stage)
+        combine_fn = MeanCombineFn(self._sum_and_count, finalize=final_stage)
       pcoll |= f'Combine{i}' >> beam.CombinePerKey(combine_fn)
       if not final_stage:
         pcoll |= f'StripBin{i}' >> beam.MapTuple(self._strip_leading_fanout_bin)

xarray_beam/_src/core.py

@@ -507,8 +507,7 @@ def _shard_inputs(self) -> list[tuple[int | None, str | None]]:
 
   def _key_to_chunks(self, key: Key) -> Iterator[tuple[Key, DatasetOrDatasets]]:
     """Convert a Key into an in-memory (Key, xarray.Dataset) pair."""
-    namespace = "xarray_beam.DatasetToChunks"
-    with inc_timer_msec(namespace, "read-msec"):
+    with inc_timer_msec(self.__class__, "read-msec"):
       sizes = {
           dim: self.expanded_chunks[dim][self.offset_index[dim][offset]]
           for dim, offset in key.offsets.items()
@@ -524,9 +523,9 @@ def _key_to_chunks(self, key: Key) -> Iterator[tuple[Key, DatasetOrDatasets]]:
         result = chunk.chunk().compute(num_workers=num_threads)
         results.append(result)
 
-    inc_counter(namespace, "read-chunks")
+    inc_counter(self.__class__, "read-chunks")
     inc_counter(
-        namespace, "read-bytes", sum(result.nbytes for result in results)
+        self.__class__, "read-bytes", sum(result.nbytes for result in results)
     )
 
     if isinstance(self.dataset, xarray.Dataset):

xarray_beam/_src/rechunk.py

@@ -314,12 +314,17 @@ class ConsolidateChunks(beam.PTransform):
   target_chunks: Mapping[str, int]
 
   def _prepend_chunk_key(self, key, chunk):
+    core.inc_counter(self.__class__, 'in-chunks')
+    core.inc_counter(self.__class__, 'in-bytes', chunk.nbytes)
     rounded_key = _round_chunk_key(key, self.target_chunks)
     return rounded_key, (key, chunk)
 
   def _consolidate(self, key, inputs):
-    ((consolidated_key, dataset),) = consolidate_chunks(inputs)
+    with core.inc_timer_msec(self.__class__, 'consolidate-msec'):
+      ((consolidated_key, dataset),) = consolidate_chunks(inputs)
     assert key == consolidated_key, (key, consolidated_key)
+    core.inc_counter(self.__class__, 'out-chunks')
+    core.inc_counter(self.__class__, 'out-bytes', dataset.nbytes)
     return consolidated_key, dataset
 
   def expand(self, pcoll):
@@ -339,17 +344,22 @@ class ConsolidateVariables(beam.PTransform):
   # of variables.
 
   def _prepend_chunk_key(self, key, chunk):
+    core.inc_counter(self.__class__, 'in-chunks')
+    core.inc_counter(self.__class__, 'in-bytes', chunk.nbytes)
     return key.replace(vars=None), (key, chunk)
 
   def _consolidate(self, key, inputs):
-    ((consolidated_key, dataset),) = consolidate_variables(inputs)
+    with core.inc_timer_msec(self.__class__, 'consolidate-msec'):
+      ((consolidated_key, dataset),) = consolidate_variables(inputs)
     assert key.offsets == consolidated_key.offsets, (key, consolidated_key)
     assert key.vars is None
     # TODO(shoyer): consider carefully whether it is better to return key or
     # consolidated_key. They are both valid in the xarray-beam data model -- the
     # difference is whether vars=None or is an explicit set of variables.
     # For now, conservatively return the version of key with vars=None so
     # users don't rely on it.
+    core.inc_counter(self.__class__, 'out-chunks')
+    core.inc_counter(self.__class__, 'out-bytes', dataset.nbytes)
     return key, dataset
 
   def expand(self, pcoll):
@@ -432,7 +442,13 @@ def _split_chunks(
     target_chunks = {
         k: v for k, v in self.target_chunks.items() if k in dataset.dims
     }
-    yield from split_chunks(key, dataset, target_chunks)
+    core.inc_counter(self.__class__, 'in-chunks')
+    core.inc_counter(self.__class__, 'in-bytes', dataset.nbytes)
+    with core.inc_timer_msec(self.__class__, 'split-msec'):
+      for new_key, new_dataset in split_chunks(key, dataset, target_chunks):
+        yield new_key, new_dataset
+        core.inc_counter(self.__class__, 'out-chunks')
+        core.inc_counter(self.__class__, 'out-bytes', new_dataset.nbytes)
 
   def expand(self, pcoll):
     return pcoll | beam.FlatMapTuple(self._split_chunks)
@@ -458,8 +474,19 @@ def split_variables(
 class SplitVariables(beam.PTransform):
   """Split existing chunks into a separate chunk per data variable."""
 
+  def _split_variables(
+      self, key: core.Key, dataset: xarray.Dataset
+  ) -> Iterator[tuple[core.Key, xarray.Dataset]]:
+    core.inc_counter(self.__class__, 'in-chunks')
+    core.inc_counter(self.__class__, 'in-bytes', dataset.nbytes)
+    with core.inc_timer_msec(self.__class__, 'split-msec'):
+      for new_key, new_dataset in split_variables(key, dataset):
+        yield new_key, new_dataset
+        core.inc_counter(self.__class__, 'out-chunks')
+        core.inc_counter(self.__class__, 'out-bytes', new_dataset.nbytes)
+
   def expand(self, pcoll):
-    return pcoll | beam.FlatMapTuple(split_variables)
+    return pcoll | beam.FlatMapTuple(self._split_variables)
 
 
 @core.export

xarray_beam/_src/zarr.py

@@ -777,11 +777,10 @@ def _validate_zarr_chunk(self, key, chunk, template=None):
 
   def _write_chunk_to_zarr(self, key, chunk, template=None):
     assert template is not None
-    namespace = 'xarray_beam.ChunksToZarr'
-    with core.inc_timer_msec(namespace, "write-msec"):
+    with core.inc_timer_msec(self.__class__, "write-msec"):
       write_chunk_to_zarr(key, chunk, self.store, template)
-    core.inc_counter(namespace, 'write-chunks')
-    core.inc_counter(namespace, 'write-bytes', chunk.nbytes)
+    core.inc_counter(self.__class__, 'write-chunks')
+    core.inc_counter(self.__class__, 'write-bytes', chunk.nbytes)
 
   def expand(self, pcoll):
     if isinstance(self.template, xarray.Dataset):