docs(examples): add 'Working with Zarr and NDArray data in SedonaDB' tutorial (#938)

Author: james-willis

.github/workflows/packaging.yml

@@ -163,6 +163,12 @@ jobs:
         run: |
             pip install "python/sedonadb-expr" -v
             pip install "python/sedonadb/[geopandas]" -v
+            # sedonadb-zarr is a pure Rust + pyo3 extension (no s2geography
+            # feature). It's installed so mkdocstrings can import it to build
+            # the Zarr extension API reference. Unset the sedonadb-only maturin
+            # args first so they don't leak into the extension build.
+            unset MATURIN_PEP517_ARGS
+            pip install "python/sedonadb-zarr/" -v
 
       - name: Validate SQL function docs coverage
         run: |

docs/image/zarr-chunks-ndarray-map.png

@@ -0,0 +1,42 @@
+<!---
+  Licensed to the Apache Software Foundation (ASF) under one
+  or more contributor license agreements.  See the NOTICE file
+  distributed with this work for additional information
+  regarding copyright ownership.  The ASF licenses this file
+  to you under the Apache License, Version 2.0 (the
+  "License"); you may not use this file except in compliance
+  with the License.  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing,
+  software distributed under the License is distributed on an
+  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+  KIND, either express or implied.  See the License for the
+  specific language governing permissions and limitations
+  under the License.
+-->
+# Zarr Extension (Python) API Reference
+
+`sedonadb-zarr` reads [Zarr](https://zarr.dev/) groups into the SedonaDB
+raster type. Install it alongside SedonaDB, register the extension on your
+connection, and read a group through its format spec:
+
+```python
+import sedona.db
+import sedonadb_zarr
+
+sd = sedona.db.connect()
+sd.register(sedonadb_zarr.ZarrExtension())
+sd.read("file:///path/to/foo.zarr").show()
+```
+
+For an end-to-end walkthrough, see
+[Working with Zarr and NDArray data in SedonaDB](../working-with-zarr-ndarray-sedonadb.md).
+
+::: sedonadb_zarr.ZarrExtension
+
+::: sedonadb_zarr.Zarr
+    options:
+      members:
+        - with_options

docs/reference/python-zarr.md

@@ -0,0 +1,330 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "cb74c064",
+   "metadata": {},
+   "source": [
+    "# Working with Zarr and NDArray data in SedonaDB\n",
+    "\n",
+    "SedonaDB's raster type is **N-dimensional**: a band isn't limited to a 2-D `(y, x)` grid — it can carry additional axes such as `time`, `year`, or `band`. This makes it a natural fit for *datacubes*: climate reanalyses, satellite time series, and model outputs.\n",
+    "\n",
+    "The `sedonadb-zarr` extension reads [Zarr](https://zarr.dev/) groups (v2 or v3) — local or in cloud object storage — directly into that raster type, so a datacube becomes a table you can query.\n",
+    "\n",
+    "This page walks through loading a real Zarr datacube from object storage, inspecting its dimensions, slicing out a 2-D plane, drawing its chunk grid on a map, and handing a plane to NumPy."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "18f72319",
+   "metadata": {},
+   "source": [
+    "## Install\n",
+    "\n",
+    "`sedonadb-zarr` is an extension, distributed separately from the core SedonaDB package. The examples below also use `sedonadb-expr` (which adds the `.rst` raster accessor used in the DataFrame expressions) and, for the map, `lonboard`:\n",
+    "\n",
+    "```bash\n",
+    "pip install \"apache-sedona[db]\" sedonadb-zarr sedonadb-expr lonboard\n",
+    "```\n",
+    "\n",
+    "`lonboard` is only needed for the map at the end; everything else works without it."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3f348994",
+   "metadata": {},
+   "source": [
+    "## Connect and load\n",
+    "\n",
+    "Register the extension on your connection, then read a Zarr group. We'll use a public, anonymously readable cube: ERA5 rainfall over 2015–2020, stored as a multiscale Zarr pyramid in EPSG:3857. We read one pyramid level and the `rain_ok` rainfall array:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "e079dfe7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sedona.db\n",
+    "import sedonadb_zarr\n",
+    "\n",
+    "sd = sedona.db.connect()\n",
+    "sd.register(sedonadb_zarr.ZarrExtension())\n",
+    "\n",
+    "url = \"https://weathermapdata.rdrn.me/era5_2015_2020_l5.zarr/2\"\n",
+    "# The path doesn't end in `.zarr`, so name the format. `arrays` selects the\n",
+    "# data array to read (the group also holds coordinate / CRS variables).\n",
+    "spec = sedonadb_zarr.Zarr().with_options({\"arrays\": [\"rain_ok\"]})\n",
+    "cube = sd.read(url, format=spec)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ddb7915c",
+   "metadata": {},
+   "source": [
+    "When a group's path *does* end in `.zarr` and needs no options, you can name the format with the string shorthand instead: `sd.read(uri, format=\"zarr\")`.\n",
+    "\n",
+    "`sedonadb-zarr` emits **one row per Zarr chunk**, so the storage layout *is* the data layout. This level tiles its `512 × 512` grid into a `4 × 4` grid of `128 × 128` chunks, and the cube is chunked one year per chunk — so it loads as `16 × 6 = 96` rows, each a single year of one spatial tile:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "71fd0f21",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "96"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "cube.count()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cb563aba",
+   "metadata": {},
+   "source": [
+    "## Inspect the dimensions\n",
+    "\n",
+    "The dimension accessors read the raster's schema only — **no pixel data is loaded** — so they return near-instantly even against a large remote cube. Each row reports its **chunk's** shape, not the full cube extent. All chunks share the same shape here, so we look at one:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "a8e37ba1",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "┌───────┬──────────────┬───────────────┬────────┐\n",
+      "│  ndim ┆     dims     ┆     shape     ┆ n_year │\n",
+      "│ int32 ┆     list     ┆      list     ┆  int64 │\n",
+      "╞═══════╪══════════════╪═══════════════╪════════╡\n",
+      "│     3 ┆ [year, y, x] ┆ [1, 128, 128] ┆      1 │\n",
+      "└───────┴──────────────┴───────────────┴────────┘\n"
+     ]
+    }
+   ],
+   "source": [
+    "cube.select(\n",
+    "    cube.raster.rst.num_dimensions().alias(\"ndim\"),\n",
+    "    cube.raster.rst.dim_names().alias(\"dims\"),\n",
+    "    cube.raster.rst.shape().alias(\"shape\"),\n",
+    "    cube.raster.rst.dim_size(\"year\").alias(\"n_year\"),\n",
+    ").show(1)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "eb79fbdf",
+   "metadata": {},
+   "source": [
+    "Each chunk is 3-dimensional (`[year, y, x]`) with a `128 × 128` spatial footprint — one tile of the full `512 × 512` grid. `n_year = 1` because the cube is chunked one year per chunk: a single row carries one year of one tile."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4ed7f1d1",
+   "metadata": {},
+   "source": [
+    "## Slice out a 2-D plane\n",
+    "\n",
+    "`RS_Slice` selects a single index along a named dimension and drops it. Here each chunk's `year` axis has length 1, so slicing index `0` collapses it, turning every `[1, 128, 128]` chunk into a 2-D `[y, x]` plane — the tile's rainfall field for its year:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "34c51b8d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "┌────────┬────────────┐\n",
+      "│  dims  ┆    shape   │\n",
+      "│  list  ┆    list    │\n",
+      "╞════════╪════════════╡\n",
+      "│ [y, x] ┆ [128, 128] │\n",
+      "└────────┴────────────┘\n"
+     ]
+    }
+   ],
+   "source": [
+    "sliced = cube.select(plane=cube.raster.rst.slice(\"year\", 0))\n",
+    "sliced.select(\n",
+    "    dims=sliced.plane.rst.dim_names(),\n",
+    "    shape=sliced.plane.rst.shape(),\n",
+    ").show(1)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "93948010",
+   "metadata": {},
+   "source": [
+    "`RS_Slice` needs pixel data, so SedonaDB resolves each row's Zarr chunk on demand before slicing — you never call a loader yourself.\n",
+    "\n",
+    "Related accessors reshape a cube in other ways:\n",
+    "\n",
+    "- `cube.raster.rst.slice_range(dim, start, end)` keeps a contiguous range of a dimension instead of a single index.\n",
+    "- `cube.raster.rst.dim_to_band(dim)` / `cube.raster.rst.band_to_dim(name)` move an axis between the dimension list and the band list."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2084a32b",
+   "metadata": {},
+   "source": [
+    "## See where the chunks are — on a map\n",
+    "\n",
+    "Every row is a chunk with a real, georeferenced footprint (the cube declares EPSG:3857), so `RS_Envelope` turns a chunk into its bounding geometry without decoding a single pixel. Reproject the footprints to lon/lat and you can draw the chunk grid straight onto a map:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "be432f29",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from lonboard import viz  # in a notebook with lonboard installed\n",
+    "\n",
+    "f = sd.funcs\n",
+    "chunks = cube.select(geom=f.st_transform(cube.raster.rst.envelope(), \"EPSG:4326\"))\n",
+    "\n",
+    "# Draw outlines only, so the basemap shows through the chunk grid.\n",
+    "viz(\n",
+    "    chunks,\n",
+    "    polygon_kwargs=dict(\n",
+    "        filled=False,\n",
+    "        stroked=True,\n",
+    "        get_line_color=[236, 64, 160],\n",
+    "        line_width_min_pixels=2,\n",
+    "    ),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f802e0f7",
+   "metadata": {},
+   "source": [
+    "![The cube's chunk envelopes drawn as a grid over a world basemap](image/zarr-chunks-ndarray-map.png)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9e0e82d9",
+   "metadata": {},
+   "source": [
+    "Because each year tiles into a `4 × 4` grid, the envelopes lay out that grid over the mapped extent — a picture of the cube's layout, drawn entirely from metadata. A `LIMIT` or row filter trims which chunks you draw (and, later, fetch)."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "69bde03e",
+   "metadata": {},
+   "source": [
+    "## Bring a plane into NumPy\n",
+    "\n",
+    "A raster band carries its bytes, shape, and pixel type, so a materialized band decodes to a correctly-shaped, correctly-typed NumPy array in one call — `Band.to_numpy()`:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "04289b8c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(128, 128) float32\n"
+     ]
+    }
+   ],
+   "source": [
+    "planes = sliced.to_arrow_table()[\"plane\"]\n",
+    "raster = planes[0].as_py()  # one 128x128 spatial tile for its year\n",
+    "band = raster.bands[0].to_numpy()\n",
+    "print(band.shape, band.dtype)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d8f29e2b",
+   "metadata": {},
+   "source": [
+    "Note that `planes[0]` currently forces a copy of the raster out of the Arrow buffer (a pyarrow limitation), so this path is **not yet zero-copy**. Rows correspond to chunks rather than a guaranteed order, so apply your own ordering (or carry a chunk identifier) if you need to know which tile and year a given plane covers."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fb471acd",
+   "metadata": {},
+   "source": [
+    "## Reading from cloud storage\n",
+    "\n",
+    "The same code reads a datacube over S3 or HTTP(S) — only the URI changes. Supported schemes are `file://` (and bare local paths), `s3://`, `http://`, and `https://`.\n",
+    "\n",
+    "For S3, credentials come from the standard AWS environment variables (`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `AWS_REGION`). To read a **public** bucket anonymously, set the region and request unsigned access:\n",
+    "\n",
+    "```bash\n",
+    "export AWS_REGION=us-west-2\n",
+    "export AWS_SKIP_SIGNATURE=true   # read public objects without credentials\n",
+    "```\n",
+    "\n",
+    "```python\n",
+    "# A real public bucket (CarbonPlan), readable with the settings above:\n",
+    "df = sd.read(\"s3://carbonplan-share/zarr-layer-examples/antarctic_era5.zarr\")\n",
+    "```\n",
+    "\n",
+    "### Selecting arrays with the `arrays` option\n",
+    "\n",
+    "By default SedonaDB discovers a group's arrays automatically — from the group's consolidated metadata when present, otherwise by listing the store. The `arrays` option names an explicit subset to read instead (as we did above):\n",
+    "\n",
+    "```python\n",
+    "spec = sedonadb_zarr.Zarr().with_options({\"arrays\": [\"rain_ok\"]})\n",
+    "df = sd.read(url, format=spec)\n",
+    "```\n",
+    "\n",
+    "Naming arrays is needed in two situations:\n",
+    "\n",
+    "- **The store can't list and has no consolidated metadata.** Plain HTTP servers generally can't list directories. Cloud Zarr groups often ship a consolidated-metadata block, so reads typically work without `arrays` — but a group (or sub-group) lacking one can't be auto-discovered over such a store, and you must name the arrays.\n",
+    "- **The group mixes arrays with different shapes or chunk grids.** Every array read together must share one chunk grid, so name a compatible subset (for example, read the data array and leave out a differently-shaped coordinate or CRS variable).\n",
+    "\n",
+    "Because each row corresponds to one chunk, a `LIMIT` or row filter directly bounds how many chunks SedonaDB fetches — handy for sampling a large remote cube before committing to a full scan."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}