Add @export to fix __module__ for Xarray-Beam's public API

Also delete the unused `_ConsolidateBase` from rechunk.py!

PiperOrigin-RevId: 820889448

Author: shoyer

xarray_beam/_src/combiners.py

@@ -64,6 +64,7 @@ def __call__(
     return sum_increment, count_increment
 
 
+@core.export
 @dataclasses.dataclass
 class MeanCombineFn(beam.transforms.CombineFn):
   """CombineFn for computing an arithmetic mean of xarray.Dataset objects."""
@@ -96,6 +97,7 @@ def extract_output(self, sum_count):
       return sum_count
 
 
+@core.export
 @dataclasses.dataclass
 class Mean(beam.PTransform):
   """Calculate the mean over one or more distributed dataset dimensions.

xarray_beam/_src/core.py

@@ -20,7 +20,7 @@
 import itertools
 import math
 import time
-from typing import Generic, TypeVar
+from typing import Any, Generic, TypeVar
 
 import apache_beam as beam
 import immutabledict
@@ -29,13 +29,21 @@
 from xarray_beam._src import threadmap
 
 
-def inc_counter(namespace: str | type, name: str, value: int = 1):
+T = TypeVar('T')
+
+
+def export(obj: T) -> T:
+  obj.__module__ = 'xarray_beam'
+  return obj
+
+
+def inc_counter(namespace: str | type[Any], name: str, value: int = 1):
   """Increments a Beam counter."""
   return beam.metrics.Metrics.counter(namespace, name).inc(value)
 
 
 @contextlib.contextmanager
-def inc_timer_msec(namespace: str | type, name: str) -> Iterator[None]:
+def inc_timer_msec(namespace: str | type[Any], name: str) -> Iterator[None]:
   """Records elapsed time in milliseconds in a Beam counter."""
   start = time.perf_counter()
   yield
@@ -46,6 +54,7 @@ def inc_timer_msec(namespace: str | type, name: str) -> Iterator[None]:
 _DEFAULT = object()
 
 
+@export
 class Key:
   """Key for keeping track of chunks of a distributed Dataset.
 
@@ -172,6 +181,7 @@ def __setstate__(self, state):
 K = TypeVar("K")
 
 
+@export
 def offsets_to_slices(
     offsets: Mapping[K, int],
     sizes: Mapping[K, int],
@@ -306,6 +316,7 @@ def normalize_expanded_chunks(
 )
 
 
+@export
 class DatasetToChunks(beam.PTransform, Generic[DatasetOrDatasets]):
   """Split one or more xarray.Datasets into keyed chunks."""
 
@@ -557,6 +568,7 @@ def _ensure_chunk_is_computed(key: Key, dataset: xarray.Dataset) -> None:
       )
 
 
+@export
 def validate_chunk(key: Key, datasets: DatasetOrDatasets) -> None:
   """Verify that a key and dataset(s) are valid for xarray-beam transforms."""
   if isinstance(datasets, xarray.Dataset):
@@ -586,6 +598,7 @@ def validate_chunk(key: Key, datasets: DatasetOrDatasets) -> None:
         )
 
 
+@export
 class ValidateEachChunk(beam.PTransform):
   """Check that keys and dataset(s) are valid for xarray-beam transforms."""
 

xarray_beam/_src/dataset.py

@@ -80,6 +80,7 @@ def _to_human_size(nbytes: int) -> str:
 UnnormalizedChunks = Mapping[str | types.EllipsisType, int | str] | int | str
 
 
+@core.export
 def normalize_chunks(
     chunks: UnnormalizedChunks,
     template: xarray.Dataset,
@@ -404,6 +405,7 @@ def apply(self, name: str) -> str:
 _get_label = _CountNamer().apply
 
 
+@core.export
 @dataclasses.dataclass
 class Dataset:
   """Experimental high-level representation of an Xarray-Beam dataset."""

xarray_beam/_src/rechunk.py

@@ -163,6 +163,7 @@ def _consolidate_chunks_in_var_group(
     ) from original_error
 
 
+@core.export
 def consolidate_chunks(
     inputs: Iterable[tuple[core.Key, xarray.Dataset]],
     combine_kwargs: Mapping[str, Any] | None = None,
@@ -196,6 +197,7 @@ def consolidate_chunks(
     yield combined_key, combined_dataset
 
 
+@core.export
 def consolidate_variables(
     inputs: Iterable[tuple[core.Key, xarray.Dataset]],
     merge_kwargs: Mapping[str, Any] | None = None,
@@ -236,6 +238,7 @@ def consolidate_variables(
     yield key, dataset
 
 
+@core.export
 def consolidate_fully(
     inputs: Iterable[tuple[core.Key, xarray.Dataset]],
     *,
@@ -286,17 +289,6 @@ def consolidate_fully(
   return core.Key(combined_offsets, combined_vars), dataset  # pytype: disable=wrong-arg-types
 
 
-class _ConsolidateBase(beam.PTransform):
-
-  def expand(self, pcoll):
-    return (
-        pcoll
-        | 'PrependTempKey' >> beam.MapTuple(self._prepend_chunk_key)
-        | 'GroupByTempKeys' >> beam.GroupByKey()
-        | 'Consolidate' >> beam.MapTuple(self._consolidate_chunks)
-    )
-
-
 def _round_chunk_key(
     key: core.Key,
     target_chunks: Mapping[str, int],
@@ -314,6 +306,7 @@ def _round_chunk_key(
   return key.replace(new_offsets)
 
 
+@core.export
 @dataclasses.dataclass
 class ConsolidateChunks(beam.PTransform):
   """Consolidate existing chunks across offsets into bigger chunks."""
@@ -338,6 +331,7 @@ def expand(self, pcoll):
     )
 
 
+@core.export
 class ConsolidateVariables(beam.PTransform):
   """Consolidate existing chunks across variables into bigger chunks."""
 
@@ -393,6 +387,7 @@ def _split_chunk_bounds(
   return list(zip([start] + breaks, breaks + [stop]))
 
 
+@core.export
 def split_chunks(
     key: core.Key,
     dataset: xarray.Dataset,
@@ -424,13 +419,16 @@ def split_chunks(
     yield new_key, new_chunk
 
 
+@core.export
 @dataclasses.dataclass
 class SplitChunks(beam.PTransform):
   """Split existing chunks into smaller chunks."""
 
   target_chunks: Mapping[str, int]
 
-  def _split_chunks(self, key, dataset):
+  def _split_chunks(
+      self, key: core.Key, dataset: xarray.Dataset
+  ) -> Iterator[tuple[core.Key, xarray.Dataset]]:
     target_chunks = {
         k: v for k, v in self.target_chunks.items() if k in dataset.dims
     }
@@ -440,6 +438,7 @@ def expand(self, pcoll):
     return pcoll | beam.FlatMapTuple(self._split_chunks)
 
 
+@core.export
 def split_variables(
     key: core.Key,
     dataset: xarray.Dataset,
@@ -454,6 +453,7 @@ def split_variables(
     yield new_key, new_dataset
 
 
+@core.export
 @dataclasses.dataclass
 class SplitVariables(beam.PTransform):
   """Split existing chunks into a separate chunk per data variable."""
@@ -462,6 +462,7 @@ def expand(self, pcoll):
     return pcoll | beam.FlatMapTuple(split_variables)
 
 
+@core.export
 def in_memory_rechunk(
     inputs: list[tuple[core.Key, xarray.Dataset]],
     target_chunks: Mapping[str, int],
@@ -489,6 +490,7 @@ def expand(self, pcoll):
     return pcoll
 
 
+@core.export
 class Rechunk(beam.PTransform):
   """Rechunk to an arbitrary new chunking scheme with bounded memory usage.
 

xarray_beam/_src/zarr.py

@@ -68,6 +68,7 @@ def _infer_chunks(dataset: xarray.Dataset) -> dict[str, int]:
   return chunks
 
 
+@core.export
 def open_zarr(
     store: ReadableStore, **kwargs: Any
 ) -> tuple[xarray.Dataset, dict[str, int]]:
@@ -102,6 +103,7 @@ def _raise_template_error():
   )
 
 
+@core.export
 def make_template(
     dataset: xarray.Dataset,
     lazy_vars: Set[str] | None = None,
@@ -144,6 +146,7 @@ def make_template(
   return result
 
 
+@core.export
 def replace_template_dims(
     template: xarray.Dataset,
     **dim_replacements: int | np.ndarray | pd.Index | xarray.DataArray,
@@ -458,6 +461,7 @@ def _setup_zarr(
   logging.info('finished setting up Zarr')
 
 
+@core.export
 def setup_zarr(
     template: xarray.Dataset,
     store: WritableStore,
@@ -511,6 +515,7 @@ def setup_zarr(
   )
 
 
+@core.export
 def validate_zarr_chunk(
     key: core.Key,
     chunk: xarray.Dataset,
@@ -584,6 +589,7 @@ def validate_zarr_chunk(
   # Note that variable names, shapes & dtypes are verified in xarray's to_zarr()
 
 
+@core.export
 def write_chunk_to_zarr(
     key: core.Key,
     chunk: xarray.Dataset,
@@ -629,6 +635,7 @@ def write_chunk_to_zarr(
     ) from e
 
 
+@core.export
 class ChunksToZarr(beam.PTransform):
   """Write keyed chunks to a Zarr store in parallel."""
 
@@ -817,6 +824,7 @@ def expand(self, pcoll):
     )
 
 
+@core.export
 @dataclasses.dataclass
 class DatasetToZarr(beam.PTransform):
   """Write an entire xarray.Dataset to a Zarr store."""