POC: Zarr weather data → Apache Arrow Fixed/Variable Shape Tensor extension types

create_zarr.py

"""
Create a realistic weather Zarr store with:
  - Coordinates: time(4), lat(8), lon(16), pressure_level(5)
  - Variables:
    - temperature(time, lat, lon, pressure_level) -- 4D, has pressure levels
    - wind_u(time, lat, lon, pressure_level)      -- 4D, has pressure levels
    - precipitation(time, lat, lon)                -- 3D, NO pressure levels
    - surface_pressure(time, lat, lon)             -- 3D, NO pressure levels

This mirrors real NWP / reanalysis data where surface variables lack a
vertical dimension while upper-air variables are defined on pressure levels.
"""

import numpy as np
import xarray as xr
from pathlib import Path
import shutil

STORE = Path("weather.zarr")


def main():
    if STORE.exists():
        shutil.rmtree(STORE)

    rng = np.random.default_rng(42)

    # -- coordinates ----------------------------------------------------------
    time = np.arange("2024-01-01", "2024-01-05", dtype="datetime64[D]")  # 4 steps
    lat = np.linspace(-90, 90, 8)
    lon = np.linspace(-180, 180, 16, endpoint=False)
    pressure_level = np.array([1000, 850, 500, 300, 200], dtype=np.int32)  # hPa

    # -- 4D upper-air variables (time, lat, lon, pressure_level) --------------
    shape_4d = (len(time), len(lat), len(lon), len(pressure_level))
    temperature = 250 + 50 * rng.random(shape_4d, dtype=np.float32)
    wind_u = -20 + 40 * rng.random(shape_4d, dtype=np.float32)

    # -- 3D surface variables (time, lat, lon) — no pressure_level ------------
    shape_3d = (len(time), len(lat), len(lon))
    precipitation = 50 * rng.random(shape_3d, dtype=np.float32)
    surface_pressure = 950 + 100 * rng.random(shape_3d, dtype=np.float32)

    ds = xr.Dataset(
        {
            "temperature": (["time", "lat", "lon", "pressure_level"], temperature, {
                "units": "K",
                "long_name": "Air Temperature",
            }),
            "wind_u": (["time", "lat", "lon", "pressure_level"], wind_u, {
                "units": "m/s",
                "long_name": "U-component of Wind",
            }),
            "precipitation": (["time", "lat", "lon"], precipitation, {
                "units": "mm",
                "long_name": "Total Precipitation",
            }),
            "surface_pressure": (["time", "lat", "lon"], surface_pressure, {
                "units": "hPa",
                "long_name": "Surface Pressure",
            }),
        },
        coords={
            "time": time,
            "lat": lat,
            "lon": lon,
            "pressure_level": pressure_level,
        },
    )

    ds.to_zarr(STORE)
    print(f"Created {STORE}")
    print(ds)


if __name__ == "__main__":
    main()

pyproject.toml

[project]
name = "arrow-tensor-zarr"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.12"
dependencies = [
    "numpy>=2.4.4",
    "pyarrow>=23.0.1",
    "xarray>=2026.2.0",
    "zarr>=3.1.6",
]

zarr_to_arrow_tensor.py

"""
POC: Read a Zarr weather store into Apache Arrow tensor extension types.

Strategy
--------
The Zarr store has variables with *different dimensionality*:
  - 4D: temperature, wind_u         → (time, lat, lon, pressure_level)
  - 3D: precipitation, surface_pressure → (time, lat, lon)

We slice along the time axis so each row = one timestep.

1. **Fixed Shape Tensor** (arrow.fixed_shape_tensor)
   Used when every row has an identical shape, i.e. we handle each variable
   independently.  A 4D variable sliced per-timestep gives (lat, lon, plev);
   a 3D variable gives (lat, lon).

2. **Variable Shape Tensor** (arrow.variable_shape_tensor)
   Used when we want a *single column* that holds tensors of different shape
   per row.  We demonstrate this by packing ALL variables for a given timestep
   into one column — some rows are (lat, lon, plev) and others are (lat, lon).
   We keep (lat, lon) as uniform dims and mark the pressure_level axis as
   variable (present=plev, absent=1 or 0).

Both approaches are shown and the resulting Arrow tables are printed, then
written to IPC files so downstream readers can verify the extension metadata.
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pyarrow as pa
import zarr


# ---------------------------------------------------------------------------
# Variable Shape Tensor — custom ExtensionType (not yet a pyarrow builtin)
# ---------------------------------------------------------------------------

class VariableShapeTensorType(pa.ExtensionType):
    """Arrow canonical variable_shape_tensor implemented as a pyarrow ExtensionType."""

    def __init__(
        self,
        value_type: pa.DataType,
        ndim: int,
        dim_names: list[str] | None = None,
        permutation: list[int] | None = None,
        uniform_shape: list[int | None] | None = None,
    ):
        self._ndim = ndim
        self._dim_names = dim_names
        self._permutation = permutation
        self._uniform_shape = uniform_shape

        storage = pa.struct([
            pa.field("data", pa.list_(value_type)),
            pa.field("shape", pa.list_(pa.int32(), ndim)),
        ])
        super().__init__(storage, "arrow.variable_shape_tensor")

    def __arrow_ext_serialize__(self) -> bytes:
        md: dict = {}
        if self._dim_names:
            md["dim_names"] = self._dim_names
        if self._permutation:
            md["permutation"] = self._permutation
        if self._uniform_shape:
            md["uniform_shape"] = self._uniform_shape
        return json.dumps(md).encode()

    @classmethod
    def __arrow_ext_deserialize__(cls, storage_type, serialized):
        md = json.loads(serialized.decode())
        vtype = storage_type.field("data").type.value_type
        ndim = storage_type.field("shape").type.list_size
        return cls(vtype, ndim, **md)


def _register_vst():
    try:
        pa.unregister_extension_type("arrow.variable_shape_tensor")
    except KeyError:
        pass
    pa.register_extension_type(VariableShapeTensorType(pa.float32(), 1))


# ---------------------------------------------------------------------------
# helpers
# ---------------------------------------------------------------------------

def _zarr_store(path: str | Path) -> zarr.Group:
    return zarr.open_group(path, mode="r")


def _load_array(grp: zarr.Group, name: str) -> np.ndarray:
    return np.asarray(grp[name])


# ---------------------------------------------------------------------------
# Approach 1 — Fixed Shape Tensor (one column per variable)
# ---------------------------------------------------------------------------

def build_fixed_shape_table(store_path: Path) -> pa.Table:
    """Each variable becomes its own FixedShapeTensorArray column, sliced by time."""

    grp = _zarr_store(store_path)

    columns: dict[str, pa.Array] = {}
    fields: list[pa.Field] = []

    # time coordinate — zarr v3 stores as raw int; reconstruct datetime
    time_dt = _load_time(grp)
    n_time = len(time_dt)
    time_arr = pa.array(time_dt.astype("datetime64[ms]"), type=pa.timestamp("ms"))
    fields.append(pa.field("time", pa.timestamp("ms")))
    columns["time"] = time_arr

    # 4D vars: per-timestep shape = (lat, lon, pressure_level)
    for var_name in ("temperature", "wind_u"):
        data = _load_array(grp, var_name)  # (time, lat, lon, plev)
        per_step_shape = list(data.shape[1:])
        dim_names = _get_dim_names(grp, var_name)[1:]  # drop 'time'

        tensor_type = pa.fixed_shape_tensor(
            pa.float32(), per_step_shape, dim_names=dim_names,
        )
        flat = data.reshape(n_time, -1).astype(np.float32)
        storage = pa.FixedSizeListArray.from_arrays(
            pa.array(flat.ravel(), type=pa.float32()),
            int(np.prod(per_step_shape)),
        )
        arr = pa.ExtensionArray.from_storage(tensor_type, storage)

        fields.append(pa.field(var_name, tensor_type))
        columns[var_name] = arr

    # 3D vars: per-timestep shape = (lat, lon)
    for var_name in ("precipitation", "surface_pressure"):
        data = _load_array(grp, var_name)  # (time, lat, lon)
        per_step_shape = list(data.shape[1:])
        dim_names = _get_dim_names(grp, var_name)[1:]

        tensor_type = pa.fixed_shape_tensor(
            pa.float32(), per_step_shape, dim_names=dim_names,
        )
        flat = data.reshape(n_time, -1).astype(np.float32)
        storage = pa.FixedSizeListArray.from_arrays(
            pa.array(flat.ravel(), type=pa.float32()),
            int(np.prod(per_step_shape)),
        )
        arr = pa.ExtensionArray.from_storage(tensor_type, storage)

        fields.append(pa.field(var_name, tensor_type))
        columns[var_name] = arr

    schema = pa.schema(fields)
    return pa.table(columns, schema=schema)


def _get_dim_names(grp: zarr.Group, var_name: str) -> list[str]:
    """Read dimension_names from Zarr v3 consolidated metadata or variable metadata."""
    # Zarr v3: dimension_names stored in array metadata (exposed via consolidated metadata)
    store_path = grp.store.root
    if store_path is None:
        store_path = grp.store.path
    import json as _json
    meta_path = Path(store_path) / "zarr.json"
    if meta_path.exists():
        with open(meta_path) as f:
            root_meta = _json.load(f)
        cm = root_meta.get("consolidated_metadata", {}).get("metadata", {})
        if var_name in cm and "dimension_names" in cm[var_name]:
            return cm[var_name]["dimension_names"]
    # Fallback: check per-array zarr.json
    arr_meta_path = Path(store_path) / var_name / "zarr.json"
    if arr_meta_path.exists():
        with open(arr_meta_path) as f:
            arr_meta = _json.load(f)
        if "dimension_names" in arr_meta:
            return arr_meta["dimension_names"]
    # Last resort: xarray v2 convention
    attrs = dict(grp[var_name].attrs)
    return attrs.get("_ARRAY_DIMENSIONS", [])


def _load_time(grp: zarr.Group) -> np.ndarray:
    """Reconstruct datetime64 from Zarr time coordinate (may be raw int under v3)."""
    raw = _load_array(grp, "time")
    attrs = dict(grp["time"].attrs)
    units = attrs.get("units", "")
    if "days since" in units:
        origin = np.datetime64(units.split("since")[-1].strip())
        return origin + raw.astype("timedelta64[D]")
    if raw.dtype.kind == "M":
        return raw
    # fallback: assume epoch seconds
    return (np.datetime64("1970-01-01") + raw.astype("timedelta64[s]")).astype("datetime64[ms]")


# ---------------------------------------------------------------------------
# Approach 2 — Variable Shape Tensor (mixed-shape column)
# ---------------------------------------------------------------------------

def build_variable_shape_table(store_path: Path) -> pa.Table:
    """
    One row per (timestep, variable).  All variables share a single
    VariableShapeTensorType column.  4D vars produce (lat, lon, plev) tensors;
    3D vars produce (lat, lon, 1) tensors — the third dim is variable.
    """
    grp = _zarr_store(store_path)

    time_dt = _load_time(grp)
    var_names = ["temperature", "wind_u", "precipitation", "surface_pressure"]

    ndim = 3  # max rank after dropping time
    lat_size = len(_load_array(grp, "lat"))
    lon_size = len(_load_array(grp, "lon"))

    vst = VariableShapeTensorType(
        pa.float32(),
        ndim,
        dim_names=["lat", "lon", "pressure_level"],
        uniform_shape=[lat_size, lon_size, None],  # pressure_level varies
    )

    row_times = []
    row_varnames = []
    data_lists: list[list[float]] = []
    shape_vals: list[int] = []  # flat list of all shape ints

    for var_name in var_names:
        raw = _load_array(grp, var_name).astype(np.float32)
        dims = _get_dim_names(grp, var_name)

        for t_idx in range(len(time_dt)):
            row_times.append(time_dt[t_idx])
            row_varnames.append(var_name)

            slab = raw[t_idx]  # (lat, lon[, plev])

            if "pressure_level" not in dims:
                # Expand to (lat, lon, 1) so ndim matches
                slab = slab[..., np.newaxis]

            shape_vals.extend(slab.shape)
            data_lists.append(slab.ravel().tolist())

    # Build storage arrays
    n_rows = len(data_lists)
    data_arr = pa.array(data_lists, type=pa.list_(pa.float32()))
    assert len(shape_vals) == n_rows * ndim, (
        f"shape_vals has {len(shape_vals)} ints but expected {n_rows * ndim}"
    )
    shape_arr = pa.FixedSizeListArray.from_arrays(
        pa.array(shape_vals, type=pa.int32()), ndim,
    )
    storage = pa.StructArray.from_arrays(
        [data_arr, shape_arr], names=["data", "shape"],
    )
    tensor_col = pa.ExtensionArray.from_storage(vst, storage)

    table = pa.table({
        "time": pa.array(
            np.array(row_times).astype("datetime64[ms]"),
            type=pa.timestamp("ms"),
        ),
        "variable": pa.array(row_varnames, type=pa.utf8()),
        "tensor": tensor_col,
    })
    return table


# ---------------------------------------------------------------------------
# round-trip verification
# ---------------------------------------------------------------------------

def write_ipc(table: pa.Table, path: Path):
    with pa.ipc.new_file(path, table.schema) as writer:
        writer.write_table(table)
    print(f"  Wrote {path}  ({path.stat().st_size:,} bytes)")


def read_ipc(path: Path) -> pa.Table:
    with pa.ipc.open_file(path) as reader:
        return reader.read_all()


def verify_roundtrip(original: pa.Table, path: Path):
    """Write → read → compare."""
    write_ipc(original, path)
    restored = read_ipc(path)
    assert original.schema.equals(restored.schema), "Schema mismatch!"
    assert original.num_rows == restored.num_rows, "Row count mismatch!"
    for col_name in original.column_names:
        orig_col = original.column(col_name)
        rest_col = restored.column(col_name)
        assert orig_col.equals(rest_col), f"Column {col_name} mismatch!"
    print(f"  Round-trip OK  ({original.num_rows} rows, {len(original.schema)} cols)")


# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------

def main():
    _register_vst()
    store = Path("weather.zarr")

    print("=" * 70)
    print("APPROACH 1: Fixed Shape Tensor  (one column per variable)")
    print("=" * 70)
    fixed_table = build_fixed_shape_table(store)
    print(fixed_table.schema)
    print()
    print(fixed_table)
    print()
    verify_roundtrip(fixed_table, Path("fixed_shape.arrow"))

    print()
    print("=" * 70)
    print("APPROACH 2: Variable Shape Tensor  (mixed-shape single column)")
    print("=" * 70)
    vst_table = build_variable_shape_table(store)
    print(vst_table.schema)
    print()
    print(vst_table)
    print()
    verify_roundtrip(vst_table, Path("variable_shape.arrow"))

    # -- Demonstrate reading back a single tensor from each approach ----------
    print()
    print("=" * 70)
    print("TENSOR EXTRACTION DEMO")
    print("=" * 70)

    # Fixed shape: reconstruct numpy from row 0 of temperature
    temp_col = fixed_table.column("temperature")
    tensor_type = temp_col.type
    row0_flat = temp_col[0].as_py()
    row0_np = np.array(row0_flat, dtype=np.float32).reshape(tensor_type.shape)
    print(f"\n  temperature[t=0] shape: {row0_np.shape}  "
          f"dims: {tensor_type.dim_names}")
    print(f"  min={row0_np.min():.1f}  max={row0_np.max():.1f}  mean={row0_np.mean():.1f}")

    # Variable shape: reconstruct numpy from an arbitrary row
    tensor_col = vst_table.column("tensor")
    for i in [0, 8]:  # row 0 = temperature t=0 (4D), row 8 = precipitation t=0 (3D)
        row = tensor_col[i].as_py()
        shape = row["shape"]
        data = np.array(row["data"], dtype=np.float32).reshape(shape)
        vname = vst_table.column("variable")[i].as_py()
        print(f"\n  {vname}[t=0] shape: {tuple(shape)}  "
              f"min={data.min():.1f}  max={data.max():.1f}")


if __name__ == "__main__":
    main()