spatialdata-js-util 0.1.0__py3-none-any.whl

spatialdata_js_util/runners.py

@@ -0,0 +1,458 @@
+from __future__ import annotations
+
+import json
+import shutil
+import sys
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Sequence, TypeVar
+
+import pandas as pd
+import pyarrow.dataset as ds
+
+from .errors import WriterCommandError
+from .index_permutations import (
+    DEFAULT_CONDITIONS,
+    FEATURE_SELECTIVITY_CONDITIONS,
+    IndexCondition,
+    write_index_permutations,
+)
+from .points import EncodingPolicy
+from .points_zarr import (
+    DEFAULT_CHUNK_ROWS,
+    DEFAULT_SHARD_CHUNKS,
+    SIDECAR_GROUP,
+    PointsSource,
+    add_points_zarr_to_store,
+    write_points_zarr_store,
+)
+from .points import (
+    build_spatialdata_multiscale_metadata,
+    write_morton_points_parquet,
+    write_multiscale_points_parquet,
+)
+from .store import (
+    copy_points_element_metadata,
+    experimental_points_output_path,
+    list_points_keys,
+    points_parquet_path,
+    read_points_dataframe,
+    read_points_element_attrs,
+    register_points_elements_in_consolidated_metadata,
+    validate_points_key,
+)
+
+_T = TypeVar("_T")
+
+
+@dataclass(frozen=True)
+class MortonZarrOutput:
+    source_parquet: Path
+    output_parquet: Path
+    in_place: bool
+    output_points_key: str | None
+    collection: str
+
+
+def _as_command_error(action: Callable[[], _T]) -> _T:
+    try:
+        return action()
+    except WriterCommandError:
+        raise
+    except (FileNotFoundError, ValueError) as exc:
+        raise WriterCommandError(str(exc)) from exc
+
+
+def read_input_dataframe(path: str | Path) -> pd.DataFrame:
+    input_path = Path(path)
+    if input_path.is_dir():
+        return read_points_dataframe(input_path)
+    suffix = input_path.suffix.lower()
+    if suffix == ".csv":
+        return pd.read_csv(input_path)
+    if suffix in {".parquet", ".pq"}:
+        return pd.read_parquet(input_path)
+    raise WriterCommandError(
+        f"Unsupported input: {path}\n"
+        "Expected a .csv file, .parquet file, or a directory of Parquet parts."
+    )
+
+
+def run_list_points(zarr: str | Path) -> dict[str, Any]:
+    keys = list_points_keys(zarr)
+    if not keys:
+        raise WriterCommandError(f"No Points elements found under {Path(zarr) / 'points'}")
+    return {"zarr": str(zarr), "points_keys": keys}
+
+
+def run_morton_points(
+    input_path: str | Path,
+    output_path: str | Path,
+    *,
+    feature_key: str | None = None,
+    row_group_size: int = 50_000,
+    compression: str = "zstd",
+) -> dict[str, Any]:
+    df = _as_command_error(lambda: read_input_dataframe(input_path))
+    sorted_df = _as_command_error(
+        lambda: write_morton_points_parquet(
+            df,
+            output_path,
+            feature_key=feature_key,
+            row_group_size=row_group_size,
+            compression=compression,
+        )
+    )
+    return {
+        "format": "morton-points",
+        "rows": int(len(sorted_df)),
+        "output": str(output_path),
+        "row_group_size": row_group_size,
+    }
+
+
+def run_points_zarr(
+    input_path: str | Path,
+    output_store: str | Path,
+    *,
+    feature_key: str,
+    coordinates: Sequence[str] = ("x", "y"),
+    points_key: str = "transcripts",
+    tile_size: float = 100.0,
+    n_levels: int | None = None,
+    seed: int = 0,
+    chunk_rows: int = DEFAULT_CHUNK_ROWS,
+    shard_chunks: int = DEFAULT_SHARD_CHUNKS,
+    overwrite: bool = False,
+) -> dict[str, Any]:
+    if Path(output_store).exists() and not overwrite:
+        raise WriterCommandError(
+            f"Output already exists: {output_store}\nChoose another path or enable overwrite."
+        )
+    source = _as_command_error(
+        lambda: _points_source(input_path, [*coordinates, feature_key])
+    )
+    return _as_command_error(
+        lambda: write_points_zarr_store(
+            source,
+            output_store,
+            points_key=points_key,
+            overwrite=overwrite,
+            feature_key=feature_key,
+            coordinates=coordinates,
+            tile_size=tile_size,
+            n_levels=n_levels,
+            seed=seed,
+            chunk_rows=chunk_rows,
+            shard_chunks=shard_chunks,
+            progress=lambda message: print(message, file=sys.stderr, flush=True),
+        )
+    )
+
+
+def run_add_points_zarr(
+    store: str | Path,
+    *,
+    points_key: str,
+    tile_size: float = 100.0,
+    n_levels: int | None = None,
+    seed: int = 0,
+    chunk_rows: int = DEFAULT_CHUNK_ROWS,
+    shard_chunks: int = DEFAULT_SHARD_CHUNKS,
+    overwrite: bool = False,
+) -> dict[str, Any]:
+    existing = Path(store) / "points" / points_key / SIDECAR_GROUP
+    if existing.exists() and not overwrite:
+        raise WriterCommandError(
+            f"Output already exists: {existing}\nChoose another element or enable overwrite."
+        )
+    return _as_command_error(
+        lambda: add_points_zarr_to_store(
+            None,
+            store,
+            points_key=points_key,
+            overwrite=overwrite,
+            tile_size=tile_size,
+            n_levels=n_levels,
+            seed=seed,
+            chunk_rows=chunk_rows,
+            shard_chunks=shard_chunks,
+            progress=lambda message: print(message, file=sys.stderr, flush=True),
+        )
+    )
+
+
+def _points_source(path: str | Path, columns: list[str]) -> pd.DataFrame | PointsSource:
+    """Parquet is streamed a batch at a time, so it need not fit in memory; CSV is read whole."""
+    input_path = Path(path)
+    if input_path.suffix.lower() == ".csv":
+        return read_input_dataframe(input_path)
+    if not input_path.exists():
+        raise FileNotFoundError(f"Input not found: {input_path}")
+    dataset = ds.dataset(input_path, format="parquet")
+
+    def batches() -> Iterator[pd.DataFrame]:
+        for batch in dataset.to_batches(columns=columns, batch_size=4_000_000):
+            yield batch.to_pandas()
+
+    return batches
+
+
+def run_multiscale_points(
+    input_path: str | Path,
+    output_path: str | Path,
+    *,
+    metadata_json: str | Path | None = None,
+    row_group_size: int = 50_000,
+    compression: str = "zstd",
+) -> dict[str, Any]:
+    df = _as_command_error(lambda: read_input_dataframe(input_path))
+    if metadata_json:
+        metadata = _as_command_error(lambda: json.loads(Path(metadata_json).read_text()))
+    else:
+        metadata = _as_command_error(lambda: build_spatialdata_multiscale_metadata(df))
+    _as_command_error(
+        lambda: write_multiscale_points_parquet(
+            df,
+            output_path,
+            metadata=metadata,
+            row_group_size=row_group_size,
+            compression=compression,
+        )
+    )
+    return {
+        "format": "spatialdata_multiscale_points",
+        "rows": int(len(df)),
+        "output": str(output_path),
+        "row_group_size": row_group_size,
+    }
+
+
+def resolve_morton_from_zarr_output(
+    zarr_path: Path,
+    points_key: str,
+    *,
+    output: str | Path | None = None,
+    output_points_key: str | None = None,
+    experimental: bool = False,
+) -> MortonZarrOutput:
+    if output is not None and output_points_key is not None:
+        raise WriterCommandError("Pass either output or output_points_key, not both.")
+    source_parquet = points_parquet_path(zarr_path, points_key)
+    if output:
+        resolved_output = Path(output)
+        return MortonZarrOutput(
+            source_parquet=source_parquet,
+            output_parquet=resolved_output,
+            in_place=resolved_output == source_parquet,
+            output_points_key=None,
+            collection="path",
+        )
+    if experimental:
+        resolved_output_key = validate_points_key(output_points_key or points_key)
+        resolved_output = experimental_points_output_path(zarr_path, resolved_output_key)
+        return MortonZarrOutput(
+            source_parquet=source_parquet,
+            output_parquet=resolved_output,
+            in_place=False,
+            output_points_key=resolved_output_key,
+            collection="points.experimental",
+        )
+    resolved_output_key = validate_points_key(output_points_key or points_key)
+    resolved_output = points_parquet_path(zarr_path, resolved_output_key)
+    return MortonZarrOutput(
+        source_parquet=source_parquet,
+        output_parquet=resolved_output,
+        in_place=resolved_output == source_parquet,
+        output_points_key=resolved_output_key,
+        collection="points",
+    )
+
+
+def run_morton_points_from_zarr(
+    zarr: str | Path,
+    *,
+    points_key: str | None = None,
+    experimental: bool = False,
+    output: str | Path | None = None,
+    output_points_key: str | None = None,
+    feature_key: str | None = None,
+    overwrite: bool = False,
+    row_group_size: int = 50_000,
+    compression: str = "zstd",
+) -> dict[str, Any]:
+    zarr_path = Path(zarr)
+    keys = list_points_keys(zarr_path)
+    if not keys:
+        raise WriterCommandError(f"No Points elements found under {zarr_path / 'points'}")
+
+    resolved_key = points_key
+    if resolved_key is None:
+        if len(keys) == 1:
+            resolved_key = keys[0]
+        else:
+            raise WriterCommandError(
+                "Multiple Points elements found; pass points_key. "
+                f"Available keys: {', '.join(keys)}"
+            )
+    if resolved_key not in keys:
+        raise WriterCommandError(
+            f"Unknown Points element {resolved_key!r}. Available: {', '.join(keys)}"
+        )
+
+    attrs = _as_command_error(lambda: read_points_element_attrs(zarr_path, resolved_key))
+    resolved_feature_key = feature_key or attrs.get("feature_key")
+    output_spec = resolve_morton_from_zarr_output(
+        zarr_path,
+        resolved_key,
+        output=output,
+        output_points_key=output_points_key,
+        experimental=experimental,
+    )
+    if (
+        output_spec.output_parquet.exists()
+        and not output_spec.in_place
+        and not overwrite
+    ):
+        raise WriterCommandError(
+            f"Output already exists: {output_spec.output_parquet}\n"
+            "Choose another element name/path or enable overwrite."
+        )
+    if (
+        output_spec.collection == "points"
+        and output_spec.output_points_key is not None
+        and output_spec.output_points_key != resolved_key
+    ):
+        target_element_dir = zarr_path / "points" / output_spec.output_points_key
+        if target_element_dir.exists() and not overwrite:
+            raise WriterCommandError(
+                f"Points element already exists: points/{output_spec.output_points_key}\n"
+                "Choose another element name or enable overwrite."
+            )
+
+    df = _as_command_error(lambda: read_points_dataframe(output_spec.source_parquet))
+    if output_spec.output_parquet.exists():
+        if output_spec.output_parquet.is_dir():
+            shutil.rmtree(output_spec.output_parquet)
+        else:
+            output_spec.output_parquet.unlink()
+    if (
+        output_spec.collection == "points"
+        and output_spec.output_points_key is not None
+        and output_spec.output_points_key != resolved_key
+    ):
+        _as_command_error(
+            lambda: copy_points_element_metadata(
+                zarr_path,
+                source_key=resolved_key,
+                dest_key=output_spec.output_points_key,
+            )
+        )
+    sorted_df = _as_command_error(
+        lambda: write_morton_points_parquet(
+            df,
+            output_spec.output_parquet,
+            feature_key=resolved_feature_key,
+            row_group_size=row_group_size,
+            compression=compression,
+        )
+    )
+    if (
+        output_spec.collection == "points"
+        and output_spec.output_points_key is not None
+        and output_spec.output_points_key != resolved_key
+    ):
+        _as_command_error(
+            lambda: register_points_elements_in_consolidated_metadata(
+                zarr_path,
+                [resolved_key, output_spec.output_points_key],
+                template_key=resolved_key,
+            )
+        )
+    return {
+        "format": "morton-points",
+        "zarr": str(zarr_path),
+        "points_key": resolved_key,
+        "output_points_key": output_spec.output_points_key,
+        "output_collection": output_spec.collection,
+        "source": str(output_spec.source_parquet),
+        "output": str(output_spec.output_parquet),
+        "in_place": output_spec.in_place,
+        "feature_key": resolved_feature_key,
+        "rows": int(len(sorted_df)),
+        "row_group_size": row_group_size,
+    }
+
+
+def run_write_index_permutations(
+    source_zarr: str | Path,
+    dest_zarr: str | Path,
+    *,
+    points_key: str | None = None,
+    max_rows: int | None = None,
+    condition_ids: Sequence[str] | None = None,
+    overwrite: bool = False,
+    row_group_size: int = 50_000,
+    compression: str = "zstd",
+    encodings: EncodingPolicy = "auto",
+    write_page_index: bool = True,
+) -> dict[str, Any]:
+    selected: tuple[IndexCondition, ...] | None = None
+    if condition_ids:
+        by_id = {
+            condition.id: condition
+            for condition in (*DEFAULT_CONDITIONS, *FEATURE_SELECTIVITY_CONDITIONS)
+        }
+        missing = [value for value in condition_ids if value not in by_id]
+        if missing:
+            raise WriterCommandError(f"Unknown conditions: {', '.join(missing)}")
+        selected = tuple(by_id[value] for value in condition_ids)
+
+    return _as_command_error(
+        lambda: write_index_permutations(
+            source_zarr,
+            dest_zarr,
+            points_key=points_key,
+            max_rows=max_rows,
+            conditions=selected,
+            overwrite=overwrite,
+            row_group_size=row_group_size,
+            compression=compression,
+            encodings=encodings,
+            write_page_index=write_page_index,
+        )
+    )
+
+
+def run_tables_to_csc(
+    source: str | Path,
+    dest: str | Path | None = None,
+    *,
+    tables: Sequence[str] | None = None,
+    layers: bool = True,
+    densify: bool = False,
+    overwrite: bool = False,
+) -> dict[str, Any]:
+    from .tables import convert_store_tables_to_csc
+
+    result = _as_command_error(
+        lambda: convert_store_tables_to_csc(
+            source,
+            dest,
+            tables=tables,
+            layers=layers,
+            densify=densify,
+            overwrite=overwrite,
+        )
+    )
+    return result.manifest
+
+
+def run_list_tables(zarr: str | Path) -> dict[str, Any]:
+    from .store import list_element_keys
+
+    keys = list_element_keys(zarr, "tables")
+    if not keys:
+        raise WriterCommandError(f"No table elements found under {Path(zarr) / 'tables'}")
+    return {"zarr": str(zarr), "table_keys": keys}