Tweak xbeam.Dataset.__repr__()

PiperOrigin-RevId: 813333527

Author: shoyer

xarray_beam/_src/dataset.py

@@ -33,9 +33,11 @@
 import dataclasses
 import functools
 import itertools
+import math
 import operator
 import os.path
 import tempfile
+import textwrap
 from typing import Any, Callable, Literal
 
 import apache_beam as beam
@@ -45,6 +47,25 @@
 from xarray_beam._src import zarr
 
 
+def _at_least_two_digits(n: int | float) -> str:
+  if isinstance(n, int):
+    return str(n)
+  elif round(n, 2) < 10:
+    return f'{n:.1f}'
+  else:
+    return f'{n:.0f}'
+
+
+def _to_human_size(nbytes: int) -> str:
+  """Convert a number of bytes to a human-readable string."""
+  for unit in ['B', 'kB', 'MB', 'GB', 'TB', 'PB', 'EB']:
+    if nbytes < 1000:
+      return f'{_at_least_two_digits(nbytes)}{unit}'
+    nbytes /= 1000
+  nbytes *= 1000
+  return f'{_at_least_two_digits(nbytes)}EB'
+
+
 def _infer_new_chunks(
     old_sizes: Mapping[str, int],
     old_chunks: Mapping[str, int],
@@ -149,6 +170,50 @@ class Dataset:
   def __post_init__(self):
     self.chunks = rechunk.normalize_chunks(self.chunks, self.sizes)
 
+  @property
+  def bytes_per_chunk(self) -> int:
+    """Estimate of the number of bytes per chunk."""
+    variable_sizes = [
+        v.dtype.itemsize * math.prod(self.chunks[d] for d in v.dims)
+        for v in self.template.values()
+    ]
+    return max(variable_sizes) if self.split_vars else sum(variable_sizes)
+
+  @property
+  def chunk_count(self) -> int:
+    """Count the number of chunks in this dataset."""
+    if self.split_vars:
+      total = 0
+      for variable in self.template.values():
+        total += math.prod(
+            math.ceil(self.sizes[d] / self.chunks[d])
+            for d in variable.dims
+        )
+      return total
+    else:
+      return math.prod(
+          math.ceil(self.sizes[d] / self.chunks[d])
+          for d in self.sizes
+      )
+
+  def __repr__(self):
+    base = repr(self.template)
+    chunks_str = ', '.join(
+        [f'{k}: {v}' for k, v in self.chunks.items()]
+        + [f'split_vars={self.split_vars}']
+    )
+    chunk_size = _to_human_size(self.bytes_per_chunk)
+    total_size = _to_human_size(self.template.nbytes)
+    chunk_count = self.chunk_count
+    plural = 's' if chunk_count != 1 else ''
+    return (
+        f'<xarray_beam.Dataset>\n'
+        f'PTransform: {self.ptransform}\n'
+        f'Chunks:     {chunk_size} ({chunks_str})\n'
+        f'Template:   {total_size} ({chunk_count} chunk{plural})\n'
+        + textwrap.indent('\n'.join(base.split('\n')[1:]), ' ' * 4)
+    )
+
   @classmethod
   def from_xarray(
       cls,
@@ -351,13 +416,5 @@ def head(self, **indexers_kwargs: int) -> Dataset:
   transpose = _whole_dataset_method('transpose')
 
   def pipe(self, func, *args, **kwargs):
+    """Apply a function to this dataset, like xarray.Dataset.pipe()."""
     return func(*args, **kwargs)
-
-  def __repr__(self):
-    base = repr(self.template)
-    chunks_str = ', '.join(f'{k}: {v}' for k, v in self.chunks.items())
-    return (
-        f'<xarray_beam.Dataset[{chunks_str}][split_vars={self.split_vars}]>'
-        + f'\nPTransform: {self.ptransform}\n'
-        + '\n'.join(base.split('\n')[1:])
-    )

xarray_beam/_src/dataset_test.py

@@ -23,8 +23,47 @@
 from xarray_beam._src import test_util
 
 
+class ToHumanSizeTest(test_util.TestCase):
+
+  @parameterized.named_parameters(
+      dict(testcase_name='zero', size=0, expected='0B'),
+      dict(testcase_name='one_byte', size=1, expected='1B'),
+      dict(testcase_name='nine_bytes', size=9, expected='9B'),
+      dict(testcase_name='ten_bytes', size=10, expected='10B'),
+      dict(testcase_name='ninety_nine_bytes', size=99, expected='99B'),
+      dict(testcase_name='one_hundred_bytes', size=100, expected='100B'),
+      dict(testcase_name='almost_one_kb', size=999, expected='999B'),
+      dict(testcase_name='one_kb', size=1000, expected='1.0kB'),
+      dict(testcase_name='round_to_10_kb', size=9996, expected='10kB'),
+      dict(testcase_name='100_mb', size=10**8, expected='100MB'),
+      dict(testcase_name='one_mb', size=10**6, expected='1.0MB'),
+      dict(testcase_name='one_gb', size=10**9, expected='1.0GB'),
+      dict(testcase_name='one_tb', size=10**12, expected='1.0TB'),
+      dict(testcase_name='one_pb', size=10**15, expected='1.0PB'),
+      dict(testcase_name='one_eb', size=10**18, expected='1.0EB'),
+      dict(testcase_name='one_thousand_eb', size=10**21, expected='1000EB'),
+      dict(testcase_name='ten_thousand_eb', size=10**22, expected='10000EB'),
+  )
+  def test_to_human_size(self, size, expected):
+    self.assertEqual(xbeam_dataset._to_human_size(size), expected)
+
+
 class DatasetTest(test_util.TestCase):
 
+  def test_repr(self):
+    ds = xarray.Dataset({'foo': ('x', np.arange(10))})
+    beam_ds = xbeam.Dataset.from_xarray(ds, {'x': 5})
+    self.assertRegex(
+        repr(beam_ds),
+        re.escape(
+            '<xarray_beam.Dataset>\n'
+            'PTransform: <DatasetToChunks>\n'
+            'Chunks:     40B (x: 5, split_vars=False)\n'
+            'Template:   80B (2 chunks)\n'
+            '    Dimensions:'
+        ).replace('DatasetToChunks', 'DatasetToChunks.*'),
+    )
+
   def test_from_xarray(self):
     ds = xarray.Dataset({'foo': ('x', np.arange(10))})
     beam_ds = xbeam.Dataset.from_xarray(ds, {'x': 5})
@@ -33,11 +72,9 @@ def test_from_xarray(self):
     self.assertEqual(beam_ds.template.keys(), {'foo'})
     self.assertEqual(beam_ds.chunks, {'x': 5})
     self.assertFalse(beam_ds.split_vars)
+    self.assertEqual(beam_ds.bytes_per_chunk, 40)
+    self.assertEqual(beam_ds.chunk_count, 2)
     self.assertRegex(beam_ds.ptransform.label, r'^from_xarray_\d+$')
-    self.assertEqual(
-        repr(beam_ds).split('\n')[0],
-        '<xarray_beam.Dataset[x: 5][split_vars=False]>',
-    )
     expected = [
         (xbeam.Key({'x': 0}), ds.head(x=5)),
         (xbeam.Key({'x': 5}), ds.tail(x=5)),
@@ -240,8 +277,8 @@ def test_infer_new_chunks_uneven_new_size_error(self):
     with self.assertRaisesWithLiteralMatch(
         ValueError,
         "cannot infer new chunks for dimension 'x' with changed size "
-        "10 -> 3: the 2 chunks along this dimension do not evenly divide "
-        "the new size 3",
+        '10 -> 3: the 2 chunks along this dimension do not evenly divide '
+        'the new size 3',
     ):
       xbeam_dataset._infer_new_chunks(
           old_sizes={'x': 10}, old_chunks={'x': 5}, new_sizes={'x': 3}
@@ -378,11 +415,36 @@ def test_rechunk_split_vars(self):
 
 class EndToEndTest(test_util.TestCase):
 
+  def test_bytes_per_chunk_and_chunk_count(self):
+    source_ds = test_util.dummy_era5_surface_dataset(
+        variables=2, latitudes=73, longitudes=144, times=365, freq='24H'
+    )
+
+    xbeam_ds = xbeam.Dataset.from_xarray(
+        source_ds, {'time': 90}, split_vars=False
+    )
+    self.assertEqual(
+        xbeam_ds.chunks, {'time': 90, 'latitude': 73, 'longitude': 144}
+    )
+    self.assertEqual(xbeam_ds.bytes_per_chunk, 2 * 73 * 144 * 90 * 4)
+    self.assertEqual(xbeam_ds.chunk_count, 5)
+
+    xbeam_ds = xbeam.Dataset.from_xarray(
+        source_ds, {'time': 90}, split_vars=True
+    )
+    self.assertEqual(
+        xbeam_ds.chunks, {'time': 90, 'latitude': 73, 'longitude': 144}
+    )
+    self.assertEqual(xbeam_ds.bytes_per_chunk, 73 * 144 * 90 * 4)
+    self.assertEqual(xbeam_ds.chunk_count, 5 * 2)
+
   def test_docstring_example(self):
     input_path = self.create_tempdir('source').full_path
     output_path = self.create_tempdir('output').full_path
 
-    source_ds = test_util.dummy_era5_surface_dataset(times=365, freq='24H')
+    source_ds = test_util.dummy_era5_surface_dataset(
+        variables=2, latitudes=73, longitudes=144, times=365, freq='24H'
+    )
     source_ds.chunk({'time': 90}).to_zarr(input_path)
 
     transform = (