diff --git a/README.md b/README.md index 4804393..540f818 100644 --- a/README.md +++ b/README.md @@ -224,6 +224,60 @@ The writer preserves source pixel dtype by default. To normalize stored pixel buffers explicitly, pass `pixel_dtype`, for example `pixel_dtype="uint16"`. Integer casts clamp by default; pass `clamp=False` to use NumPy casting behavior directly. +## Shape tables + +OME-Zarr and OME-NGFF are strong fits for dense image pixels and label rasters. +The gap is object-level analytics: cells, nuclei, ROIs, detections, tracks, +measurements, and relationships are naturally queried as tables. + +OME-Arrow Shapes fills that gap without replacing OME-Zarr. +Each row is one biological object, geometry is a single logical column, dense +label masks remain canonical label images, and object rows reference labels with +`label_image_id` plus `label_value`. +Measurements remain ordinary Arrow columns for DuckDB, Polars, DataFusion, +PyArrow, and Parquet workflows. + +```python +from ome_arrow import make_relationship_table, make_shape_table + +shapes = make_shape_table( + [ + { + "object_id": "cell-1", + "image_id": "image-1", + "label_image_id": "labels-1", + "label_value": 7, + "geometry": [128.0, 256.0], + "centroid": [128.0, 256.0], + "class": "cell", + "area_um2": 84.2, + } + ], + geometry_encoding="geoarrow.point", + axes=("y", "x"), + units=("pixel", "pixel"), +) + +relationships = make_relationship_table( + [ + { + "parent_id": "cell-1", + "child_id": "nucleus-1", + "relationship_type": "contains", + } + ] +) +``` + +Measurements remain normal Arrow columns (`area_um2` above). +Label masks are represented by reference with `label_image_id` and +`label_value`, keeping the segmentation raster canonical instead of embedding +mask payloads in each object row. + +See the dedicated Shapes docs: [`docs/src/shapes.md`](docs/src/shapes.md). +For local Parquet read/write benchmarks, run: +`python benchmarks/benchmark_shapes_parquet.py --repeats 5 --warmup 1`. + ## Tensor ingest (PyTorch/JAX) You can ingest torch or JAX arrays directly with `OMEArrow(...)`. diff --git a/benchmarks/benchmark_shapes_parquet.py b/benchmarks/benchmark_shapes_parquet.py new file mode 100644 index 0000000..47a908e --- /dev/null +++ b/benchmarks/benchmark_shapes_parquet.py @@ -0,0 +1,320 @@ +"""Benchmark OME-Arrow Shapes Parquet read/write paths. + +The benchmark uses synthetic but biologically interpretable object tables: + +- one row is one segmented object +- each object belongs to one source image +- label-reference rows point at canonical label rasters by label value +- measurement columns model common morphology and intensity features + +The results are directional local signals, not universal format rankings. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import tempfile +import time +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Callable + +import pyarrow as pa + +from ome_arrow import make_shape_table, read_shape_parquet, write_shape_parquet + + +@dataclass(frozen=True) +class ShapeFixture: + """Synthetic scientific shape workload description.""" + + name: str + geometry_encoding: str + image_count: int + objects_per_image: int + measurement_count: int + + @property + def object_count(self) -> int: + """Return the total number of object rows.""" + return self.image_count * self.objects_per_image + + +@dataclass(frozen=True) +class ShapeBenchmarkResult: + """One shape Parquet benchmark result.""" + + fixture: str + geometry_encoding: str + compression: str | None + image_count: int + object_count: int + measurement_count: int + write_ms: float + full_read_ms: float + projected_read_ms: float + filtered_read_ms: float + file_size_mb: float + projected_columns: int + filtered_rows: int + + +def _time(fn: Callable[[], object], *, repeats: int, warmup: int) -> float: + """Return median runtime in milliseconds.""" + for _ in range(warmup): + fn() + times_ms = [] + for _ in range(repeats): + start = time.perf_counter() + fn() + times_ms.append((time.perf_counter() - start) * 1000.0) + return statistics.median(times_ms) + + +def _measurement_payload(index: int) -> dict[str, float]: + """Create deterministic morphology and intensity measurements.""" + return { + "area_um2": float(50.0 + (index % 200) * 0.5), + "perimeter_um": float(20.0 + (index % 120) * 0.25), + "mean_intensity_dna": float(100.0 + (index % 4096)), + "mean_intensity_mito": float(80.0 + ((index * 3) % 2048)), + "eccentricity": float((index % 100) / 100.0), + "solidity": float(0.75 + (index % 25) / 100.0), + } + + +def _point_rows(fixture: ShapeFixture) -> list[dict[str, object]]: + """Build centroid-like point geometry rows.""" + rows: list[dict[str, object]] = [] + for image_idx in range(fixture.image_count): + image_id = f"image-{image_idx:03d}" + label_image_id = f"{image_id}-labels" + for object_idx in range(fixture.objects_per_image): + index = image_idx * fixture.objects_per_image + object_idx + y = float((object_idx // 100) * 12 + 6) + x = float((object_idx % 100) * 12 + 6) + rows.append( + { + "object_id": f"{image_id}-cell-{object_idx:05d}", + "image_id": image_id, + "label_image_id": label_image_id, + "label_value": object_idx + 1, + "geometry": [y, x], + "centroid": [y, x], + "class": "cell", + "confidence": 0.95, + **_measurement_payload(index), + } + ) + return rows + + +def _labelmask_rows(fixture: ShapeFixture) -> list[dict[str, object]]: + """Build rows that reference canonical label rasters.""" + rows: list[dict[str, object]] = [] + for image_idx in range(fixture.image_count): + image_id = f"image-{image_idx:03d}" + label_image_id = f"{image_id}-nuclear-labels" + for object_idx in range(fixture.objects_per_image): + index = image_idx * fixture.objects_per_image + object_idx + label_value = object_idx + 1 + rows.append( + { + "object_id": f"{image_id}-nucleus-{object_idx:05d}", + "image_id": image_id, + "label_image_id": label_image_id, + "label_value": label_value, + "geometry": { + "label_image_id": label_image_id, + "label_value": label_value, + }, + "class": "nucleus", + "confidence": 0.98, + **_measurement_payload(index), + } + ) + return rows + + +def _shape_table(fixture: ShapeFixture) -> pa.Table: + """Create one benchmark shape table.""" + if fixture.geometry_encoding == "ome.labelmask": + rows = _labelmask_rows(fixture) + else: + rows = _point_rows(fixture) + return make_shape_table( + rows, + geometry_encoding=fixture.geometry_encoding, + axes=("y", "x"), + units=("pixel", "pixel"), + ) + + +def _benchmark_fixture( + fixture: ShapeFixture, + *, + compression: str | None, + repeats: int, + warmup: int, + workdir: Path, +) -> ShapeBenchmarkResult: + """Run write/read benchmark cases for one shape fixture.""" + table = _shape_table(fixture) + out = workdir / f"{fixture.name}.{compression or 'none'}.ome-shapes.parquet" + projected_columns = [ + "object_id", + "image_id", + "label_value", + "area_um2", + "mean_intensity_dna", + "eccentricity", + ] + filtered_image_id = f"image-{fixture.image_count // 2:03d}" + + def write() -> None: + write_shape_parquet( + table, + out, + compression=compression, + row_group_size=fixture.objects_per_image, + use_dictionary=[ + "object_id", + "image_id", + "label_image_id", + "class", + ], + ) + + write() + write_ms = _time(write, repeats=repeats, warmup=warmup) + + def full_read() -> pa.Table: + return read_shape_parquet(out) + + def projected_read() -> pa.Table: + return read_shape_parquet(out, columns=projected_columns) + + def filtered_read() -> pa.Table: + return read_shape_parquet( + out, + columns=projected_columns, + filters=[("image_id", "==", filtered_image_id)], + ) + + full = full_read() + projected = projected_read() + filtered = filtered_read() + if full.num_rows != fixture.object_count: + raise AssertionError(f"full read returned {full.num_rows} rows") + if projected.num_columns != len(projected_columns): + raise AssertionError("projected read returned unexpected columns") + if filtered.num_rows != fixture.objects_per_image: + raise AssertionError(f"filtered read returned {filtered.num_rows} rows") + + return ShapeBenchmarkResult( + fixture=fixture.name, + geometry_encoding=fixture.geometry_encoding, + compression=compression, + image_count=fixture.image_count, + object_count=fixture.object_count, + measurement_count=fixture.measurement_count, + write_ms=write_ms, + full_read_ms=_time(full_read, repeats=repeats, warmup=warmup), + projected_read_ms=_time(projected_read, repeats=repeats, warmup=warmup), + filtered_read_ms=_time(filtered_read, repeats=repeats, warmup=warmup), + file_size_mb=out.stat().st_size / (1024 * 1024), + projected_columns=len(projected_columns), + filtered_rows=filtered.num_rows, + ) + + +def run( + *, + repeats: int, + warmup: int, + image_count: int, + objects_per_image: int, +) -> list[ShapeBenchmarkResult]: + """Run shape Parquet benchmarks.""" + fixtures = [ + ShapeFixture( + name="cell-centroids", + geometry_encoding="geoarrow.point", + image_count=image_count, + objects_per_image=objects_per_image, + measurement_count=6, + ), + ShapeFixture( + name="nucleus-label-refs", + geometry_encoding="ome.labelmask", + image_count=image_count, + objects_per_image=objects_per_image, + measurement_count=6, + ), + ] + results: list[ShapeBenchmarkResult] = [] + with tempfile.TemporaryDirectory(prefix="ome_arrow_shapes_parquet_") as tmp: + workdir = Path(tmp) + for fixture in fixtures: + for compression in (None, "zstd"): + results.append( + _benchmark_fixture( + fixture, + compression=compression, + repeats=repeats, + warmup=warmup, + workdir=workdir, + ) + ) + return results + + +def _print_results(results: list[ShapeBenchmarkResult]) -> None: + """Print benchmark results as a compact scientific table.""" + print("") + print("OME-Arrow Shapes Parquet benchmark") + print( + f"{'fixture':20} {'geometry':16} {'codec':8} " + f"{'images':>6} {'objects':>8} {'meas':>5} " + f"{'write':>9} {'full':>9} {'project':>9} {'filter':>9} " + f"{'size MB':>9}" + ) + print("-" * 118) + for result in results: + codec = result.compression or "none" + print( + f"{result.fixture:20} {result.geometry_encoding:16} {codec:8} " + f"{result.image_count:6d} {result.object_count:8d} " + f"{result.measurement_count:5d} {result.write_ms:9.2f} " + f"{result.full_read_ms:9.2f} {result.projected_read_ms:9.2f} " + f"{result.filtered_read_ms:9.2f} {result.file_size_mb:9.2f}" + ) + + +def main() -> None: + """Run the command-line benchmark.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--image-count", type=int, default=24) + parser.add_argument("--objects-per-image", type=int, default=2_000) + parser.add_argument("--json-out", type=Path, default=None) + args = parser.parse_args() + + results = run( + repeats=args.repeats, + warmup=args.warmup, + image_count=args.image_count, + objects_per_image=args.objects_per_image, + ) + _print_results(results) + if args.json_out is not None: + args.json_out.write_text( + json.dumps({"results": [asdict(result) for result in results]}, indent=2) + ) + + +if __name__ == "__main__": + main() diff --git a/docs/src/_static/generate_schema_image.py b/docs/src/_static/generate_schema_image.py new file mode 100644 index 0000000..f9252f5 --- /dev/null +++ b/docs/src/_static/generate_schema_image.py @@ -0,0 +1,334 @@ +"""Generate the README schema example image.""" + +from __future__ import annotations + +import math +from pathlib import Path + +from PIL import Image, ImageDraw, ImageFont + +SCALE = 2 +OUT = Path(__file__).with_name("various_ome_arrow_schema.png") + + +def _font(size: int, *, bold: bool = False) -> ImageFont.ImageFont: + """Load a readable system font with a Pillow fallback.""" + candidates = ( + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "", + "/System/Library/Fonts/Supplemental/Arial.ttf", + "/Library/Fonts/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "", + "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ) + for candidate in candidates: + if candidate and Path(candidate).exists(): + return ImageFont.truetype(candidate, size * SCALE) + return ImageFont.load_default(size * SCALE) + + +def _text_center( + draw: ImageDraw.ImageDraw, + box: tuple[int, int, int, int], + text: str, + font: ImageFont.ImageFont, + *, + fill: str = "#2b2f33", +) -> None: + """Draw centered text inside a box.""" + x0, y0, x1, y1 = (v * SCALE for v in box) + bbox = draw.multiline_textbbox((0, 0), text, font=font, spacing=7 * SCALE) + tw = bbox[2] - bbox[0] + th = bbox[3] - bbox[1] + draw.multiline_text( + (x0 + (x1 - x0 - tw) / 2, y0 + (y1 - y0 - th) / 2), + text, + font=font, + fill=fill, + align="center", + spacing=7 * SCALE, + ) + + +def _rounded_rect( + draw: ImageDraw.ImageDraw, + box: tuple[int, int, int, int], + *, + fill: str, + outline: str, + width: int = 2, + radius: int = 8, +) -> None: + """Draw a scaled rounded rectangle.""" + draw.rounded_rectangle( + tuple(v * SCALE for v in box), + radius=radius * SCALE, + fill=fill, + outline=outline, + width=width * SCALE, + ) + + +def _table( + draw: ImageDraw.ImageDraw, + x: int, + y: int, + col_widths: tuple[int, ...], + rows: int, + *, + headers: tuple[str, ...], + row_h: int = 76, + header_h: int = 48, +) -> list[list[tuple[int, int, int, int]]]: + """Draw a light table and return body cell boxes.""" + border = "#cfd5db" + total_w = sum(col_widths) + total_h = header_h + rows * row_h + draw.rectangle( + (x * SCALE, y * SCALE, (x + total_w) * SCALE, (y + total_h) * SCALE), + fill="#ffffff", + outline=border, + width=6 * SCALE, + ) + + cx = x + header_font = _font(17) + for idx, width in enumerate(col_widths): + if idx: + draw.line( + ( + cx * SCALE, + y * SCALE, + cx * SCALE, + (y + total_h) * SCALE, + ), + fill=border, + width=6 * SCALE, + ) + _text_center(draw, (cx, y, cx + width, y + header_h), headers[idx], header_font) + cx += width + + draw.line( + ( + x * SCALE, + (y + header_h) * SCALE, + (x + total_w) * SCALE, + (y + header_h) * SCALE, + ), + fill=border, + width=6 * SCALE, + ) + for row in range(1, rows): + yy = y + header_h + row * row_h + draw.line( + (x * SCALE, yy * SCALE, (x + total_w) * SCALE, yy * SCALE), + fill=border, + width=6 * SCALE, + ) + + cells: list[list[tuple[int, int, int, int]]] = [] + for row in range(rows): + row_cells = [] + cx = x + for width in col_widths: + row_cells.append( + ( + cx, + y + header_h + row * row_h, + cx + width, + y + header_h + (row + 1) * row_h, + ) + ) + cx += width + cells.append(row_cells) + return cells + + +def _draw_measurement( + draw: ImageDraw.ImageDraw, + box: tuple[int, int, int, int], +) -> None: + """Draw a compact measurement icon.""" + x0, y0, x1, y1 = box + cx = (x0 + x1) // 2 + cy = (y0 + y1) // 2 + _rounded_rect( + draw, + (cx - 20, cy - 18, cx + 20, cy + 18), + fill="#7ea2c5", + outline="#607f9f", + radius=7, + ) + font = _font(15, bold=True) + _text_center( + draw, + (cx - 20, cy - 19, cx + 20, cy + 18), + "12\n34", + font, + fill="#fff", + ) + + +def _draw_image_tile( + draw: ImageDraw.ImageDraw, + box: tuple[int, int, int, int], +) -> None: + """Draw a small microscopy image tile.""" + x0, y0, x1, y1 = box + cx = (x0 + x1) // 2 + cy = (y0 + y1) // 2 + tile = (cx - 23, cy - 23, cx + 23, cy + 23) + draw.rectangle( + tuple(v * SCALE for v in tile), + fill="#f6c23e", + outline="#d98a1a", + width=3 * SCALE, + ) + draw.rectangle( + tuple((v + 5 if i < 2 else v - 5) * SCALE for i, v in enumerate(tile)), + fill="#fdf5c9", + outline="#f0ad2e", + width=2 * SCALE, + ) + points = [] + for step in range(26): + t = step / 25 + px = cx - 14 + int(t * 28) + py = cy + int(math.sin(t * math.tau * 1.5) * 9) + points.append((px * SCALE, py * SCALE)) + draw.line(points, fill="#1fba6d", width=4 * SCALE) + for px, py in ((cx - 9, cy - 5), (cx + 1, cy + 5), (cx + 10, cy - 3)): + draw.ellipse( + ((px - 2) * SCALE, (py - 2) * SCALE, (px + 2) * SCALE, (py + 2) * SCALE), + fill="#229e5a", + ) + + +def _draw_label_tile( + draw: ImageDraw.ImageDraw, + box: tuple[int, int, int, int], +) -> None: + """Draw a compact label raster icon.""" + x0, y0, x1, y1 = box + cx = (x0 + x1) // 2 + cy = (y0 + y1) // 2 + tile = (cx - 23, cy - 23, cx + 23, cy + 23) + draw.rectangle( + tuple(v * SCALE for v in tile), + fill="#f6c23e", + outline="#d98a1a", + width=3 * SCALE, + ) + draw.rectangle( + tuple((v + 5 if i < 2 else v - 5) * SCALE for i, v in enumerate(tile)), + fill="#fdf5c9", + outline="#f0ad2e", + width=2 * SCALE, + ) + for dx, dy, color in ((-8, -4, "#56c2a4"), (4, 8, "#7bd389"), (9, -7, "#a5d76e")): + draw.ellipse( + ( + (cx + dx - 10) * SCALE, + (cy + dy - 8) * SCALE, + (cx + dx + 10) * SCALE, + (cy + dy + 8) * SCALE, + ), + fill=color, + outline="#258a5a", + width=2 * SCALE, + ) + + +def _draw_shape_icon( + draw: ImageDraw.ImageDraw, + box: tuple[int, int, int, int], +) -> None: + """Draw a vector shape icon.""" + x0, y0, x1, y1 = box + cx = (x0 + x1) // 2 + cy = (y0 + y1) // 2 + polygon = [ + ((cx - 20) * SCALE, (cy - 3) * SCALE), + ((cx - 8) * SCALE, (cy - 19) * SCALE), + ((cx + 16) * SCALE, (cy - 13) * SCALE), + ((cx + 21) * SCALE, (cy + 11) * SCALE), + ((cx - 4) * SCALE, (cy + 19) * SCALE), + ] + draw.polygon(polygon, fill="#d8efff", outline="#2c7fb8") + draw.line([*polygon, polygon[0]], fill="#2c7fb8", width=3 * SCALE) + for px, py in ( + (cx - 20, cy - 3), + (cx - 8, cy - 19), + (cx + 16, cy - 13), + (cx + 21, cy + 11), + (cx - 4, cy + 19), + ): + draw.ellipse( + ((px - 3) * SCALE, (py - 3) * SCALE, (px + 3) * SCALE, (py + 3) * SCALE), + fill="#0f5c99", + ) + + +def main() -> None: + """Generate the image asset.""" + width, height = 1000, 252 + image = Image.new("RGBA", (width * SCALE, height * SCALE), "#ffffff") + draw = ImageDraw.Draw(image) + + left = _table( + draw, + 0, + 6, + (132, 118), + 2, + headers=("Measurements", "Images"), + row_h=80, + header_h=50, + ) + _draw_measurement(draw, left[0][0]) + _draw_image_tile(draw, left[0][1]) + _draw_measurement(draw, left[1][0]) + _draw_image_tile(draw, left[1][1]) + + _text_center(draw, (270, 96, 312, 142), "or", _font(18)) + + center = _table( + draw, + 335, + 6, + (105, 105, 105), + 2, + headers=("Images", "Labels", "Shapes"), + row_h=80, + header_h=50, + ) + for row in center: + _draw_image_tile(draw, row[0]) + _draw_label_tile(draw, row[1]) + _draw_shape_icon(draw, row[2]) + + _text_center(draw, (670, 96, 712, 142), "or", _font(18)) + + draw.rectangle( + (745 * SCALE, 6 * SCALE, 995 * SCALE, 246 * SCALE), + fill="#ffffff", + outline="#cfd5db", + width=6 * SCALE, + ) + _text_center(draw, (768, 44, 972, 126), "Your\nschema\nhere!", _font(23)) + chip_font = _font(12, bold=True) + chips = [ + (770, 155, 828, 181, "images", "#dff2ff", "#2c7fb8"), + (836, 155, 892, 181, "labels", "#eaf8dc", "#4e9a33"), + (900, 155, 962, 181, "shapes", "#e8efff", "#4569b2"), + (804, 192, 930, 218, "measurements", "#edf1f5", "#61707f"), + ] + for x0, y0, x1, y1, label, fill, outline in chips: + _rounded_rect(draw, (x0, y0, x1, y1), fill=fill, outline=outline, radius=6) + _text_center(draw, (x0, y0, x1, y1), label, chip_font, fill="#25313b") + + image = image.resize((width, height), Image.Resampling.LANCZOS) + image.save(OUT) + + +if __name__ == "__main__": + main() diff --git a/docs/src/_static/various_ome_arrow_schema.png b/docs/src/_static/various_ome_arrow_schema.png index 22f1672..12eca96 100644 Binary files a/docs/src/_static/various_ome_arrow_schema.png and b/docs/src/_static/various_ome_arrow_schema.png differ diff --git a/docs/src/index.md b/docs/src/index.md index cf5e62b..e797295 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -14,6 +14,7 @@ maxdepth: 3 --- python-api dlpack +shapes why-ome-arrow presentations examples/learning_to_fly_with_ome-arrow diff --git a/docs/src/python-api.md b/docs/src/python-api.md index 8e110dd..c3f243f 100644 --- a/docs/src/python-api.md +++ b/docs/src/python-api.md @@ -7,6 +7,18 @@ ome_arrow :members: :undoc-members: :show-inheritance: + :exclude-members: OME_ARROW_RELATIONSHIPS_METADATA_KEY, + OME_ARROW_SHAPES_METADATA_KEY, + OMEArrowShapes, + geometry_storage_type, + make_relationship_table, + make_shape_table, + read_shape_parquet, + relationship_schema, + shape_schema, + validate_relationship_table, + validate_shape_table, + write_shape_parquet ``` ```{eval-rst} @@ -46,6 +58,15 @@ ome_arrow.meta :show-inheritance: ``` +```{eval-rst} +ome_arrow.shapes +------------------- +.. automodule:: ome_arrow.shapes + :members: + :undoc-members: + :show-inheritance: +``` + ```{eval-rst} ome_arrow.tensor ------------------- diff --git a/docs/src/shapes.md b/docs/src/shapes.md new file mode 100644 index 0000000..62cb08b --- /dev/null +++ b/docs/src/shapes.md @@ -0,0 +1,212 @@ +# OME-Arrow Shapes + +OME-Arrow Shapes is an Arrow-native convention for bioimaging object tables. +It is intended to complement OME-Zarr and OME-NGFF, not replace them. + +The core model is: + +- one row is one biological object +- geometry is one logical column +- measurements are ordinary Arrow columns +- dense label rasters remain canonical OME-Zarr labels +- shape rows reference labels instead of duplicating mask payloads +- coordinate metadata follows OME axis, unit, and coordinate-space concepts + +## The Gap + +OME has strong, mature conventions for image pixels and label images. +OME-Zarr and OME-NGFF are a natural fit for dense multiscale arrays, including +segmentation rasters. + +Object tables sit in a different access pattern. +Common bioimage analysis needs to query, filter, join, and aggregate objects: + +- load all cell centroids +- filter detections by class or confidence +- join cells to nuclei +- compute population summaries +- associate each object with a source image and label value +- use DuckDB, Polars, DataFusion, or PyArrow without unpacking image arrays + +Those workflows are columnar and relational. +They benefit from Arrow and Parquet, but still need OME-compatible semantics for +axes, units, labels, provenance, and relationships. + +OME-Arrow Shapes fills that table-shaped gap. + +## Alignment With OME + +OME-Arrow Shapes stays with OME's existing direction in several ways. + +Images stay dense. +Pixels should remain in OME-Zarr, OME-TIFF, or OME-Arrow image storage. + +Labels stay dense. +Segmentation masks should remain canonical label rasters. +Shape rows can reference label objects with `label_image_id` and `label_value`. + +Coordinates stay explicit. +Shape schemas store axes, units, coordinate space, and geometry encoding in Arrow +schema metadata. +This keeps table values interpretable without inventing a separate coordinate +system model. + +Measurements stay columnar. +Area, volume, intensity, texture, model scores, and morphology features are plain +Arrow columns. +No special nested measurement encoding is required. + +## Shape Schema + +Use `make_shape_table` to create a validated Arrow table. + +```python +from ome_arrow import make_shape_table + +shapes = make_shape_table( + [ + { + "object_id": "cell-1", + "image_id": "image-1", + "label_image_id": "labels-1", + "label_value": 7, + "geometry": [128.0, 256.0], + "centroid": [128.0, 256.0], + "class": "cell", + "confidence": 0.98, + "area_um2": 84.2, + } + ], + geometry_encoding="geoarrow.point", + axes=("y", "x"), + units=("pixel", "pixel"), + coordinate_space="pixel", +) +``` + +The canonical columns are: + +| Column | Purpose | +| ---------------- | --------------------------------------- | +| `object_id` | stable object identifier | +| `image_id` | source image identifier | +| `label_image_id` | source label image identifier | +| `label_value` | integer label value in the label raster | +| `geometry` | one logical geometry value | +| `centroid` | coordinate vector for object center | +| `bbox` | min/max coordinate bounds | +| `class` | object class or annotation category | +| `confidence` | detection or classification confidence | + +Any additional columns are measurements. +They remain ordinary Arrow columns and can be queried directly. + +## Geometry Encodings + +The current registry includes: + +| Encoding | Storage intent | +| ----------------------- | -------------------------- | +| `geoarrow.point` | point coordinate vector | +| `geoarrow.linestring` | list of coordinate vectors | +| `geoarrow.polygon` | list of rings | +| `geoarrow.multipolygon` | list of polygons | +| `ome.labelmask` | label image reference | +| `ome.pointcloud` | list of coordinate vectors | +| `ome.boundingbox` | min/max coordinate vectors | +| `ome.mesh3d` | vertices and faces | + +The registry lets 2D shapes reuse GeoArrow-like nested Arrow layouts while +leaving room for bioimaging-specific geometry such as label masks, point clouds, +and meshes. + +## Label References + +For segmentation objects, prefer references over embedded masks. + +```python +from ome_arrow import make_shape_table + +objects = make_shape_table( + [ + { + "object_id": "nucleus-1", + "image_id": "image-1", + "label_image_id": "nuclear-labels", + "label_value": 42, + "geometry": { + "label_image_id": "nuclear-labels", + "label_value": 42, + }, + "class": "nucleus", + } + ], + geometry_encoding="ome.labelmask", +) +``` + +This avoids duplicating raster masks in every object row. +The label image remains the canonical segmentation, and the object table remains +small, queryable, and provenance-preserving. + +## Relationships + +Object relationships are also ordinary Arrow rows. + +```python +from ome_arrow import make_relationship_table + +relationships = make_relationship_table( + [ + { + "parent_id": "cell-1", + "child_id": "nucleus-1", + "relationship_type": "contains", + "confidence": 1.0, + } + ] +) +``` + +Supported relationship types are: + +- `contains` +- `adjacent` +- `touches` +- `parent` +- `track` +- `derived_from` + +## Performance Contract + +Shape tables are built on Arrow-native arrays and schema metadata. +The implementation avoids binary geometry blobs where nested Arrow types are +practical, which keeps column projection, filtering, and Parquet serialization +available to Arrow-compatible engines. + +The test suite includes pytest performance canaries for shape table construction +and image-id filtering. +They are intentionally lightweight regression checks rather than rigorous +hardware-specific benchmarks. + +Run them with: + +```sh +python -m pytest tests/test_shapes.py tests/test_shapes_performance.py +``` + +For Parquet read/write performance, run: + +```sh +python benchmarks/benchmark_shapes_parquet.py --repeats 5 --warmup 1 +``` + +This benchmark reports scientifically interpretable workload dimensions: +image count, object count, measurement count, geometry encoding, projected +column count, filtered row count, file size, and read/write timings. +It includes both centroid-like `geoarrow.point` rows and `ome.labelmask` rows +that reference canonical label rasters. + +Projected and filtered reads exercise the performance path most analysis tools +need: load only the object IDs, image IDs, label values, and measurements needed +for a specific query. diff --git a/src/ome_arrow/__init__.py b/src/ome_arrow/__init__.py index bd256a0..e5acfbb 100644 --- a/src/ome_arrow/__init__.py +++ b/src/ome_arrow/__init__.py @@ -36,6 +36,20 @@ OME_ARROW_TAG_TYPE, OME_ARROW_TAG_VERSION, ) +from ome_arrow.shapes import ( + OME_ARROW_RELATIONSHIPS_METADATA_KEY, + OME_ARROW_SHAPES_METADATA_KEY, + OMEArrowShapes, + geometry_storage_type, + make_relationship_table, + make_shape_table, + read_shape_parquet, + relationship_schema, + shape_schema, + validate_relationship_table, + validate_shape_table, + write_shape_parquet, +) from ome_arrow.tensor import LazyTensorView, TensorView from ome_arrow.utils import describe_ome_arrow, verify_ome_arrow from ome_arrow.view import view_matplotlib, view_pyvista diff --git a/src/ome_arrow/shapes.py b/src/ome_arrow/shapes.py new file mode 100644 index 0000000..6704a0a --- /dev/null +++ b/src/ome_arrow/shapes.py @@ -0,0 +1,670 @@ +"""Arrow-native shape and relationship tables for OME-Arrow.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from os import PathLike +from typing import Any, Iterable, Literal, Sequence, get_args + +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.parquet as pq + +from ome_arrow.meta import OME_ARROW_TAG_VERSION + +GeometryEncoding = Literal[ + "geoarrow.point", + "geoarrow.linestring", + "geoarrow.polygon", + "geoarrow.multipolygon", + "ome.mesh3d", + "ome.labelmask", + "ome.pointcloud", + "ome.boundingbox", +] + +SUPPORTED_GEOMETRY_ENCODINGS: frozenset[str] = frozenset(get_args(GeometryEncoding)) + +OME_ARROW_SHAPES_METADATA_KEY = b"ome.arrow.shapes" +OME_ARROW_RELATIONSHIPS_METADATA_KEY = b"ome.arrow.relationships" + +DEFAULT_SHAPE_COLUMNS = { + "object_id", + "image_id", + "label_image_id", + "label_value", + "geometry", + "centroid", + "bbox", + "class", + "confidence", +} + +RELATIONSHIP_TYPES: frozenset[str] = frozenset( + { + "contains", + "adjacent", + "touches", + "parent", + "track", + "derived_from", + } +) + +RELATIONSHIP_SCHEMA = pa.schema( + [ + pa.field("parent_id", pa.string(), nullable=False), + pa.field("child_id", pa.string(), nullable=False), + pa.field("relationship_type", pa.string(), nullable=False), + pa.field("confidence", pa.float32()), + ] +) + + +def _coordinate_type(dimensions: int) -> pa.ListType: + """Return the Arrow coordinate vector type for a geometry dimension.""" + if dimensions < 1: + raise ValueError("geometry dimensions must be at least 1.") + return pa.list_(pa.float64()) + + +def geometry_storage_type( + geometry_encoding: GeometryEncoding | str, + *, + dimensions: int = 2, +) -> pa.DataType: + """Return the Arrow storage type for a registered geometry encoding. + + Args: + geometry_encoding: Registered OME-Arrow geometry encoding name. + dimensions: Number of coordinate dimensions for coordinate vectors. + + Returns: + Arrow data type for the logical geometry column. + + Raises: + ValueError: If the encoding is unknown or dimensions are invalid. + """ + if geometry_encoding not in SUPPORTED_GEOMETRY_ENCODINGS: + raise ValueError(f"Unsupported geometry_encoding: {geometry_encoding!r}.") + + coordinate = _coordinate_type(dimensions) + if geometry_encoding == "geoarrow.point": + storage_type = coordinate + elif geometry_encoding in {"geoarrow.linestring", "ome.pointcloud"}: + storage_type = pa.list_(coordinate) + elif geometry_encoding == "geoarrow.polygon": + storage_type = pa.list_(pa.list_(coordinate)) + elif geometry_encoding == "geoarrow.multipolygon": + storage_type = pa.list_(pa.list_(pa.list_(coordinate))) + elif geometry_encoding == "ome.boundingbox": + storage_type = pa.struct( + [ + pa.field("min", coordinate, nullable=False), + pa.field("max", coordinate, nullable=False), + ] + ) + elif geometry_encoding == "ome.labelmask": + storage_type = pa.struct( + [ + pa.field("label_image_id", pa.string(), nullable=False), + pa.field("label_value", pa.int64(), nullable=False), + ] + ) + else: + storage_type = pa.struct( + [ + pa.field("vertices", pa.list_(_coordinate_type(3)), nullable=False), + pa.field("faces", pa.list_(pa.list_(pa.int32())), nullable=False), + ] + ) + return storage_type + + +def shape_metadata( + *, + geometry_encoding: GeometryEncoding | str, + axes: Sequence[str] = ("y", "x"), + units: Sequence[str] | None = None, + coordinate_space: str = "pixel", + geometry_column: str = "geometry", +) -> dict[str, Any]: + """Build JSON-serializable schema metadata for a shape table. + + Args: + geometry_encoding: Registered OME-Arrow geometry encoding name. + axes: Coordinate axis names for geometry values. + units: Units aligned to axes. Defaults to ``"pixel"`` for each axis. + coordinate_space: Name of the coordinate space for geometry values. + geometry_column: Name of the logical geometry column. + + Returns: + Metadata dictionary stored under ``OME_ARROW_SHAPES_METADATA_KEY``. + + Raises: + ValueError: If the encoding, axes, or units are invalid. + """ + if geometry_encoding not in SUPPORTED_GEOMETRY_ENCODINGS: + raise ValueError(f"Unsupported geometry_encoding: {geometry_encoding!r}.") + if not axes: + raise ValueError("axes must contain at least one axis name.") + if units is None: + units = tuple("pixel" for _ in axes) + if len(units) != len(axes): + raise ValueError("units must have the same length as axes.") + + return { + "type": "ome.arrow.shapes", + "version": str(OME_ARROW_TAG_VERSION), + "geometry_column": geometry_column, + "geometry_encoding": str(geometry_encoding), + "axes": list(axes), + "units": list(units), + "coordinate_space": coordinate_space, + } + + +def _schema_with_json_metadata( + schema: pa.Schema, + *, + key: bytes, + payload: dict[str, Any], +) -> pa.Schema: + """Attach compact JSON metadata to an Arrow schema.""" + metadata = dict(schema.metadata or {}) + metadata[key] = json.dumps(payload, sort_keys=True).encode("utf-8") + return schema.with_metadata(metadata) + + +def shape_schema( + geometry_encoding: GeometryEncoding | str, + *, + axes: Sequence[str] = ("y", "x"), + units: Sequence[str] | None = None, + coordinate_space: str = "pixel", + geometry_column: str = "geometry", + measurement_fields: Iterable[pa.Field] | None = None, +) -> pa.Schema: + """Create an OME-Arrow shape table schema. + + Args: + geometry_encoding: Registered OME-Arrow geometry encoding name. + axes: Coordinate axis names for geometry, centroid, and bounding boxes. + units: Units aligned to axes. Defaults to ``"pixel"`` for each axis. + coordinate_space: Name of the coordinate space for geometry values. + geometry_column: Name of the logical geometry column. + measurement_fields: Extra Arrow fields for ordinary measurement columns. + + Returns: + Arrow schema with OME-Arrow shape metadata attached. + """ + metadata = shape_metadata( + geometry_encoding=geometry_encoding, + axes=axes, + units=units, + coordinate_space=coordinate_space, + geometry_column=geometry_column, + ) + coordinate = _coordinate_type(len(axes)) + fields = [ + pa.field("object_id", pa.string(), nullable=False), + pa.field("image_id", pa.string()), + pa.field("label_image_id", pa.string()), + pa.field("label_value", pa.int64()), + pa.field( + geometry_column, + geometry_storage_type(geometry_encoding, dimensions=len(axes)), + ), + pa.field("centroid", coordinate), + pa.field( + "bbox", + pa.struct( + [ + pa.field("min", coordinate, nullable=False), + pa.field("max", coordinate, nullable=False), + ] + ), + ), + pa.field("class", pa.string()), + pa.field("confidence", pa.float32()), + ] + if measurement_fields is not None: + fields.extend(measurement_fields) + + return _schema_with_json_metadata( + pa.schema(fields), + key=OME_ARROW_SHAPES_METADATA_KEY, + payload=metadata, + ) + + +def _infer_measurement_fields( + rows: Sequence[dict[str, Any]], + *, + geometry_column: str = "geometry", +) -> list[pa.Field]: + """Infer measurement fields for columns outside the canonical shape columns.""" + if not rows: + return [] + + fields: list[pa.Field] = [] + row_columns = set().union(*(row.keys() for row in rows)) + reserved_columns = DEFAULT_SHAPE_COLUMNS | {geometry_column} + for name in sorted(row_columns - reserved_columns): + values = [row.get(name) for row in rows] + fields.append(pa.field(name, pa.array(values).type)) + return fields + + +def make_shape_table( + rows: Sequence[dict[str, Any]], + *, + geometry_encoding: GeometryEncoding | str, + axes: Sequence[str] = ("y", "x"), + units: Sequence[str] | None = None, + coordinate_space: str = "pixel", + geometry_column: str = "geometry", + validate: bool = True, +) -> pa.Table: + """Create an OME-Arrow shape table from Python row dictionaries. + + Args: + rows: Shape rows, where each row represents one biological object. + geometry_encoding: Registered OME-Arrow geometry encoding name. + axes: Coordinate axis names for geometry, centroid, and bounding boxes. + units: Units aligned to axes. Defaults to ``"pixel"`` for each axis. + coordinate_space: Name of the coordinate space for geometry values. + geometry_column: Name of the logical geometry column. + validate: Validate the table after construction. + + Returns: + Arrow table with OME-Arrow shape schema metadata. + """ + row_list = list(rows) + schema = shape_schema( + geometry_encoding, + axes=axes, + units=units, + coordinate_space=coordinate_space, + geometry_column=geometry_column, + measurement_fields=_infer_measurement_fields( + row_list, + geometry_column=geometry_column, + ), + ) + table = pa.Table.from_pylist(row_list, schema=schema) + if validate: + validate_shape_table(table) + return table + + +def _shape_metadata_from_schema(schema: pa.Schema) -> dict[str, Any]: + """Read OME-Arrow shape JSON metadata from a schema.""" + raw_metadata = schema.metadata or {} + raw_payload = raw_metadata.get(OME_ARROW_SHAPES_METADATA_KEY) + if raw_payload is None: + raise ValueError("Shape table schema metadata is missing OME-Arrow shapes.") + metadata = json.loads(raw_payload.decode("utf-8")) + if metadata.get("type") != "ome.arrow.shapes": + raise ValueError("Shape table metadata type must be 'ome.arrow.shapes'.") + return metadata + + +def _is_coordinate(value: Any, dimensions: int) -> bool: + """Return whether a Python value is a coordinate of the expected arity.""" + return isinstance(value, list) and len(value) == dimensions + + +def _validate_coordinate(value: Any, dimensions: int, path: str) -> None: + """Validate one coordinate vector. + + Args: + value: Python value to validate. + dimensions: Expected coordinate length. + path: Human-readable location used in error messages. + + Raises: + ValueError: If the value is not a coordinate of the expected arity. + """ + if not _is_coordinate(value, dimensions): + raise ValueError(f"{path} must contain {dimensions} coordinates.") + + +def _validate_geometry_value( + value: Any, + *, + geometry_encoding: str, + dimensions: int, + row_index: int, +) -> None: + """Validate coordinate arity for one geometry value.""" + if value is None or geometry_encoding == "ome.labelmask": + return + if geometry_encoding == "geoarrow.point": + _validate_coordinate(value, dimensions, f"geometry row {row_index}") + elif geometry_encoding in {"geoarrow.linestring", "ome.pointcloud"}: + for point_index, point in enumerate(value): + _validate_coordinate( + point, + dimensions, + f"geometry row {row_index} point {point_index}", + ) + elif geometry_encoding == "geoarrow.polygon": + for ring_index, ring in enumerate(value): + for point_index, point in enumerate(ring): + _validate_coordinate( + point, + dimensions, + f"geometry row {row_index} ring {ring_index} point {point_index}", + ) + elif geometry_encoding == "geoarrow.multipolygon": + for polygon_index, polygon in enumerate(value): + for ring_index, ring in enumerate(polygon): + for point_index, point in enumerate(ring): + _validate_coordinate( + point, + dimensions, + "geometry row " + f"{row_index} polygon {polygon_index} " + f"ring {ring_index} point {point_index}", + ) + elif geometry_encoding == "ome.boundingbox": + _validate_coordinate(value.get("min"), dimensions, f"bbox row {row_index} min") + _validate_coordinate(value.get("max"), dimensions, f"bbox row {row_index} max") + elif geometry_encoding == "ome.mesh3d": + for vertex_index, vertex in enumerate(value.get("vertices", [])): + _validate_coordinate( + vertex, 3, f"mesh row {row_index} vertex {vertex_index}" + ) + + +def _validate_coordinate_columns(table: pa.Table, metadata: dict[str, Any]) -> None: + """Validate coordinate arity for geometry, centroid, and bounding boxes.""" + dimensions = len(metadata.get("axes", [])) + if dimensions < 1: + raise ValueError("Shape table axes metadata must contain at least one axis.") + + geometry_column = metadata.get("geometry_column", "geometry") + geometry_encoding = str(metadata.get("geometry_encoding")) + for row_index, value in enumerate(table[geometry_column].to_pylist()): + _validate_geometry_value( + value, + geometry_encoding=geometry_encoding, + dimensions=dimensions, + row_index=row_index, + ) + + if "centroid" in table.column_names: + for row_index, value in enumerate(table["centroid"].to_pylist()): + if value is not None: + _validate_coordinate(value, dimensions, f"centroid row {row_index}") + + if "bbox" in table.column_names: + for row_index, value in enumerate(table["bbox"].to_pylist()): + if value is not None: + _validate_coordinate( + value.get("min"), dimensions, f"bbox row {row_index} min" + ) + _validate_coordinate( + value.get("max"), dimensions, f"bbox row {row_index} max" + ) + + +def validate_shape_table(table: pa.Table) -> None: + """Validate an OME-Arrow shape table. + + Args: + table: Arrow table to validate. + + Raises: + ValueError: If required metadata, columns, encoding, or IDs are invalid. + """ + metadata = _shape_metadata_from_schema(table.schema) + geometry_column = metadata.get("geometry_column", "geometry") + geometry_encoding = metadata.get("geometry_encoding") + axes = metadata.get("axes", []) + + if geometry_encoding not in SUPPORTED_GEOMETRY_ENCODINGS: + raise ValueError(f"Unsupported geometry_encoding: {geometry_encoding!r}.") + if "object_id" not in table.column_names: + raise ValueError("Shape table must contain an object_id column.") + if geometry_column not in table.column_names: + raise ValueError(f"Shape table must contain {geometry_column!r} column.") + if table.schema.field(geometry_column).type != geometry_storage_type( + geometry_encoding, + dimensions=len(axes), + ): + raise ValueError("Shape table geometry column does not match metadata.") + _validate_coordinate_columns(table, metadata) + if pc.any(pc.is_null(table["object_id"])).as_py(): + raise ValueError("Shape table object_id values must not be null.") + + +def write_shape_parquet( + table: pa.Table, + path: str | PathLike[str], + *, + compression: str | None = "zstd", + row_group_size: int | None = 65_536, + use_dictionary: bool | list[str] = True, + validate: bool = True, +) -> None: + """Write an OME-Arrow shape table to Parquet. + + Args: + table: OME-Arrow shape table to write. + path: Output Parquet path. + compression: Parquet compression codec, or ``None`` for uncompressed. + row_group_size: Number of rows per Parquet row group. + use_dictionary: Dictionary-encode eligible columns. This is useful for + repeated scientific labels such as image IDs, label image IDs, and + object classes. + validate: Validate the table before writing. + + Raises: + ValueError: If validation fails. + """ + if validate: + validate_shape_table(table) + pq.write_table( + table, + path, + compression=compression, + row_group_size=row_group_size, + use_dictionary=use_dictionary, + ) + + +def read_shape_parquet( + path: str | PathLike[str], + *, + columns: Sequence[str] | None = None, + filters: Any | None = None, + memory_map: bool = True, + validate: bool = True, +) -> pa.Table: + """Read an OME-Arrow shape Parquet table. + + Args: + path: Input Parquet path. + columns: Optional column projection for analytical reads. + filters: Optional PyArrow Parquet filters for predicate pushdown. + memory_map: Use memory mapping where supported. + validate: Validate complete shape tables after reading. Projected reads + that omit required columns still validate schema metadata but skip + full table validation. + + Returns: + Arrow table read from Parquet. + + Raises: + ValueError: If OME-Arrow shape metadata or complete-table validation + fails. + """ + schema = pq.read_schema(path, memory_map=memory_map) + metadata = _shape_metadata_from_schema(schema) + table = pq.read_table( + path, + columns=columns, + filters=filters, + memory_map=memory_map, + ) + if not validate: + return table + + geometry_column = metadata.get("geometry_column", "geometry") + required = {"object_id", geometry_column} + if required.issubset(table.column_names): + validate_shape_table(table) + return table + + +def relationship_metadata() -> dict[str, Any]: + """Build JSON-serializable schema metadata for object relationships.""" + return { + "type": "ome.arrow.relationships", + "version": str(OME_ARROW_TAG_VERSION), + "relationship_types": sorted(RELATIONSHIP_TYPES), + } + + +def relationship_schema() -> pa.Schema: + """Create an OME-Arrow relationship table schema. + + Returns: + Arrow schema with OME-Arrow relationship metadata attached. + """ + return _schema_with_json_metadata( + RELATIONSHIP_SCHEMA, + key=OME_ARROW_RELATIONSHIPS_METADATA_KEY, + payload=relationship_metadata(), + ) + + +def make_relationship_table( + rows: Sequence[dict[str, Any]], + *, + validate: bool = True, +) -> pa.Table: + """Create an OME-Arrow relationship table from edge rows. + + Args: + rows: Relationship rows with parent, child, and relationship type. + validate: Validate the table after construction. + + Returns: + Arrow table with OME-Arrow relationship metadata. + """ + table = pa.Table.from_pylist(list(rows), schema=relationship_schema()) + if validate: + validate_relationship_table(table) + return table + + +def validate_relationship_table(table: pa.Table) -> None: + """Validate an OME-Arrow relationship table. + + Args: + table: Arrow table to validate. + + Raises: + ValueError: If required metadata, columns, IDs, or relationship types fail. + """ + raw_payload = (table.schema.metadata or {}).get( + OME_ARROW_RELATIONSHIPS_METADATA_KEY + ) + if raw_payload is None: + raise ValueError( + "Relationship table schema metadata is missing OME-Arrow relationships." + ) + metadata = json.loads(raw_payload.decode("utf-8")) + if metadata.get("type") != "ome.arrow.relationships": + raise ValueError( + "Relationship table metadata type must be 'ome.arrow.relationships'." + ) + + for name in ("parent_id", "child_id", "relationship_type"): + if name not in table.column_names: + raise ValueError(f"Relationship table must contain a {name} column.") + if pc.any(pc.is_null(table[name])).as_py(): + raise ValueError(f"Relationship table {name} values must not be null.") + + unknown = set(table["relationship_type"].to_pylist()) - RELATIONSHIP_TYPES + if unknown: + raise ValueError(f"Unsupported relationship_type values: {sorted(unknown)}.") + + +@dataclass(frozen=True) +class OMEArrowShapes: + """Convenience wrapper around a validated OME-Arrow shape table.""" + + table: pa.Table + + def __post_init__(self) -> None: + """Validate the wrapped shape table.""" + validate_shape_table(self.table) + + @classmethod + def from_rows( + cls, + rows: Sequence[dict[str, Any]], + *, + geometry_encoding: GeometryEncoding | str, + axes: Sequence[str] = ("y", "x"), + units: Sequence[str] | None = None, + coordinate_space: str = "pixel", + geometry_column: str = "geometry", + ) -> OMEArrowShapes: + """Create a shape wrapper from Python row dictionaries. + + Args: + rows: Shape rows, where each row represents one biological object. + geometry_encoding: Registered OME-Arrow geometry encoding name. + axes: Coordinate axis names for geometry, centroid, and bounding boxes. + units: Units aligned to axes. Defaults to ``"pixel"`` for each axis. + coordinate_space: Name of the coordinate space for geometry values. + geometry_column: Name of the logical geometry column. + + Returns: + Validated OME-Arrow shapes wrapper. + """ + return cls( + make_shape_table( + rows, + geometry_encoding=geometry_encoding, + axes=axes, + units=units, + coordinate_space=coordinate_space, + geometry_column=geometry_column, + ) + ) + + @property + def metadata(self) -> dict[str, Any]: + """Return decoded OME-Arrow shapes metadata.""" + return _shape_metadata_from_schema(self.table.schema) + + @property + def geometry_encoding(self) -> str: + """Return the registered geometry encoding for the table.""" + return str(self.metadata["geometry_encoding"]) + + @property + def axes(self) -> tuple[str, ...]: + """Return coordinate axis names for the shape table.""" + return tuple(self.metadata["axes"]) + + def for_image(self, image_id: str) -> OMEArrowShapes: + """Return shapes that reference one image ID. + + Args: + image_id: Image identifier to filter on. + + Returns: + New wrapper containing only matching shape rows. + """ + if "image_id" not in self.table.column_names: + return type(self)(self.table.slice(0, 0)) + mask = pc.equal(self.table["image_id"], image_id) + return type(self)(self.table.filter(mask)) diff --git a/tests/test_shapes.py b/tests/test_shapes.py new file mode 100644 index 0000000..a917686 --- /dev/null +++ b/tests/test_shapes.py @@ -0,0 +1,353 @@ +"""Tests for Arrow-native OME-Arrow shape tables.""" + +import json +import pathlib + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +from ome_arrow import ( + OME_ARROW_SHAPES_METADATA_KEY, + OMEArrowShapes, + make_relationship_table, + make_shape_table, + read_shape_parquet, + shape_schema, + validate_relationship_table, + validate_shape_table, + write_shape_parquet, +) +from ome_arrow.shapes import SUPPORTED_GEOMETRY_ENCODINGS + + +def test_make_shape_table_stores_single_geometry_column_with_metadata() -> None: + """Create a shape table with one logical geometry value per object row.""" + table = make_shape_table( + [ + { + "object_id": "cell-1", + "image_id": "image-1", + "label_image_id": "labels-1", + "label_value": 7, + "geometry": [[0.0, 0.0], [2.0, 0.0], [2.0, 3.0]], + "centroid": [1.25, 1.5], + "bbox": {"min": [0.0, 0.0], "max": [2.0, 3.0]}, + "class": "cell", + "confidence": 0.98, + "area": 5.5, + } + ], + geometry_encoding="geoarrow.linestring", + axes=("y", "x"), + units=("pixel", "pixel"), + coordinate_space="pixel", + ) + + metadata = json.loads(table.schema.metadata[OME_ARROW_SHAPES_METADATA_KEY]) + + assert table.num_rows == 1 + assert table.column_names.count("geometry") == 1 + assert table.schema.field("geometry").type == pa.list_(pa.list_(pa.float64())) + assert table.schema.field("area").type == pa.float64() + assert metadata["type"] == "ome.arrow.shapes" + assert metadata["geometry_encoding"] == "geoarrow.linestring" + assert metadata["axes"] == ["y", "x"] + assert metadata["coordinate_space"] == "pixel" + assert validate_shape_table(table) is None + + +def test_shape_table_supports_labelmask_geometry_references() -> None: + """Represent label masks by reference instead of embedding raster masks.""" + table = make_shape_table( + [ + { + "object_id": "nucleus-1", + "image_id": "image-1", + "label_image_id": "nuclear-labels", + "label_value": 42, + "geometry": { + "label_image_id": "nuclear-labels", + "label_value": 42, + }, + "class": "nucleus", + } + ], + geometry_encoding="ome.labelmask", + ) + + geometry = table["geometry"].to_pylist()[0] + + assert geometry == {"label_image_id": "nuclear-labels", "label_value": 42} + assert validate_shape_table(table) is None + + +def test_shape_table_rejects_invalid_geometry_encoding() -> None: + """Fail before creating tables for unknown geometry encodings.""" + assert "geoarrow.polygon" in SUPPORTED_GEOMETRY_ENCODINGS + + with pytest.raises(ValueError, match="Unsupported geometry_encoding"): + shape_schema("wkb") + + +def test_validate_shape_table_requires_metadata_and_identity() -> None: + """Require shape metadata and object identity columns.""" + table = pa.table({"geometry": [[[0.0, 1.0]]]}) + + with pytest.raises(ValueError, match="schema metadata"): + validate_shape_table(table) + + schema = shape_schema("geoarrow.point").remove(0) + table = pa.Table.from_pylist( + [{"geometry": [0.0, 1.0]}], + schema=schema, + ) + + with pytest.raises(ValueError, match="object_id"): + validate_shape_table(table) + + +def test_validate_shape_table_rejects_wrong_coordinate_arity() -> None: + """Reject geometry, centroid, and bbox coordinates that do not match axes.""" + with pytest.raises(ValueError, match="geometry row 0"): + make_shape_table( + [ + { + "object_id": "cell-1", + "geometry": [0.0, 1.0, 2.0], + } + ], + geometry_encoding="geoarrow.point", + ) + + with pytest.raises(ValueError, match="centroid row 0"): + make_shape_table( + [ + { + "object_id": "cell-1", + "geometry": [0.0, 1.0], + "centroid": [0.0, 1.0, 2.0], + } + ], + geometry_encoding="geoarrow.point", + ) + + with pytest.raises(ValueError, match="bbox row 0"): + make_shape_table( + [ + { + "object_id": "cell-1", + "geometry": [0.0, 1.0], + "bbox": {"min": [0.0], "max": [1.0, 2.0]}, + } + ], + geometry_encoding="geoarrow.point", + ) + + +def test_make_shape_table_reserves_custom_geometry_column() -> None: + """Do not infer custom geometry columns as measurement fields.""" + table = make_shape_table( + [ + { + "object_id": "cell-1", + "shape": [0.0, 1.0], + "area": 12.5, + } + ], + geometry_encoding="geoarrow.point", + geometry_column="shape", + ) + + assert table.column_names == [ + "object_id", + "image_id", + "label_image_id", + "label_value", + "shape", + "centroid", + "bbox", + "class", + "confidence", + "area", + ] + assert validate_shape_table(table) is None + + +def test_ome_arrow_shapes_wrapper_exposes_metadata_and_filtering() -> None: + """Wrap shape tables with small convenience accessors.""" + shapes = OMEArrowShapes.from_rows( + [ + { + "object_id": "cell-1", + "image_id": "image-1", + "label_image_id": "labels", + "label_value": 1, + "geometry": [0.0, 1.0], + }, + { + "object_id": "cell-2", + "image_id": "image-2", + "label_image_id": "labels", + "label_value": 2, + "geometry": [2.0, 3.0], + }, + ], + geometry_encoding="geoarrow.point", + ) + + subset = shapes.for_image("image-1") + + assert shapes.geometry_encoding == "geoarrow.point" + assert shapes.axes == ("y", "x") + assert subset.table["object_id"].to_pylist() == ["cell-1"] + + +def test_relationship_table_models_object_edges() -> None: + """Build object relationship edges as ordinary Arrow rows.""" + table = make_relationship_table( + [ + { + "parent_id": "cell-1", + "child_id": "nucleus-1", + "relationship_type": "contains", + "confidence": 1.0, + } + ] + ) + + assert table.schema.field("relationship_type").type == pa.string() + assert table["parent_id"].to_pylist() == ["cell-1"] + assert validate_relationship_table(table) is None + + +def test_validate_relationship_table_rejects_invalid_tables() -> None: + """Reject missing metadata, missing IDs, null IDs, and unknown relationships.""" + table = pa.table( + { + "parent_id": ["cell-1"], + "child_id": ["nucleus-1"], + "relationship_type": ["contains"], + } + ) + with pytest.raises(ValueError, match="schema metadata"): + validate_relationship_table(table) + + missing_parent = make_relationship_table( + [{"child_id": "nucleus-1", "relationship_type": "contains"}], + validate=False, + ).drop(["parent_id"]) + with pytest.raises(ValueError, match="parent_id"): + validate_relationship_table(missing_parent) + + null_child = make_relationship_table( + [ + { + "parent_id": "cell-1", + "child_id": None, + "relationship_type": "contains", + } + ], + validate=False, + ) + with pytest.raises(ValueError, match="child_id values"): + validate_relationship_table(null_child) + + missing_relationship_type = make_relationship_table( + [{"parent_id": "cell-1", "child_id": "nucleus-1"}], + validate=False, + ).drop(["relationship_type"]) + with pytest.raises(ValueError, match="relationship_type"): + validate_relationship_table(missing_relationship_type) + + unsupported_relationship = make_relationship_table( + [ + { + "parent_id": "cell-1", + "child_id": "nucleus-1", + "relationship_type": "overlaps", + } + ], + validate=False, + ) + with pytest.raises(ValueError, match="Unsupported relationship_type"): + validate_relationship_table(unsupported_relationship) + + +def test_shape_parquet_roundtrip_preserves_metadata( + tmp_path: pathlib.Path, +) -> None: + """Write and read complete shape tables with schema metadata intact.""" + path = tmp_path / "cells.ome-shapes.parquet" + table = make_shape_table( + [ + { + "object_id": "cell-1", + "image_id": "image-1", + "label_image_id": "labels-1", + "label_value": 1, + "geometry": [10.0, 20.0], + "area_um2": 42.5, + } + ], + geometry_encoding="geoarrow.point", + axes=("y", "x"), + units=("micrometer", "micrometer"), + coordinate_space="physical", + ) + + write_shape_parquet(table, path) + roundtrip = read_shape_parquet(path) + + assert roundtrip.schema.metadata[OME_ARROW_SHAPES_METADATA_KEY] + assert roundtrip["object_id"].to_pylist() == ["cell-1"] + assert roundtrip["area_um2"].to_pylist() == [42.5] + assert pq.ParquetFile(path).metadata.num_rows == 1 + + +def test_shape_parquet_projection_and_filters_support_label_references( + tmp_path: pathlib.Path, +) -> None: + """Read only needed label-reference columns with Parquet predicate filters.""" + path = tmp_path / "label_refs.ome-shapes.parquet" + table = make_shape_table( + [ + { + "object_id": "nucleus-1", + "image_id": "image-1", + "label_image_id": "nuclear-labels", + "label_value": 1, + "geometry": { + "label_image_id": "nuclear-labels", + "label_value": 1, + }, + "class": "nucleus", + }, + { + "object_id": "nucleus-2", + "image_id": "image-2", + "label_image_id": "nuclear-labels", + "label_value": 2, + "geometry": { + "label_image_id": "nuclear-labels", + "label_value": 2, + }, + "class": "nucleus", + }, + ], + geometry_encoding="ome.labelmask", + ) + + write_shape_parquet(table, path, row_group_size=1) + projected = read_shape_parquet( + path, + columns=["object_id", "image_id", "label_value"], + filters=[("image_id", "==", "image-2")], + ) + + assert projected.column_names == ["object_id", "image_id", "label_value"] + assert projected.to_pydict() == { + "object_id": ["nucleus-2"], + "image_id": ["image-2"], + "label_value": [2], + } diff --git a/tests/test_shapes_performance.py b/tests/test_shapes_performance.py new file mode 100644 index 0000000..04f6e61 --- /dev/null +++ b/tests/test_shapes_performance.py @@ -0,0 +1,99 @@ +"""Performance canaries for Arrow-native OME-Arrow shape tables.""" + +from __future__ import annotations + +import pathlib +import statistics +import time +from collections.abc import Callable + +from ome_arrow import ( + OMEArrowShapes, + make_shape_table, + read_shape_parquet, + write_shape_parquet, +) + + +def _shape_rows(count: int) -> list[dict[str, object]]: + """Create deterministic point rows for shape performance tests.""" + return [ + { + "object_id": f"cell-{i}", + "image_id": f"image-{i % 5}", + "label_image_id": "labels", + "label_value": i, + "geometry": [float(i), float(i + 1)], + "centroid": [float(i), float(i + 1)], + "class": "cell", + "area": float(i % 997), + "mean_intensity": float(i % 255), + } + for i in range(count) + ] + + +def _median_seconds(fn: Callable[[], object], repeats: int = 3) -> float: + """Return median runtime across repeated calls.""" + times = [] + for _ in range(repeats): + start = time.perf_counter() + fn() + times.append(time.perf_counter() - start) + return statistics.median(times) + + +def test_make_shape_table_handles_large_object_tables_quickly() -> None: + """Keep shape table construction on Arrow's fast path.""" + rows = _shape_rows(20_000) + + table = make_shape_table(rows, geometry_encoding="geoarrow.point") + elapsed = _median_seconds( + lambda: make_shape_table(rows, geometry_encoding="geoarrow.point") + ) + + assert table.num_rows == 20_000 + assert table.schema.field("area").type.bit_width == 64 + assert elapsed < 5.0 + + +def test_shape_filtering_stays_vectorized_for_image_ids() -> None: + """Keep image-id filtering backed by Arrow compute operations.""" + table = make_shape_table(_shape_rows(20_000), geometry_encoding="geoarrow.point") + shapes = OMEArrowShapes(table) + + subset = shapes.for_image("image-3") + elapsed = _median_seconds(lambda: shapes.for_image("image-3")) + + assert subset.table.num_rows == 4_000 + assert subset.table["image_id"].to_pylist() == ["image-3"] * 4_000 + assert elapsed < 2.0 + + +def test_shape_parquet_projection_stays_fast(tmp_path: pathlib.Path) -> None: + """Keep projected shape reads on Parquet's columnar fast path.""" + path = tmp_path / "cells.ome-shapes.parquet" + table = make_shape_table(_shape_rows(20_000), geometry_encoding="geoarrow.point") + + write_shape_parquet(table, path, row_group_size=5_000) + write_elapsed = _median_seconds( + lambda: write_shape_parquet(table, path, row_group_size=5_000) + ) + + projected = read_shape_parquet( + path, + columns=["object_id", "image_id", "area", "mean_intensity"], + filters=[("image_id", "==", "image-3")], + ) + read_elapsed = _median_seconds( + lambda: read_shape_parquet( + path, + columns=["object_id", "image_id", "area", "mean_intensity"], + filters=[("image_id", "==", "image-3")], + ) + ) + + assert projected.num_rows == 4_000 + assert projected.column_names == ["object_id", "image_id", "area", "mean_intensity"] + assert write_elapsed < 5.0 + assert read_elapsed < 2.0