[xarray-beam] fix docs & bug in make_template

Reusing dask.array names for arrays with different shapes was the indirect cause of the error on our docs pages.

I've also updated the build dependencies for our docs to fix other build issues, and added a Python 3.11 run of our unit tests.

Fixes #76

PiperOrigin-RevId: 536022594

Author: shoyer

.github/workflows/ci-build.yml

@@ -16,7 +16,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.9", "3.10"]
+        python-version: ["3.9", "3.10", "3.11"]
     steps:
     - name: Cancel previous
       uses: styfle/[email protected]

.readthedocs.yaml

@@ -6,15 +6,16 @@
 version: 2
 
 build:
-  image: latest
+  os: ubuntu-22.04
+  tools:
+    python: "3.10"
 
 # Build documentation in the docs/ directory with Sphinx
 sphinx:
   configuration: docs/conf.py
 
-# Optionally set the version of Python and requirements required to build your docs
+# Optionally declare the Python requirements required to build your docs
 python:
-  version: "3.8"
   install:
     - method: pip
       path: .

docs/requirements.txt

@@ -1,18 +1,18 @@
 # doc requirements
 Jinja2==3.1.2
-myst-nb==0.17.1
+myst-nb==0.17.2
 myst-parser==0.18.1
-sphinx_rtd_theme==1.1.1
+sphinx_rtd_theme==1.2.1
 sphinx==5.3.0
-scipy==1.10.0
+scipy==1.10.1
 
 # xarray-beam requirements
-apache-beam==2.44.0
-dask==2023.1.0
-immutabledict==2.2.3
-numpy==1.22.4
+apache-beam==2.47.0
+dask==2023.5.1
+immutabledict==2.2.4
+numpy==1.24.3
 pandas==1.5.3
-pooch==1.6.0
-rechunker==0.5.0
-xarray==2023.1.0
-zarr==2.13.6
+pooch==1.7.0
+rechunker==0.5.1
+xarray==2023.5.0
+zarr==2.14.2

setup.py

@@ -42,7 +42,7 @@
 
 setuptools.setup(
     name='xarray-beam',
-    version='0.6.0',
+    version='0.6.1',
     license='Apache 2.0',
     author='Google LLC',
     author_email='[email protected]',

xarray_beam/__init__.py

@@ -45,4 +45,4 @@
     DatasetToZarr,
 )
 
-__version__ = '0.6.0'
+__version__ = '0.6.1'

xarray_beam/_src/zarr.py

@@ -130,9 +130,11 @@ def make_template(
 
   # override the lazy variables
   delayed = dask.delayed(_raise_template_error)()
-  name = 'make_template'
   for k, v in dataset.variables.items():
     if k in lazy_vars:
+      # names of dask arrays are used for keeping track of results, so arrays
+      # with the same name cannot have different shape or dtype
+      name = f"make_template_{'x'.join(map(str, v.shape))}_{v.dtype}"
       result[k].data = dask.array.from_delayed(
           delayed, v.shape, v.dtype, name=name
       )

xarray_beam/_src/zarr_test.py

@@ -337,6 +337,20 @@ def test_infer_zarr_chunks(self):
     ):
       xbeam._src.zarr._infer_zarr_chunks(dataset.chunk({'x': (3, 2, 1)}))
 
+  def test_chunks_to_zarr_docs_demo(self):
+    # verify that the ChunksToChunk demo from our docs works
+    data = np.random.RandomState(0).randn(2920, 25, 53)
+    ds = xarray.Dataset({'temperature': (('time', 'lat', 'lon'), data)})
+    chunks = {'time': 1000, 'lat': 25, 'lon': 53}
+    temp_dir = self.create_tempdir().full_path
+    (
+        test_util.EagerPipeline()
+        | xbeam.DatasetToChunks(ds, chunks)
+        | xbeam.ChunksToZarr(temp_dir)
+    )
+    result = xarray.open_zarr(temp_dir)
+    xarray.testing.assert_identical(result, ds)
+
 
 if __name__ == '__main__':
   absltest.main()