Update discussion of rechunking, with a new illustration

PiperOrigin-RevId: 814436240

Author: shoyer

docs/_static/pancake-vs-pencil.svg

@@ -55,15 +55,15 @@
         "import pandas as pd\n",
         "\n",
         "xarray_ds = xarray.Dataset(\n",
-        "    {'temperature': (('time', 'longitude', 'latitude'), np.random.randn(365, 180, 90))},\n",
+        "    {'temperature': (('time', 'longitude', 'latitude'), np.random.randn(365, 360, 180))},\n",
         "    coords={'time': pd.date_range('2025-01-01', freq='1D', periods=365)},\n",
         ")\n",
-        "chunks = {'time': 100, 'longitude': 90, 'latitude': 90}\n",
+        "chunks = {'time': '30MB', 'longitude': -1, 'latitude': -1}\n",
         "xbeam_ds = xbeam.Dataset.from_xarray(xarray_ds, chunks)\n",
         "xbeam_ds"
       ],
       "outputs": [],
-      "execution_count": 3
+      "execution_count": 2
     },
     {
       "metadata": {
@@ -83,7 +83,7 @@
         "xarray_ds.chunk(chunks).to_zarr('example_data.zarr', mode='w')"
       ],
       "outputs": [],
-      "execution_count": 4
+      "execution_count": 3
     },
     {
       "metadata": {
@@ -102,56 +102,59 @@
         "\n",
         "All non-trivial computation happens via the embarrasingly parallel `map_blocks` method.\n",
         "\n",
+        "### Chunking strategies\n",
+        "\n",
         "In order for `map_blocks` to work, data needs to be appropriately chunked. Here are a few typical chunking patterns that work well for most needs:\n",
         "\n",
-        "- \"Pencil\" chunks, which group together all times, and parallelize over space. These long and skinny chunks look like a box of pencils:"
+        "- \"Pencil\" chunks, which group together all times, and parallelize over space. These long and skinny chunks look like a box of pencils.\n",
+        "- \"Pancake\" chunks, which group together all spatial locations, and parallelize over time. These flat and wide chunks look like a stack of pancakes.\n",
+        "- Intermediate \"compromise\" chunks, somewhere between these extremes. These optimize for worst-case behavior. "
       ]
     },
     {
       "metadata": {
-        "id": "EpOpwzZS9Rte"
-      },
-      "cell_type": "code",
-      "source": [
-        "xarray_ds.temperature.chunk({'time': -1, 'latitude': 20, 'longitude': 20}).data"
-      ],
-      "outputs": [],
-      "execution_count": 5
-    },
-    {
-      "metadata": {
-        "id": "74t0dh3E9kny"
+        "id": "e7dkTphQyLzz"
       },
       "cell_type": "markdown",
       "source": [
-        "- \"Pancake\" chunks, which group together all spatial locations, and parallelize over time. These flat and wide chunks look like a stack of pancakes:"
+        "![Pancake vs Pencil chunks](./_static/pancake-vs-pencil.svg)\n"
       ]
     },
     {
       "metadata": {
-        "id": "fho4ub-69mD5"
+        "id": "t3APp6uB9yjT"
       },
-      "cell_type": "code",
+      "cell_type": "markdown",
       "source": [
-        "xarray_ds.temperature.chunk({'time': 1, 'latitude': -1, 'longitude': -1}).data"
-      ],
-      "outputs": [],
-      "execution_count": 6
+        "Weather/climate datasets are typically generated and stored in \"pancake\" chunks, but \"pencil\" chunks are more suitable for most analytics queries, which requires large histories of weather at small numbers of locations.\n",
+        "\n",
+        "Using the right chunks is *absolutely essentially* for efficient operations with Xarray-Beam and Zarr. For example, reading data from a single location across all times (a \"pencil\" query) is extremely inefficient for a dataset stored in \"pancake\" chunks -- it would require loading the entire dataset from disk!\n",
+        "\n",
+        "[Like Dask](https://docs.dask.org/en/latest/array-best-practices.html), Xarray-Beam works best with chunks that are 10s to 100s of MB in size, large enough to amortize Python overhead but small enough that chunks fit easily into memory. Smaller chunk sizes in Zarr stores (closer to 10MB) can be convenient to increase the flexiblity of chunking methods when reading the data from disk later\n",
+        "\n",
+        "You can explicitly adjusts chunk sizes with {py:meth}`~xarray_beam.Dataset.rechunk`. The syntax is a mapping from dimension names (or `...`, for all other dimensions) to chunk sizes, which can be any of `-1` (to indicate \"size of the full dimension\"), a positive integer or a string indicating a target size in bytes like `'100MB'`, e.g.,\n",
+        "- `chunks=-1`: no chunking.\n",
+        "- `chunks='100 MB'`: chunk size of ~100 MB, without dimension-specific constraints.\n",
+        "- `chunks={'time': 100}`: fixed size of 100 along time, preserving original chunks along other dimensions.\n",
+        "- `chunks={'time': -1, ...: '100 MB'}`: pencil chunks of ~100 MB each.\n",
+        "- `chunks={'latitude': -1, 'longitude': -1, ...: '100 MB'}`: pancake chunks of ~100 MB each.\n",
+        "\n",
+        "See {py:func}`~xarray_beam.normalize_chunks` for full details."
+      ]
     },
     {
       "metadata": {
-        "id": "t3APp6uB9yjT"
+        "id": "noNiRxGrGX3L"
       },
       "cell_type": "markdown",
       "source": [
-        "Weather/climate datasets are typically generated and stored in \"pancake\" chunks, but \"pencil\" chunks are more suitable for most analytics queries, which requires large histories of weather at small numbers of locations.\n",
         "\n",
-        "Using the right chunks is *absolutely essentially* for efficient operations with Xarray-Beam and Zarr. For example, reading data from a single location across all times (a \"pencil\" query) is extremely inefficient for a dataset stored in \"pancake\" chunks -- it would require loading the entire dataset from disk!\n",
-        "\n",
-        "Rechunking is a fundamentally an expensive operation (it requires multiple complete reads/writes of a dataset from disk), but in Xarray-Beam it's straightforward, via {py:meth}`~xarray_beam.Dataset.rechunk`.\n",
+        "```{warning}\n",
+        "Rechunking between very different schemes is fundamentally an expensive operation (it requires multiple complete reads/writes of a dataset from disk). If performance and flexibility are critical, it may be worth storing multiple copies of your data in different chunking formats.\n",
+        "```\n",
         "\n",
         "```{tip}\n",
-        "Intermediate \"compromise\" chunks can sometimes be a good idea, although if performance and flexibility are critical it may be worth storing multiple copies of your data in different formats. Using Zarr v3's sharding feature in {py:meth}`~xarray_beam.Dataset.to_zarr` to group smaller chunks into shards can also help mitigate the challenges of picking an optimal chunk size.\n",
+        "Using Zarr v3's sharding feature in {py:meth}`~xarray_beam.Dataset.to_zarr` to group smaller chunks (~1 MB) into shards can also help mitigate the challenges of picking an optimal chunk size.\n",
         "```"
       ]
     },
@@ -175,14 +178,15 @@
         "with beam.Pipeline() as p:\n",
         "  p | (\n",
         "      xbeam.Dataset.from_zarr('example_data.zarr')\n",
-        "      .rechunk({'time': -1, 'latitude': 30, 'longitude': 30})\n",
+        "      .rechunk({'time': -1, ...: '100 MB'})\n",
         "      .map_blocks(lambda ds: ds.groupby('time.month').mean())\n",
+        "      .rechunk('10 MB')  # ensure a reasonable min chunk-size for Zarr\n",
         "      .to_zarr('example_climatology.zarr')\n",
         "  )\n",
         "xarray.open_zarr('example_climatology.zarr')"
       ],
       "outputs": [],
-      "execution_count": 7
+      "execution_count": 6
     },
     {
       "metadata": {
@@ -204,14 +208,14 @@
         "with beam.Pipeline() as p:\n",
         "  p | (\n",
         "      xbeam.Dataset.from_zarr('example_data.zarr')\n",
-        "      .rechunk({'time': 10, 'latitude': -1, 'longitude': -1})\n",
+        "      .rechunk({'time': '30MB', 'latitude': -1, 'longitude': -1})\n",
         "      .map_blocks(lambda ds: ds.coarsen(latitude=2, longitude=2).mean())\n",
         "      .to_zarr('example_regrid.zarr')\n",
         "  )\n",
         "xarray.open_zarr('example_regrid.zarr')"
       ],
       "outputs": [],
-      "execution_count": 8
+      "execution_count": 7
     },
     {
       "metadata": {
@@ -241,7 +245,7 @@
         "  print(f'{type(e).__name__}: {e}')"
       ],
       "outputs": [],
-      "execution_count": 9
+      "execution_count": 8
     },
     {
       "metadata": {
@@ -262,7 +266,7 @@
         "ds_beam.map_blocks(lambda ds: ds.compute(), template=ds_beam.template)"
       ],
       "outputs": [],
-      "execution_count": 10
+      "execution_count": 9
     },
     {
       "metadata": {
@@ -293,7 +297,7 @@
         "%ls *.nc"
       ],
       "outputs": [],
-      "execution_count": 14
+      "execution_count": 10
     }
   ],
   "metadata": {