feat(spatial): let readPolygon reuse a pre-resolved selection

src/spatial/polygon-reader.ts

@@ -53,6 +53,31 @@ export interface PolygonReadOptions {
    * falls back to `"center"`.
    */
   selection?: "center" | "cover";
+  /**
+   * Pre-resolved selection to reuse instead of scanning again. When supplied,
+   * {@link readPolygon} skips its internal `resolveSelection` (the O(bbox)
+   * membership scan) and streams straight from these `cells` + `bbox`. Intended
+   * for a caller that reads MANY same-shape variables through one polygon: run
+   * {@link resolvePolygonCells} ONCE and thread the result into every read
+   * instead of re-scanning per variable.
+   *
+   * MUST have been produced by {@link resolvePolygonCells} for the SAME `arr`
+   * shape/rank and the SAME `polygon` and `spatialLayout` — the reader trusts
+   * the `cells` + `bbox` verbatim and does not re-derive them, so a selection
+   * from a differently-shaped array yields wrong values. As a guard against the
+   * most common misuse, the reader bounds-checks the supplied `bbox`/`cells`
+   * against this array's spatial extents and throws {@link SliceError} on a
+   * shape mismatch rather than slicing out of range into silent fill/NaN.
+   *
+   * `spatialLayout` is still required (the per-step block read maps the `bbox`
+   * through the layout resolver). `polygon` is also still required, but only to
+   * pass {@link validatePolygonReadInput}; it does NOT feed the read once this
+   * is set. `timeRange` is applied fresh per read and need not match the one
+   * used to resolve. `maxCells` has no effect on the read (the stride is already
+   * baked into the supplied selection), but if present it must still be a valid
+   * positive integer or validation rejects it.
+   */
+  resolvedSelection?: PolygonSelection;
   /** Forwarded to ZarrArray.get (memoryCache, concurrency, maxInFlightBytes, observability, ...). */
   readOptions?: ReadOptions;
 }
@@ -358,6 +383,65 @@ function assertArrayRank(arr: ZarrArray, resolver: LayoutResolver): number {
   return k;
 }
 
+/**
+ * Guard a caller-supplied {@link PolygonReadOptions.resolvedSelection} before
+ * the reader trusts it verbatim. `resolvedSelection` bypasses the internal
+ * membership scan, so nothing else checks that its `bbox`/`cells` actually fit
+ * `arr` — a selection resolved against a different-shaped array would otherwise
+ * slice out of range and silently yield fill/NaN values (arr.get pads the
+ * out-of-range region) instead of erroring.
+ *
+ * We validate against `arr`'s ACTUAL trailing spatial extents — not the
+ * resolver's declared `nRows`/`nCols`. The resolver's extents come from
+ * `spatialLayout` (the GridIndex / lat-lon coords), so checking against them
+ * would pass whenever the caller supplies a matching `spatialLayout`, even if
+ * `arr` itself has a different-shaped spatial grid — precisely the misuse this
+ * guard exists to catch. The trailing spatial dims sit after the leading time
+ * axis and the `k` collapsed singleton middle dims: for a 3-D `[time, rows,
+ * cols]` layout they are the last two axes; for 2-D `[time, npoints]` the point
+ * axis maps to rows (cols is a synthetic single column). A mismatch is a clean
+ * {@link SliceError} at the trust boundary.
+ */
+function assertSelectionInBounds(
+  sel: PolygonSelection,
+  resolver: LayoutResolver,
+  arr: ZarrArray,
+): void {
+  const { rMin, rMax, cMin, cMax } = sel.bbox;
+  // Trailing spatial axes of `arr`, past the time axis and any collapsed
+  // singleton middle dims. `resolver.ndim` is 3 for row/col layouts (last two
+  // axes are the grid) and 2 for npoints (single trailing point axis; the
+  // resolver models cols as a size-1 synthetic column).
+  const rank = arr.shape.length;
+  const arrRows =
+    resolver.ndim === 3 ? arr.shape[rank - 2] : arr.shape[rank - 1];
+  const arrCols = resolver.ndim === 3 ? arr.shape[rank - 1] : resolver.nCols;
+  if (
+    rMin < 0 ||
+    cMin < 0 ||
+    rMax > arrRows ||
+    cMax > arrCols ||
+    rMax < rMin ||
+    cMax < cMin
+  ) {
+    throw new SliceError(
+      `polygon reader: resolvedSelection.bbox ` +
+        `[${rMin}, ${rMax})x[${cMin}, ${cMax}) is out of range for the ` +
+        `${arrRows}x${arrCols} spatial grid; the selection must ` +
+        `have been produced by resolvePolygonCells for this array's shape`,
+    );
+  }
+  for (const cell of sel.cells) {
+    if (cell.i < rMin || cell.i >= rMax || cell.j < cMin || cell.j >= cMax) {
+      throw new SliceError(
+        `polygon reader: resolvedSelection cell (${cell.i}, ${cell.j}) falls ` +
+          `outside its own bbox [${rMin}, ${rMax})x[${cMin}, ${cMax}); the ` +
+          `selection must have been produced by resolvePolygonCells`,
+      );
+    }
+  }
+}
+
 /** Bounding box in polygon (lat/lon) space. */
 function polygonEnvelope(ring: Array<[number, number]>): {
   latMin: number;
@@ -953,7 +1037,9 @@ function withSingletonMiddleDims(
  * yielding only the in-polygon cell values for that step (FR-001).
  *
  * The selection is resolved once up front (same as {@link resolvePolygonCells})
- * and each step is read as a single bounding-box block via {@link ZarrArray.get}.
+ * -- or, when `opts.resolvedSelection` is supplied, reused verbatim so a caller
+ * reading many same-shape variables through one polygon scans the bbox only once
+ * -- and each step is read as a single bounding-box block via {@link ZarrArray.get}.
  * A single {@link MemoryCache} is shared across all per-step reads, so a backing
  * chunk that spans the full time axis is fetched/decompressed at most once and
  * reused for every later step (FR-005) — and because only one time slice is
@@ -990,7 +1076,17 @@ export async function* readPolygon(
   // the trailing spatial dims (e.g. depth in `[time, 1, lat, lon]`); each is
   // collapsed by selecting index 0 in the per-step read below.
   const singletonMiddleDims = assertArrayRank(arr, resolver);
-  const { cells, bbox } = resolveSelection(opts, resolver);
+  // Reuse a caller-supplied selection (resolved ONCE for many same-shape
+  // variables) instead of re-scanning the bbox per read; see
+  // `PolygonReadOptions.resolvedSelection`. Absent it, resolve as before. A
+  // supplied selection is trusted verbatim but bounds-checked against this
+  // array's spatial extents first, so a mismatched shape errors cleanly rather
+  // than slicing out of range into silent fill/NaN values.
+  if (opts.resolvedSelection) {
+    assertSelectionInBounds(opts.resolvedSelection, resolver, arr);
+  }
+  const { cells, bbox } =
+    opts.resolvedSelection ?? resolveSelection(opts, resolver);
   if (cells.length === 0) return;
 
   const [tStart, tEnd] = opts.timeRange ?? [0, nTime];

tests/unit/polygon-reader.test.ts

@@ -620,6 +620,175 @@ describe("readPolygon streaming (2-D supplied grid)", () => {
   });
 });
 
+// ── resolvedSelection: reuse a pre-resolved selection across reads ───────────
+
+describe("readPolygon with a pre-resolved selection", () => {
+  it("streams byte-identically whether the selection is resolved inline or reused", async () => {
+    const { arr } = await makeArray({
+      shape: [3, 5, 5],
+      chunks: [3, 5, 5],
+      filler: (t, r, c) => t * 100 + r * 10 + c,
+    });
+    const opts = {
+      polygon: CONCAVE_POLY,
+      spatialLayout: { kind: "2d" as const, grid: GRID_5x5 },
+    };
+    // Resolve once, then feed that selection back into the read.
+    const sel = resolvePolygonCells(arr, opts);
+    const inline = await collect(readPolygon(arr, opts));
+    const reused = await collect(
+      readPolygon(arr, { ...opts, resolvedSelection: sel }),
+    );
+
+    expect(reused.map((s) => s.t)).toEqual(inline.map((s) => s.t));
+    reused.forEach((step, k) => {
+      expect(Array.from(step.values)).toEqual(Array.from(inline[k].values));
+    });
+  });
+
+  it("trusts the supplied selection verbatim without re-scanning the bbox", async () => {
+    const { arr } = await makeArray({
+      shape: [2, 5, 5],
+      chunks: [2, 5, 5],
+      filler: (t, r, c) => t * 100 + r * 10 + c,
+    });
+    const opts = {
+      polygon: CONCAVE_POLY,
+      spatialLayout: { kind: "2d" as const, grid: GRID_5x5 },
+    };
+    // A hand-picked selection that a fresh scan would NEVER produce: two cells,
+    // in a non-row-major order, one of them (4,4) OUTSIDE the concave polygon.
+    // The reader can only emit exactly these — in this order — if it skipped its
+    // own membership scan and streamed straight from `resolvedSelection`.
+    const forced: PolygonSelection = {
+      cells: [
+        { i: 4, j: 4, lat: 4, lon: 4 },
+        { i: 0, j: 0, lat: 0, lon: 0 },
+      ],
+      bbox: { rMin: 0, rMax: 5, cMin: 0, cMax: 5 },
+      stride: 1,
+    };
+    const steps = await collect(
+      readPolygon(arr, { ...opts, resolvedSelection: forced }),
+    );
+
+    expect(steps.map((s) => s.t)).toEqual([0, 1]);
+    for (const step of steps) {
+      expect(step.values.length).toBe(2);
+      // value = t*100 + i*10 + j, aligned to the FORCED cell order.
+      expect(step.values[0]).toBe(step.t * 100 + 44); // (4,4) first
+      expect(step.values[1]).toBe(step.t * 100 + 0); //  (0,0) second
+    }
+    // Sanity: (4,4) is genuinely outside the concave polygon, so an inline scan
+    // would have dropped it — the forced selection really is un-scannable.
+    expect(pointInPolygon(4, 4, CONCAVE_POLY)).toBe(false);
+  });
+
+  it("yields nothing for a supplied empty selection", async () => {
+    const { arr } = await makeArray({
+      shape: [3, 5, 5],
+      chunks: [3, 5, 5],
+      filler: () => 1,
+    });
+    const steps = await collect(
+      readPolygon(arr, {
+        polygon: CONCAVE_POLY,
+        spatialLayout: { kind: "2d", grid: GRID_5x5 },
+        resolvedSelection: {
+          cells: [],
+          bbox: { rMin: 0, rMax: 0, cMin: 0, cMax: 0 },
+          stride: 1,
+        },
+      }),
+    );
+    expect(steps).toEqual([]);
+  });
+
+  it("throws on a bbox that exceeds the array's spatial extents", async () => {
+    // A selection resolved against a 5×5 grid, reused against a 4×4 array — the
+    // exact same-shape misuse the contract warns about. Without the guard,
+    // arr.get would pad the out-of-range column with fill values and the read
+    // would return NaN silently; the guard must turn that into a clean error.
+    const { arr } = await makeArray({
+      shape: [2, 4, 4],
+      chunks: [2, 4, 4],
+      filler: (t, r, c) => t * 100 + r * 10 + c,
+    });
+    const grid4 = makeGrid(
+      4,
+      4,
+      (i) => i,
+      (j) => j,
+    );
+    const oversized: PolygonSelection = {
+      cells: [{ i: 0, j: 0, lat: 0, lon: 0 }],
+      bbox: { rMin: 0, rMax: 5, cMin: 0, cMax: 5 }, // 5 > nCols/nRows === 4
+      stride: 1,
+    };
+    await expect(
+      collect(
+        readPolygon(arr, {
+          polygon: CONCAVE_POLY,
+          spatialLayout: { kind: "2d", grid: grid4 },
+          resolvedSelection: oversized,
+        }),
+      ),
+    ).rejects.toBeInstanceOf(SliceError);
+  });
+
+  it("throws when the selection fits the spatialLayout but not the array's shape", async () => {
+    // The subtle misuse: the caller passes a `spatialLayout` (5×5 grid) that
+    // MATCHES the pre-resolved selection, but reuses it against a 4×4 array.
+    // Because the guard would derive its extents from the grid, a layout-based
+    // check would pass and arr.get would silently pad the out-of-range row/col
+    // with fill values. The guard must key off `arr`'s actual trailing spatial
+    // dims, so this is rejected before any read.
+    const { arr } = await makeArray({
+      shape: [2, 4, 4],
+      chunks: [2, 4, 4],
+      filler: (t, r, c) => t * 100 + r * 10 + c,
+    });
+    const sel: PolygonSelection = {
+      cells: [{ i: 4, j: 4, lat: 4, lon: 4 }],
+      bbox: { rMin: 0, rMax: 5, cMin: 0, cMax: 5 }, // fits GRID_5x5, not the 4×4 arr
+      stride: 1,
+    };
+    await expect(
+      collect(
+        readPolygon(arr, {
+          polygon: CONCAVE_POLY,
+          spatialLayout: { kind: "2d", grid: GRID_5x5 },
+          resolvedSelection: sel,
+        }),
+      ),
+    ).rejects.toBeInstanceOf(SliceError);
+  });
+
+  it("throws on a cell that falls outside its own bbox", async () => {
+    const { arr } = await makeArray({
+      shape: [2, 5, 5],
+      chunks: [2, 5, 5],
+      filler: (t, r, c) => t * 100 + r * 10 + c,
+    });
+    const stray: PolygonSelection = {
+      // (9, 0) lies outside the supplied bbox (rMax === 5) — reading it would
+      // index past the block and emit NaN. The guard rejects it before any read.
+      cells: [{ i: 9, j: 0, lat: 9, lon: 0 }],
+      bbox: { rMin: 0, rMax: 5, cMin: 0, cMax: 5 },
+      stride: 1,
+    };
+    await expect(
+      collect(
+        readPolygon(arr, {
+          polygon: CONCAVE_POLY,
+          spatialLayout: { kind: "2d", grid: GRID_5x5 },
+          resolvedSelection: stray,
+        }),
+      ),
+    ).rejects.toBeInstanceOf(SliceError);
+  });
+});
+
 // ── US2: chunk-read-once + memory bounding ───────────────────────────────────
 
 // 6×6 grid, lat=i, lon=j; a polygon spanning multiple 3×3 spatial tiles.