Add spatialdata-experimental-writer Python package

python/spatialdata-experimental-writer/README.md

@@ -0,0 +1,132 @@
+# spatialdata-experimental-writer
+
+Experimental vector optimization writers for browser-oriented SpatialData
+rendering.
+
+The initial writer targets Vitessce-compatible Morton-sorted Points Parquet:
+
+- `x`, `y`, optional `z` coordinates are preserved.
+- `morton_code_2d` is added using 16 bits per axis.
+- the first 2–4 rows are sentinel/extreme rows with `morton_code_2d == 0`;
+  readers can infer the full point bounding box from these rows.
+- `{feature_key}_codes` (for example `feature_name_codes`) are added when
+  `feature_key` is set in element attrs.
+- string/categorical columns are placed at the right side of the table.
+- row-group size is controlled when writing Parquet.
+- intermediate Morton uint columns are not persisted in the output Parquet.
+
+Morton v1 belongs on the **canonical** element path
+`points/<key>/points.parquet`. Use `--experimental` only for layouts that
+standard readers cannot consume (see
+[ADR 0002](../../docs/adr/0002-spatially-aware-vector-loading.md)).
+
+Feature / gene filtering in the browser is documented in ADR 0002; pass integer
+`featureCodes` through `@spatialdata/core` `loadPointsInBounds()` and
+`PointsLayerConfig.featureCodes` in `@spatialdata/vis`.
+
+## Install
+
+```bash
+cd python/spatialdata-experimental-writer
+uv sync
+```
+
+For the interactive TUI:
+
+```bash
+uv sync --group tui
+```
+
+## Interactive TUI
+
+```bash
+uv run spatialdata-experimental-writer tui
+uv run spatialdata-experimental-writer tui ~/data/xenium_rep1_io.zarr
+```
+
+The TUI wraps all writer commands:
+
+1. Pick a command from the home menu.
+2. Enter paths and options on guided forms (Zarr store path is pre-filled when
+   passed on the command line).
+3. Confirm before any in-place overwrite of canonical `points/<key>/points.parquet`.
+4. Watch run output, then review post-write verification checks.
+
+Morton verification checks after Morton writes:
+
+| Check | Meaning |
+|-------|---------|
+| `column_present` | `morton_code_2d` column exists |
+| `sentinel_prefix` | First 2–4 rows have `morton_code_2d == 0` |
+| `sentinel_bbox` | Sentinel rows encode full dataset x/y bounds |
+| `morton_monotonic` | Morton codes non-decreasing after sentinels |
+| `row_group_sentinels` | Row group 0 contains only sentinel rows |
+| `no_uint_intermediates` | No persisted `*_uint` staging columns |
+
+Multiscale and index-permutation runs show schema/manifest checks instead.
+
+## Commands
+
+```bash
+# List Points elements in a store
+uv run spatialdata-experimental-writer list-points ~/data/xenium.zarr
+
+# Morton-sort transcripts in-place on canonical points/<key>/points.parquet
+uv run spatialdata-experimental-writer morton-points-from-zarr \
+  ~/data/xenium.zarr --points-key transcripts
+
+# Optional: write to points.experimental/ instead of canonical path
+uv run spatialdata-experimental-writer morton-points-from-zarr \
+  ~/data/xenium.zarr --points-key transcripts --experimental
+
+# Build a derivative store with transcript index sort permutations
+uv run spatialdata-experimental-writer write-index-permutations \
+  ~/data/xenium_rep1_io.zarr \
+  ~/data/xenium_rep1_index-permutations.zarr
+
+# Morton-sort a CSV or Parquet file
+uv run spatialdata-experimental-writer morton-points input.csv output.parquet \
+  --feature-key feature_name
+```
+
+## Xenium workflow
+
+Standard sandbox datasets are listed in the
+[spatialdata datasets docs](https://spatialdata.scverse.org/en/stable/tutorials/notebooks/datasets/README.html):
+
+| Dataset | URL |
+|---------|-----|
+| `xenium_rep1_io.zarr` | `https://s3.embl.de/spatialdata/spatialdata-sandbox/xenium_rep1_io.zarr/` |
+| `xenium_rep2_io.zarr` | `https://s3.embl.de/spatialdata/spatialdata-sandbox/xenium_rep2_io.zarr/` |
+| `visium_associated_xenium_io.zarr` | `https://s3.embl.de/spatialdata/spatialdata-sandbox/visium_associated_xenium_io.zarr/` |
+
+After downloading a store locally:
+
+```bash
+uv run spatialdata-experimental-writer morton-points-from-zarr \
+  ~/data/spatialdata/sdata_inputs/xenium_rep1_io.zarr \
+  --points-key transcripts
+```
+
+This replaces `points/transcripts/points.parquet` in place (single-file output;
+multipart source directories are replaced). Open the store in `@spatialdata/vis`
+with `experimentalOptimizations="auto"` to use TileLayer row-group reads.
+
+For sort-strategy benchmarks on a **copy** of the store:
+
+```bash
+uv run spatialdata-experimental-writer write-index-permutations \
+  ~/data/spatialdata/sdata_inputs/xenium_rep1_io.zarr \
+  ~/data/spatialdata/sdata_inputs/xenium_rep1_index-permutations.zarr \
+  --max-rows 500000
+
+uv run python scripts/benchmark_points_index.py \
+  ~/data/spatialdata/sdata_inputs/xenium_rep1_index-permutations.zarr
+```
+
+## Multiscale hook
+
+The package also includes a Padua-style multiscale Parquet writer that stores
+`spatialdata_multiscale` JSON metadata in the Parquet schema. That layout is
+non-standard for morton-points v1 and belongs under `points.experimental/` if
+persisted.

python/spatialdata-experimental-writer/pyproject.toml

@@ -0,0 +1,37 @@
+[project]
+name = "spatialdata-experimental-writer"
+version = "0.1.0"
+description = "Experimental SpatialData vector optimization writers"
+requires-python = ">=3.12"
+readme = "README.md"
+license = { text = "MIT" }
+authors = [{ name = "SpatialData.js contributors" }]
+dependencies = [
+    "numpy>=2.0",
+    "pandas>=2.2",
+    "pyarrow>=18",
+]
+
+[project.scripts]
+spatialdata-experimental-writer = "spatialdata_experimental_writer.cli:main"
+
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.uv]
+package = true
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[dependency-groups]
+dev = [
+    "pytest>=8.0",
+]
+tui = [
+    "textual>=1.0",
+]

python/spatialdata-experimental-writer/scripts/benchmark_points_index.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python3
+"""Benchmark points index permutations using index-manifest.json."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from pathlib import Path
+
+import pandas as pd
+import pyarrow.parquet as pq
+
+
+def _load_bounds(manifest: dict, scenario_id: str | None) -> dict[str, float]:
+    scenarios = manifest.get("benchmark_scenarios") or []
+    if scenario_id:
+        for scenario in scenarios:
+            if scenario.get("id") == scenario_id:
+                return scenario["bounds"]
+        raise SystemExit(f"Unknown scenario id: {scenario_id}")
+    if scenarios:
+        return scenarios[0]["bounds"]
+    raise SystemExit("Manifest has no benchmark_scenarios")
+
+
+def _feature_codes(manifest: dict, scenario_id: str | None) -> list[int] | None:
+    scenarios = manifest.get("benchmark_scenarios") or []
+    if not scenario_id:
+        return None
+    for scenario in scenarios:
+        if scenario.get("id") == scenario_id:
+            codes = scenario.get("feature_codes")
+            return list(codes) if codes is not None else None
+    return None
+
+
+def _parquet_path(store: Path, element_path: str) -> Path:
+    return store / element_path / "points.parquet"
+
+
+def _read_rows_in_bounds(
+    parquet_path: Path,
+    bounds: dict[str, float],
+    feature_codes: list[int] | None,
+    feature_key: str | None,
+) -> tuple[int, int]:
+    if parquet_path.is_dir():
+        parts = sorted(parquet_path.glob("part.*.parquet"))
+        if not parts:
+            raise FileNotFoundError(f"No parquet parts under {parquet_path}")
+        frames = [pd.read_parquet(part) for part in parts]
+        df = pd.concat(frames, ignore_index=True)
+        bytes_read = sum(part.stat().st_size for part in parts)
+    else:
+        bytes_read = parquet_path.stat().st_size
+        df = pd.read_parquet(parquet_path)
+
+    mask = (
+        (df["x"] >= bounds["minX"])
+        & (df["x"] <= bounds["maxX"])
+        & (df["y"] >= bounds["minY"])
+        & (df["y"] <= bounds["maxY"])
+    )
+    if feature_codes is not None:
+        code_column = f"{feature_key}_codes" if feature_key else "feature_name_codes"
+        if code_column not in df.columns:
+            raise KeyError(f"Missing feature code column {code_column!r}")
+        mask &= df[code_column].isin(feature_codes)
+    return int(mask.sum()), int(bytes_read)
+
+
+def _estimate_row_group_bytes(parquet_path: Path, bounds: dict[str, float]) -> int | None:
+    if not parquet_path.is_file():
+        return None
+    if "morton_code_2d" not in pq.ParquetFile(parquet_path).schema_arrow.names:
+        return None
+    # Upper bound only: full file size when row-group APIs are unavailable in this script.
+    return parquet_path.stat().st_size
+
+
+def benchmark_store(
+    store: Path,
+    *,
+    scenario_id: str | None,
+    conditions: list[str] | None,
+) -> list[dict]:
+    manifest_path = store / "index-manifest.json"
+    if not manifest_path.exists():
+        raise FileNotFoundError(f"Missing index-manifest.json under {store}")
+    manifest = json.loads(manifest_path.read_text())
+    bounds = _load_bounds(manifest, scenario_id)
+    feature_codes = _feature_codes(manifest, scenario_id)
+    feature_key = manifest.get("feature_key")
+    selected = conditions or [entry["id"] for entry in manifest.get("conditions", [])]
+
+    results: list[dict] = []
+    for condition in manifest.get("conditions", []):
+        condition_id = condition["id"]
+        if condition_id not in selected:
+            continue
+        element_path = condition["element_path"]
+        parquet_path = _parquet_path(store, element_path)
+        started = time.perf_counter()
+        try:
+            rows, bytes_read = _read_rows_in_bounds(
+                parquet_path, bounds, feature_codes, feature_key
+            )
+            row_group_hint = _estimate_row_group_bytes(parquet_path, bounds)
+        except Exception as error:  # noqa: BLE001 - report per condition
+            results.append(
+                {
+                    "condition": condition_id,
+                    "element_path": element_path,
+                    "error": str(error),
+                }
+            )
+            continue
+        elapsed_ms = (time.perf_counter() - started) * 1000
+        results.append(
+            {
+                "condition": condition_id,
+                "element_path": element_path,
+                "sort_order": condition.get("sort_order"),
+                "tiling_kind": condition.get("tiling_kind"),
+                "rows_in_bounds": rows,
+                "bytes_read_estimate": bytes_read,
+                "morton_row_group_bytes_upper_bound": row_group_hint,
+                "latency_ms": round(elapsed_ms, 2),
+                "bounds": bounds,
+                "feature_codes": feature_codes,
+            }
+        )
+    return results
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Benchmark points index permutations from index-manifest.json"
+    )
+    parser.add_argument("store", type=Path, help="Derivative Zarr store path")
+    parser.add_argument("--scenario", metavar="ID", help="benchmark_scenarios id")
+    parser.add_argument(
+        "--conditions",
+        metavar="IDS",
+        help="Comma-separated condition ids (default: all in manifest)",
+    )
+    args = parser.parse_args()
+    condition_ids = args.conditions.split(",") if args.conditions else None
+    results = benchmark_store(args.store, scenario_id=args.scenario, conditions=condition_ids)
+    print(json.dumps({"store": str(args.store), "results": results}, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()

python/spatialdata-experimental-writer/src/spatialdata_experimental_writer/__init__.py

@@ -0,0 +1,17 @@
+from .points import (
+    MORTON_CODE_2D_COLUMN,
+    MORTON_CODE_EXTREME_VALUE_INDICATOR,
+    build_spatialdata_multiscale_metadata,
+    morton_sort_points,
+    write_morton_points_parquet,
+    write_multiscale_points_parquet,
+)
+
+__all__ = [
+    "MORTON_CODE_2D_COLUMN",
+    "MORTON_CODE_EXTREME_VALUE_INDICATOR",
+    "build_spatialdata_multiscale_metadata",
+    "morton_sort_points",
+    "write_morton_points_parquet",
+    "write_multiscale_points_parquet",
+]