Add sel helper for named-dimension selection (#382)

Zarr v3 arrays can have `dimension_names` in their metadata, but `get`
and `set` only accept positional selection arrays. This means users need
to know the exact dimension order and manually map names to indices,
which is error-prone and tightly couples code to a specific array layout.

`sel` converts a record of dimension names to slices/indices into the
positional array that `get`/`set` already accept, using the array's
`dimensionNames` metadata. Unspecified dimensions default to `null`
(select all). Unknown dimension names throw immediately to catch typos.

```ts
let arr = await zarr.open(store, { kind: "array" });
// arr.dimensionNames -> ["time", "lat", "lon"]

let selection = zarr.sel(arr, { lat: zarr.slice(100, 200), time: 0 });
// -> [0, slice(100, 200), null]

let result = await zarr.get(arr, selection);
```

Author: manzt

.changeset/add-sel-helper.md

@@ -0,0 +1,15 @@
+---
+"zarrita": minor
+---
+
+Add `sel` helper for named-dimension selection. Converts a record of dimension names to the positional selection array that `get` and `set` accept, using the array's `dimensionNames` metadata. Throws for unknown dimension names.
+
+```ts
+let arr = await zarr.open(store, { kind: "array" });
+// arr.dimensionNames -> ["time", "lat", "lon"]
+
+let selection = zarr.sel(arr, { lat: zarr.slice(100, 200), time: 0 });
+// -> [0, slice(100, 200), null]
+
+let result = await zarr.get(arr, selection);
+```

packages/zarrita/__tests__/util.test.ts

@@ -1,4 +1,6 @@
 import { afterAll, beforeAll, describe, expect, test, vi } from "vitest";
+import * as zarr from "../src/index.js";
+import { sel, slice } from "../src/indexing/util.js";
 import type { ArrayMetadata, DataType } from "../src/metadata.js";
 import {
 	BoolArray,
@@ -264,6 +266,46 @@ describe("ensure_correct_scalar", () => {
 	});
 });
 
+describe("sel", () => {
+	async function make_array(dimension_names?: string[]) {
+		let h = zarr.root();
+		return zarr.create(h.resolve("/test"), {
+			shape: [100, 200, 300],
+			chunk_shape: [10, 20, 30],
+			data_type: "float32",
+			dimension_names,
+		});
+	}
+
+	test("maps named dimensions to positional selection", async () => {
+		let arr = await make_array(["time", "lat", "lon"]);
+		expect(sel(arr, { lat: slice(10, 20), time: 0 })).toStrictEqual([
+			0,
+			{ start: 10, stop: 20, step: null },
+			null,
+		]);
+	});
+
+	test("unspecified dimensions default to null", async () => {
+		let arr = await make_array(["time", "lat", "lon"]);
+		expect(sel(arr, { lon: 5 })).toStrictEqual([null, null, 5]);
+	});
+
+	test("throws for unknown dimension name", async () => {
+		let arr = await make_array(["time", "lat", "lon"]);
+		expect(() => sel(arr, { bad: 0 })).toThrowError(
+			/Unknown dimension name: "bad"/,
+		);
+	});
+
+	test("throws when array has no dimension_names", async () => {
+		let arr = await make_array();
+		expect(() => sel(arr, { time: 0 })).toThrowError(
+			/does not have dimension_names/,
+		);
+	});
+});
+
 describe("v2_to_v3_array_metadata", () => {
 	let v2meta = {
 		zarr_format: 2 as const,

packages/zarrita/src/index.ts

@@ -26,6 +26,7 @@ export type {
 	Slice,
 } from "./indexing/types.js";
 export {
+	sel,
 	slice,
 	slice_indices as _zarrita_internal_slice_indices,
 } from "./indexing/util.js";

packages/zarrita/src/indexing/util.ts

@@ -1,3 +1,6 @@
+import type { Readable } from "@zarrita/storage";
+import type { Array as ZarrArray } from "../hierarchy.js";
+import type { DataType } from "../metadata.js";
 import type { ChunkQueue, Indices, Slice } from "./types.js";
 
 /** Similar to python's `range` function. Supports positive ranges only. */
@@ -132,6 +135,25 @@ export function slice(
 	};
 }
 
+/** @category Utility */
+export function sel<D extends DataType, Store extends Readable>(
+	arr: ZarrArray<D, Store>,
+	selection: Record<string, Slice | number | null>,
+): (Slice | number | null)[] {
+	let names = arr.dimensionNames;
+	if (!names) {
+		throw new Error("Array does not have dimension_names in its metadata.");
+	}
+	for (let key of Object.keys(selection)) {
+		if (!names.includes(key)) {
+			throw new Error(
+				`Unknown dimension name: "${key}". Available dimensions: ${names.map((n) => `"${n}"`).join(", ")}`,
+			);
+		}
+	}
+	return names.map((name) => selection[name] ?? null);
+}
+
 /** Built-in "queue" for awaiting promises. */
 export function create_queue(): ChunkQueue {
 	const promises: Promise<void>[] = [];