From 3c2a6a45626a996d948c0cce82688e4bb9c2c8ae Mon Sep 17 00:00:00 2001 From: Andre Brandao Date: Fri, 17 Jul 2026 19:41:32 -0300 Subject: [PATCH 1/3] feat(spatial): let readPolygon reuse a pre-resolved selection readPolygon re-scanned the bbox (resolveSelection -> gatherCells, O(rows*cols)) on every call. A caller reading many same-shape variables through one polygon paid that scan per variable. Add an optional resolvedSelection so the scan runs once (resolvePolygonCells) and is threaded into every read. Additive/back-compat; GridIndex reuse via spatialLayout is unchanged. --- src/spatial/polygon-reader.ts | 26 ++++++++++- tests/unit/polygon-reader.test.ts | 77 +++++++++++++++++++++++++++++++ 2 files changed, 101 insertions(+), 2 deletions(-) diff --git a/src/spatial/polygon-reader.ts b/src/spatial/polygon-reader.ts index 9492d4b..fbd3ba9 100644 --- a/src/spatial/polygon-reader.ts +++ b/src/spatial/polygon-reader.ts @@ -53,6 +53,23 @@ 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`, `spatialLayout`, `timeRange`, + * `selection`, and `maxCells` — the reader trusts it verbatim and does not + * re-derive the bbox, so a mismatched selection yields wrong values. `polygon` + * and `spatialLayout` are still required (the per-step block read uses the + * layout resolver); `maxCells` is ignored when this is set (the stride is + * already baked into the supplied selection). + */ + resolvedSelection?: PolygonSelection; /** Forwarded to ZarrArray.get (memoryCache, concurrency, maxInFlightBytes, observability, ...). */ readOptions?: ReadOptions; } @@ -953,7 +970,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 +1009,10 @@ 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. + const { cells, bbox } = opts.resolvedSelection ?? resolveSelection(opts, resolver); if (cells.length === 0) return; const [tStart, tEnd] = opts.timeRange ?? [0, nTime]; diff --git a/tests/unit/polygon-reader.test.ts b/tests/unit/polygon-reader.test.ts index 2127ad4..482f450 100644 --- a/tests/unit/polygon-reader.test.ts +++ b/tests/unit/polygon-reader.test.ts @@ -620,6 +620,83 @@ 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([]); + }); +}); + // ── US2: chunk-read-once + memory bounding ─────────────────────────────────── // 6×6 grid, lat=i, lon=j; a polygon spanning multiple 3×3 spatial tiles. From fae49db7e86e11e9e522ae15c65ddf5cf9e91d12 Mon Sep 17 00:00:00 2001 From: Andre Brandao Date: Fri, 17 Jul 2026 20:11:32 -0300 Subject: [PATCH 2/3] fix(spatial): bounds-check resolvedSelection before trusting it readPolygon consumed a caller-supplied resolvedSelection verbatim with no validation against the array's shape. A selection resolved for a differently-shaped array would slice out of range and silently yield fill/NaN values (arr.get pads the out-of-range region) instead of erroring. Add assertSelectionInBounds: validate the bbox against the resolver's spatial extents and each cell against its own bbox, throwing SliceError on mismatch. Also correct the resolvedSelection JSDoc (polygon is needed only for validation; timeRange is applied fresh per read; maxCells has no effect but must still be valid) and add tests for both guard paths. --- src/spatial/polygon-reader.ts | 72 +++++++++++++++++++++++++++---- tests/unit/polygon-reader.test.ts | 70 ++++++++++++++++++++++++++++-- 2 files changed, 131 insertions(+), 11 deletions(-) diff --git a/src/spatial/polygon-reader.ts b/src/spatial/polygon-reader.ts index fbd3ba9..7b56eeb 100644 --- a/src/spatial/polygon-reader.ts +++ b/src/spatial/polygon-reader.ts @@ -62,12 +62,20 @@ export interface PolygonReadOptions { * instead of re-scanning per variable. * * MUST have been produced by {@link resolvePolygonCells} for the SAME `arr` - * shape/rank and the SAME `polygon`, `spatialLayout`, `timeRange`, - * `selection`, and `maxCells` — the reader trusts it verbatim and does not - * re-derive the bbox, so a mismatched selection yields wrong values. `polygon` - * and `spatialLayout` are still required (the per-step block read uses the - * layout resolver); `maxCells` is ignored when this is set (the stride is - * already baked into the supplied selection). + * 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, ...). */ @@ -375,6 +383,47 @@ 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 the resolver's + * declared spatial extents (`nRows`/`nCols`) so the failure is a clean + * {@link SliceError} at the trust boundary. + */ +function assertSelectionInBounds( + sel: PolygonSelection, + resolver: LayoutResolver, +): void { + const { rMin, rMax, cMin, cMax } = sel.bbox; + if ( + rMin < 0 || + cMin < 0 || + rMax > resolver.nRows || + cMax > resolver.nCols || + rMax < rMin || + cMax < cMin + ) { + throw new SliceError( + `polygon reader: resolvedSelection.bbox ` + + `[${rMin}, ${rMax})x[${cMin}, ${cMax}) is out of range for the ` + + `${resolver.nRows}x${resolver.nCols} 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; @@ -1011,8 +1060,15 @@ export async function* readPolygon( const singletonMiddleDims = assertArrayRank(arr, 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. - const { cells, bbox } = opts.resolvedSelection ?? resolveSelection(opts, resolver); + // `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); + } + const { cells, bbox } = + opts.resolvedSelection ?? resolveSelection(opts, resolver); if (cells.length === 0) return; const [tStart, tEnd] = opts.timeRange ?? [0, nTime]; diff --git a/tests/unit/polygon-reader.test.ts b/tests/unit/polygon-reader.test.ts index 482f450..3058e1a 100644 --- a/tests/unit/polygon-reader.test.ts +++ b/tests/unit/polygon-reader.test.ts @@ -636,7 +636,9 @@ describe("readPolygon with a pre-resolved selection", () => { // 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 })); + 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) => { @@ -666,7 +668,9 @@ describe("readPolygon with a pre-resolved selection", () => { bbox: { rMin: 0, rMax: 5, cMin: 0, cMax: 5 }, stride: 1, }; - const steps = await collect(readPolygon(arr, { ...opts, resolvedSelection: forced })); + const steps = await collect( + readPolygon(arr, { ...opts, resolvedSelection: forced }), + ); expect(steps.map((s) => s.t)).toEqual([0, 1]); for (const step of steps) { @@ -690,11 +694,71 @@ describe("readPolygon with a pre-resolved selection", () => { readPolygon(arr, { polygon: CONCAVE_POLY, spatialLayout: { kind: "2d", grid: GRID_5x5 }, - resolvedSelection: { cells: [], bbox: { rMin: 0, rMax: 0, cMin: 0, cMax: 0 }, stride: 1 }, + 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 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 ─────────────────────────────────── From bff07fc13c3cc4cc330df8293590281476ed77fc Mon Sep 17 00:00:00 2001 From: Andre Brandao Date: Mon, 20 Jul 2026 09:45:00 -0300 Subject: [PATCH 3/3] fix(spatial): bounds-check resolvedSelection against arr shape, not layout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit assertSelectionInBounds validated resolvedSelection.bbox against the resolver's declared nRows/nCols, which are derived from spatialLayout (the GridIndex / lat-lon coords) rather than from arr.shape. A caller could therefore reuse a selection resolved for a different-shaped array without tripping the guard, as long as they also passed a spatialLayout that matched the selection — arr.get would then slice out of range and silently return fill/NaN, exactly what this guard exists to prevent. Key the check off arr's actual trailing spatial extents instead: - 3-D [time, rows, cols] layouts -> last two axes of arr.shape - 2-D [time, npoints] layout -> single trailing point axis Add a regression test: selection fits the spatialLayout (5x5 grid) but not the reused array (4x4) -> clean SliceError before any read. --- src/spatial/polygon-reader.ts | 30 ++++++++++++++++++++++++------ tests/unit/polygon-reader.test.ts | 28 ++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 6 deletions(-) diff --git a/src/spatial/polygon-reader.ts b/src/spatial/polygon-reader.ts index 7b56eeb..adea7c1 100644 --- a/src/spatial/polygon-reader.ts +++ b/src/spatial/polygon-reader.ts @@ -389,27 +389,45 @@ function assertArrayRank(arr: ZarrArray, resolver: LayoutResolver): number { * 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 the resolver's - * declared spatial extents (`nRows`/`nCols`) so the failure is a clean + * 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 > resolver.nRows || - cMax > resolver.nCols || + 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 ` + - `${resolver.nRows}x${resolver.nCols} spatial grid; the selection must ` + + `${arrRows}x${arrCols} spatial grid; the selection must ` + `have been produced by resolvePolygonCells for this array's shape`, ); } @@ -1065,7 +1083,7 @@ export async function* readPolygon( // 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); + assertSelectionInBounds(opts.resolvedSelection, resolver, arr); } const { cells, bbox } = opts.resolvedSelection ?? resolveSelection(opts, resolver); diff --git a/tests/unit/polygon-reader.test.ts b/tests/unit/polygon-reader.test.ts index 3058e1a..d741f55 100644 --- a/tests/unit/polygon-reader.test.ts +++ b/tests/unit/polygon-reader.test.ts @@ -736,6 +736,34 @@ describe("readPolygon with a pre-resolved selection", () => { ).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],