rayzin 0.0.1__py3-none-any.whl

rayzin/__init__.py

@@ -0,0 +1,26 @@
+__version__ = "0.0.1"
+
+from rayzin.enums import MetricType, SearchBackendType
+from rayzin.pipeline import (
+    build_manifest,
+    build_manifest_from_cogs,
+    build_manifest_from_zarr,
+    knn_cog_search,
+    knn_zarr_search,
+)
+from rayzin.types import SearchResults
+
+__all__ = [
+    "__version__",
+    # pipeline
+    "knn_zarr_search",
+    "knn_cog_search",
+    # pipeline — manifest build
+    "build_manifest",
+    "build_manifest_from_zarr",
+    "build_manifest_from_cogs",
+    # pipeline API types
+    "MetricType",
+    "SearchBackendType",
+    "SearchResults",
+]

rayzin/enums.py

@@ -0,0 +1,17 @@
+from enum import StrEnum
+
+
+class MetricType(StrEnum):
+    EUCLIDEAN = "euclidean"
+    COSINE = "cosine"
+
+
+class ReaderType(StrEnum):
+    ZARR = "zarr"
+    COG = "cog"
+
+
+class SearchBackendType(StrEnum):
+    NUMPY = "numpy"
+    FAISS_CPU = "faiss_cpu"
+    FAISS_GPU = "faiss_gpu"

rayzin/manifest/__init__.py

@@ -0,0 +1,7 @@
+from rayzin.manifest.filtering import filter_manifest
+from rayzin.manifest.schema import MANIFEST_SCHEMA
+
+__all__ = [
+    "MANIFEST_SCHEMA",
+    "filter_manifest",
+]

rayzin/manifest/build.py

@@ -0,0 +1,166 @@
+from collections.abc import Iterator
+from itertools import product
+from typing import Any
+
+import numpy as np
+import pyarrow as pa  # type: ignore[import-untyped]
+import zarr
+
+from rayzin.manifest.schema import CHUNK_SCHEMA, MANIFEST_SCHEMA, ChunkTable, ManifestTable
+from rayzin.readers.protocol import VectorReader
+from rayzin.readers.zarr_layout import index_axis_names, index_axis_positions
+from rayzin.types import (
+    COL_DIM,
+    COL_SLICE,
+    COL_START,
+    COL_STOP,
+    COL_URL,
+    ChunkRecord,
+    ChunkRef,
+    DimSlice,
+    IndexSlice,
+)
+
+
+def build_zarr_chunk_table(
+    store_url: str,
+    *,
+    array_name: str,
+    store_kwargs: dict[str, Any],
+    embedding_dim_name: str = "embedding",
+) -> ChunkTable:
+    return pa.Table.from_pylist(
+        list(
+            iter_zarr_chunk_slices(
+                store_url,
+                array_name=array_name,
+                store_kwargs=store_kwargs,
+                embedding_dim_name=embedding_dim_name,
+            )
+        ),
+        schema=CHUNK_SCHEMA,
+    )
+
+
+def iter_zarr_chunk_slices(
+    store_url: str,
+    *,
+    array_name: str,
+    store_kwargs: dict[str, Any],
+    embedding_dim_name: str,
+) -> Iterator[ChunkRef]:
+    group = zarr.open_group(
+        store=store_url,
+        mode="r",
+        storage_options=store_kwargs or None,
+    )
+    array = group[array_name]
+    assert isinstance(array, zarr.Array)
+
+    if array.ndim < 2:
+        msg = (
+            "Expected at least one index axis plus the embedding axis, "
+            f"got {array.ndim} dimensions."
+        )
+        raise ValueError(msg)
+
+    axis_positions = index_axis_positions(array, embedding_dim_name=embedding_dim_name)
+    axis_names = index_axis_names(array, embedding_dim_name=embedding_dim_name)
+    index_shape = tuple(array.shape[index] for index in axis_positions)
+    index_chunks = tuple(array.chunks[index] for index in axis_positions)
+    start_ranges = [range(0, size, chunk) for size, chunk in zip(index_shape, index_chunks)]
+
+    for starts in product(*start_ranges):
+        yield ChunkRef(
+            url=store_url,
+            slice=[
+                DimSlice(
+                    dim=axis_name,
+                    start=start,
+                    stop=min(start + chunk, size),
+                )
+                for axis_name, start, chunk, size in zip(
+                    axis_names,
+                    starts,
+                    index_chunks,
+                    index_shape,
+                )
+            ],
+        )
+
+
+def compute_chunk_summary_arrow(batch: ChunkTable, reader: VectorReader) -> ManifestTable:
+    return _compute_chunk_summary(batch, reader)
+
+
+def _compute_chunk_summary(batch: ChunkTable, reader: VectorReader) -> ManifestTable:
+    rows = _chunk_slice_rows(batch)
+    centroids: list[list[float]] = []
+    radii: list[float] = []
+    counts: list[int] = []
+
+    for row in rows:
+        chunk = _chunk_record_from_slice(row)
+        vectors, _shape = reader.read(chunk)
+        if vectors.shape[0] == 0:
+            msg = f"Chunk {chunk[COL_URL]} produced no vectors."
+            raise ValueError(msg)
+
+        centroid = vectors.mean(axis=0).astype(np.float32)
+        distances = np.linalg.norm(vectors - centroid[None, :], axis=1)
+        centroids.append(centroid.tolist())
+        radii.append(float(distances.max()))
+        counts.append(int(vectors.shape[0]))
+
+    return pa.Table.from_arrays(
+        [
+            batch.column(COL_URL),
+            batch.column(COL_SLICE),
+            pa.array(counts, type=pa.int32()),
+            pa.array(centroids, type=pa.list_(pa.float32())),
+            pa.array(radii, type=pa.float32()),
+        ],
+        schema=MANIFEST_SCHEMA,
+    )
+
+
+def _chunk_slice_rows(batch: ChunkTable) -> list[ChunkRef]:
+    urls = batch.column(COL_URL).to_pylist()
+    slices = batch.column(COL_SLICE).to_pylist()
+    return [
+        ChunkRef(
+            url=str(urls[i]),
+            slice=_coerce_index_slice(slices[i]),
+        )
+        for i in range(batch.num_rows)
+    ]
+
+
+def _chunk_record_from_slice(chunk: ChunkRef) -> ChunkRecord:
+    return ChunkRecord(
+        url=chunk[COL_URL],
+        slice=chunk[COL_SLICE],
+        count=0,
+        centroid=np.empty(0, dtype=np.float32),
+        radius=0.0,
+    )
+
+
+def _coerce_index_slice(raw: Any) -> IndexSlice:
+    if not isinstance(raw, list):
+        msg = f"Expected slice to be a list, got {type(raw).__name__}."
+        raise TypeError(msg)
+
+    parts: IndexSlice = []
+    for part in raw:
+        if not isinstance(part, dict):
+            msg = f"Expected each slice entry to be a dict, got {type(part).__name__}."
+            raise TypeError(msg)
+        parts.append(
+            DimSlice(
+                dim=str(part[COL_DIM]),
+                start=int(part[COL_START]),
+                stop=int(part[COL_STOP]),
+            )
+        )
+    return parts

rayzin/manifest/filtering.py

@@ -0,0 +1,60 @@
+from typing import Any
+
+import ray.data
+from ray.data.expressions import Expr
+from shapely.geometry.base import BaseGeometry  # type: ignore[import-untyped]
+
+from rayzin.manifest.spatial import (
+    GeoTransform,
+    chunk_from_row,
+    chunk_polygon,
+    read_zarr_geotransform,
+    validate_aoi_geometry,
+)
+
+
+def filter_manifest(
+    dataset: ray.data.Dataset,
+    *,
+    filter_expr: Expr | None = None,
+    aoi: BaseGeometry | None = None,
+    store_kwargs: dict[str, Any] | None = None,
+) -> ray.data.Dataset:
+    filtered = dataset
+    if isinstance(filter_expr, Expr):
+        filtered = filtered.filter(expr=filter_expr)
+    if aoi is not None:
+        filtered = filtered.filter(
+            ChunkIntersectsAOI,
+            fn_constructor_kwargs={
+                "aoi": aoi,
+                "store_kwargs": store_kwargs or {},
+            },
+        )
+    return filtered
+
+
+class ChunkIntersectsAOI:
+    # TODO: we can get away with opening tha zarr store once since we can extract the geotransform
+    # and analytically check for intersection given a chunks key.
+    # TODO: this should probably be renamed to something zarr specific and we should probable
+    # support cog by forcing manifest to be geoparquet and adding geometry up front per cog tile.
+    def __init__(
+        self,
+        aoi: BaseGeometry,
+        store_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        self._aoi = validate_aoi_geometry(aoi)
+        self._store_kwargs = store_kwargs or {}
+        self._transform_cache: dict[str, GeoTransform] = {}
+
+    def __call__(self, row: dict[str, Any]) -> bool:
+        chunk = chunk_from_row(row)
+        transform = self._transform_cache.get(chunk["url"])
+        if transform is None:
+            transform = read_zarr_geotransform(
+                chunk["url"],
+                store_kwargs=self._store_kwargs,
+            )
+            self._transform_cache[chunk["url"]] = transform
+        return bool(chunk_polygon(chunk, transform).intersects(self._aoi))

rayzin/manifest/schema.py

@@ -0,0 +1,77 @@
+from typing import Annotated, TypeAlias
+
+import pyarrow as pa  # type: ignore[import-untyped]
+
+from rayzin.types import (
+    COL_CENTROID,
+    COL_CHUNK_ID,
+    COL_COUNT,
+    COL_DIM,
+    COL_DISTANCE,
+    COL_LOWER_BOUNDS,
+    COL_MIN_LOWER_BOUND,
+    COL_OFFSET,
+    COL_QUERY_ID,
+    COL_RADIUS,
+    COL_SLICE,
+    COL_START,
+    COL_STOP,
+    COL_URL,
+)
+
+INDEX_SLICE_TYPE: pa.ListType = pa.list_(
+    pa.struct(
+        [
+            pa.field(COL_DIM, pa.string()),
+            pa.field(COL_START, pa.int64()),
+            pa.field(COL_STOP, pa.int64()),
+        ]
+    )
+)
+
+CHUNK_SCHEMA: pa.Schema = pa.schema(
+    [
+        pa.field(COL_URL, pa.string()),
+        pa.field(COL_SLICE, INDEX_SLICE_TYPE),
+    ]
+)
+
+MANIFEST_SCHEMA: pa.Schema = pa.schema(
+    [
+        pa.field(COL_URL, pa.string()),
+        pa.field(COL_SLICE, INDEX_SLICE_TYPE),
+        pa.field(COL_COUNT, pa.int32()),
+        pa.field(COL_CENTROID, pa.list_(pa.float32())),
+        pa.field(COL_RADIUS, pa.float32()),
+    ]
+)
+
+LOWER_BOUND_SCHEMA: pa.Schema = MANIFEST_SCHEMA.append(
+    pa.field(COL_LOWER_BOUNDS, pa.list_(pa.float32()))
+).append(pa.field(COL_MIN_LOWER_BOUND, pa.float32()))
+
+SEARCH_RESULT_SCHEMA: pa.Schema = pa.schema(
+    [
+        pa.field(COL_QUERY_ID, pa.int64()),
+        pa.field(COL_CHUNK_ID, pa.string()),
+        pa.field(COL_URL, pa.string()),
+        pa.field(COL_SLICE, INDEX_SLICE_TYPE),
+        pa.field(COL_OFFSET, pa.int64()),
+        pa.field(COL_DISTANCE, pa.float32()),
+    ]
+)
+
+BLOCK_SEARCH_SUMMARY_SCHEMA: pa.Schema = pa.schema(
+    [
+        pa.field("rows_seen", pa.int64()),
+        pa.field("rows_searched", pa.int64()),
+        pa.field("query_evaluations", pa.int64()),
+        pa.field("results_added", pa.int64()),
+    ]
+)
+
+ChunkTable: TypeAlias = Annotated[pa.Table, CHUNK_SCHEMA]
+ManifestTable: TypeAlias = Annotated[pa.Table, MANIFEST_SCHEMA]
+LowerBoundTable: TypeAlias = Annotated[pa.Table, LOWER_BOUND_SCHEMA]
+SearchResultTable: TypeAlias = Annotated[pa.Table, SEARCH_RESULT_SCHEMA]
+BlockSearchSummaryTable: TypeAlias = Annotated[pa.Table, BLOCK_SEARCH_SUMMARY_SCHEMA]

rayzin/manifest/spatial.py

@@ -0,0 +1,137 @@
+from collections.abc import Sequence
+from typing import Any
+
+import zarr
+from shapely.geometry.base import BaseGeometry  # type: ignore[import-untyped]
+
+from rayzin.types import (
+    COL_DIM,
+    COL_SLICE,
+    COL_START,
+    COL_STOP,
+    COL_URL,
+    ChunkRef,
+    DimSlice,
+)
+
+GeoTransform = tuple[float, float, float, float, float, float]
+
+
+def read_zarr_geotransform(
+    store_url: str,
+    *,
+    store_kwargs: dict[str, Any] | None = None,
+) -> GeoTransform:
+    # TODO: also support spatial zarr conventions.
+    group = zarr.open_group(
+        store=store_url,
+        mode="r",
+        storage_options=store_kwargs or None,
+    )
+    if "spatial_ref" not in group:
+        msg = f"Zarr store {store_url} is missing spatial_ref."
+        raise ValueError(msg)
+    spatial_ref = group["spatial_ref"]
+
+    raw_transform = spatial_ref.attrs.get("GeoTransform")
+    if raw_transform is None:
+        msg = f"Zarr store {store_url} is missing spatial_ref.GeoTransform."
+        raise ValueError(msg)
+
+    return parse_geotransform(raw_transform)
+
+
+def parse_geotransform(raw_transform: Any) -> GeoTransform:
+    parts: Sequence[Any]
+    if isinstance(raw_transform, str):
+        parts = raw_transform.split()
+    elif isinstance(raw_transform, Sequence):
+        parts = raw_transform
+    else:
+        msg = f"Unsupported GeoTransform value: {raw_transform!r}"
+        raise TypeError(msg)
+
+    if len(parts) != 6:
+        msg = f"Expected 6 GeoTransform coefficients, got {len(parts)}."
+        raise ValueError(msg)
+
+    origin_x = _coerce_float(parts[0])
+    pixel_width = _coerce_float(parts[1])
+    row_rotation = _coerce_float(parts[2])
+    origin_y = _coerce_float(parts[3])
+    column_rotation = _coerce_float(parts[4])
+    pixel_height = _coerce_float(parts[5])
+    return origin_x, pixel_width, row_rotation, origin_y, column_rotation, pixel_height
+
+
+def chunk_polygon(chunk: ChunkRef, transform: GeoTransform) -> BaseGeometry:
+    from shapely.geometry import Polygon  # type: ignore[import-untyped]
+
+    x_start, x_stop = _dim_interval(chunk["slice"], "x")
+    y_start, y_stop = _dim_interval(chunk["slice"], "y")
+    x0, y0 = pixel_to_world(transform, x_start, y_start)
+    x1, y1 = pixel_to_world(transform, x_stop, y_start)
+    x2, y2 = pixel_to_world(transform, x_stop, y_stop)
+    x3, y3 = pixel_to_world(transform, x_start, y_stop)
+    return Polygon([(x0, y0), (x1, y1), (x2, y2), (x3, y3)])
+
+
+def pixel_to_world(transform: GeoTransform, x: int, y: int) -> tuple[float, float]:
+    origin_x, pixel_width, row_rotation, origin_y, column_rotation, pixel_height = transform
+    world_x = origin_x + (x * pixel_width) + (y * row_rotation)
+    world_y = origin_y + (x * column_rotation) + (y * pixel_height)
+    return world_x, world_y
+
+
+def validate_aoi_geometry(aoi: BaseGeometry) -> BaseGeometry:
+    if isinstance(aoi, BaseGeometry):
+        if aoi.is_empty:
+            msg = "AOI geometry is empty."
+            raise ValueError(msg)
+        return aoi
+
+    msg = "AOI must be a shapely BaseGeometry. " f"Got {type(aoi).__name__}."
+    raise TypeError(msg)
+
+
+def chunk_from_row(row: dict[str, Any]) -> ChunkRef:
+    return ChunkRef(
+        url=str(row[COL_URL]),
+        slice=_coerce_slice(row[COL_SLICE]),
+    )
+
+
+def _dim_interval(parts: list[DimSlice], dim: str) -> tuple[int, int]:
+    for part in parts:
+        if part[COL_DIM] == dim:
+            return int(part[COL_START]), int(part[COL_STOP])
+    msg = f"Chunk is missing the required {dim!r} dimension."
+    raise KeyError(msg)
+
+
+def _coerce_slice(raw: Any) -> list[DimSlice]:
+    if not isinstance(raw, list):
+        msg = f"Expected slice to be a list, got {type(raw).__name__}."
+        raise TypeError(msg)
+
+    parts: list[DimSlice] = []
+    for part in raw:
+        if not isinstance(part, dict):
+            msg = f"Expected each slice entry to be a dict, got {type(part).__name__}."
+            raise TypeError(msg)
+        parts.append(
+            DimSlice(
+                dim=str(part[COL_DIM]),
+                start=_coerce_int(part[COL_START]),
+                stop=_coerce_int(part[COL_STOP]),
+            )
+        )
+    return parts
+
+
+def _coerce_float(value: Any) -> float:
+    return float(str(value))
+
+
+def _coerce_int(value: Any) -> int:
+    return int(str(value))

rayzin/metrics.py

@@ -0,0 +1,156 @@
+from typing import Protocol, runtime_checkable
+
+import numpy as np
+import pyarrow as pa  # type: ignore[import-untyped]
+
+from rayzin.enums import MetricType
+from rayzin.manifest.schema import LOWER_BOUND_SCHEMA, LowerBoundTable, ManifestTable
+from rayzin.types import (
+    COL_CENTROID,
+    COL_COUNT,
+    COL_RADIUS,
+    COL_SLICE,
+    COL_URL,
+    Float32Array,
+)
+
+
+@runtime_checkable
+class Metric(Protocol):
+    def distance(self, a: Float32Array, b: Float32Array) -> float: ...
+    def pairwise(self, vectors: Float32Array, query: Float32Array) -> Float32Array: ...
+    def lower_bound(self, query: Float32Array, centroid: Float32Array, radius: float) -> float: ...
+
+
+class EuclideanMetric:
+    """Squared L2 distance, matching FAISS IndexFlatL2 semantics."""
+
+    def distance(self, a: Float32Array, b: Float32Array) -> float:
+        delta = np.asarray(a - b, dtype=np.float32)
+        return float(np.dot(delta, delta))
+
+    def pairwise(self, vectors: Float32Array, query: Float32Array) -> Float32Array:
+        query_array = np.asarray(query, dtype=np.float32)
+        if query_array.ndim == 1:
+            deltas = np.asarray(vectors - query_array[None, :], dtype=np.float32)
+            return np.asarray(np.sum(deltas * deltas, axis=1, dtype=np.float32), dtype=np.float32)
+
+        deltas = np.asarray(vectors[None, :, :] - query_array[:, None, :], dtype=np.float32)
+        return np.asarray(np.sum(deltas * deltas, axis=2, dtype=np.float32), dtype=np.float32)
+
+    def lower_bound(self, query: Float32Array, centroid: Float32Array, radius: float) -> float:
+        centroid_distance = float(np.linalg.norm(query - centroid))
+        return float(max(0.0, centroid_distance - radius) ** 2)
+
+
+class CosineMetric:
+    """Cosine distance implemented as 1 - inner_product(normalized(a), normalized(b))."""
+
+    def distance(self, a: Float32Array, b: Float32Array) -> float:
+        return float(1.0 - np.dot(_normalize(a), _normalize(b)))
+
+    def pairwise(self, vectors: Float32Array, query: Float32Array) -> Float32Array:
+        normalized_vectors = _normalize(vectors)
+        normalized_query = _normalize(query)
+        if normalized_query.ndim == 1:
+            scores = normalized_vectors @ normalized_query
+            return np.asarray(1.0 - scores, dtype=np.float32)
+
+        scores = normalized_query @ normalized_vectors.T
+        return np.asarray(1.0 - scores, dtype=np.float32)
+
+    def lower_bound(self, query: Float32Array, centroid: Float32Array, radius: float) -> float:
+        return max(0.0, self.distance(query, centroid) - radius)
+
+
+EUCLIDEAN = EuclideanMetric()
+COSINE = CosineMetric()
+
+
+def make_metric(metric_type: MetricType) -> Metric:
+    if metric_type == MetricType.EUCLIDEAN:
+        return EuclideanMetric()
+    if metric_type == MetricType.COSINE:
+        return CosineMetric()
+    raise ValueError(metric_type)
+
+
+def add_lower_bounds_fn(
+    batch: ManifestTable,
+    queries: Float32Array,
+    metric_type: str,
+) -> LowerBoundTable:
+    return _add_lower_bounds(batch, queries=queries, metric_type=metric_type)
+
+
+def _add_lower_bounds(
+    batch: ManifestTable,
+    queries: Float32Array,
+    metric_type: str,
+) -> LowerBoundTable:
+    query_matrix = np.asarray(queries, dtype=np.float32)
+    if query_matrix.ndim != 2:
+        msg = f"Expected queries to have shape (nq, d), got {query_matrix.shape!r}."
+        raise ValueError(msg)
+
+    centroids = np.asarray(batch.column(COL_CENTROID).to_pylist(), dtype=np.float32)
+    radii = np.asarray(batch.column(COL_RADIUS).to_pylist(), dtype=np.float32)
+    lower_bounds = _lower_bounds(
+        query_matrix,
+        centroids,
+        radii,
+        metric_type=MetricType(metric_type),
+    )
+    min_lower_bounds = np.asarray(np.min(lower_bounds, axis=1), dtype=np.float32)
+
+    return pa.Table.from_arrays(
+        [
+            batch.column(COL_URL),
+            batch.column(COL_SLICE),
+            batch.column(COL_COUNT),
+            batch.column(COL_CENTROID),
+            batch.column(COL_RADIUS),
+            pa.array(lower_bounds.tolist(), type=pa.list_(pa.float32())),
+            pa.array(min_lower_bounds.tolist(), type=pa.float32()),
+        ],
+        schema=LOWER_BOUND_SCHEMA,
+    )
+
+
+def _lower_bounds(
+    queries: Float32Array,
+    centroids: Float32Array,
+    radii: Float32Array,
+    *,
+    metric_type: MetricType,
+) -> Float32Array:
+    if metric_type == MetricType.EUCLIDEAN:
+        deltas = np.asarray(
+            centroids[:, None, :] - queries[None, :, :],
+            dtype=np.float32,
+        )
+        centroid_distances = np.linalg.norm(deltas, axis=2)
+        return np.asarray(
+            np.square(np.maximum(0.0, centroid_distances - radii[:, None])),
+            dtype=np.float32,
+        )
+
+    metric = make_metric(metric_type)
+    distances = np.asarray(metric.pairwise(centroids, queries), dtype=np.float32).T
+    if distances.ndim != 2:
+        msg = f"Expected pairwise distances to have shape (n_rows, nq), got {distances.shape!r}."
+        raise ValueError(msg)
+    return np.asarray(np.maximum(0.0, distances - radii[:, None]), dtype=np.float32)
+
+
+def _normalize(vectors: Float32Array) -> Float32Array:
+    array = np.asarray(vectors, dtype=np.float32)
+    if array.ndim == 1:
+        denominator = np.float32(np.linalg.norm(array))
+        if denominator < np.finfo(np.float32).eps:
+            denominator = np.float32(np.finfo(np.float32).eps)
+        return np.asarray(array / denominator, dtype=np.float32)
+
+    norms = np.linalg.norm(array, axis=1, keepdims=True)
+    safe_norms = np.maximum(norms, np.finfo(np.float32).eps)
+    return np.asarray(array / safe_norms, dtype=np.float32)