From 46289c0ff6f1e00842ead73c5ae494ec9930a1b7 Mon Sep 17 00:00:00 2001 From: Cameron Fraser Date: Wed, 2 Sep 2026 17:35:04 -0700 Subject: [PATCH 01/14] set up workerized zarr chunk decoding! --- src/data_manager/sources/ome_zarr_source.ts | 31 ++++++++++++ src/data_manager/task_pool/runner.ts | 12 ++--- src/data_manager/task_pool/task_pool.ts | 2 +- src/data_manager/task_pool/types.ts | 8 ++- src/data_manager/zarr_decode_worker.ts | 54 +++++++++++++++++++++ 5 files changed, 98 insertions(+), 9 deletions(-) create mode 100644 src/data_manager/zarr_decode_worker.ts diff --git a/src/data_manager/sources/ome_zarr_source.ts b/src/data_manager/sources/ome_zarr_source.ts index d81545d2..37122c33 100644 --- a/src/data_manager/sources/ome_zarr_source.ts +++ b/src/data_manager/sources/ome_zarr_source.ts @@ -17,6 +17,9 @@ import { import { unitNameToSymbol } from "../../loaders/VolumeLoaderUtils.js"; import type { Chunk, LocalChunkId } from "../types.js"; import type { NumberType, TypedArray } from "../../types.js"; +import { taskHandle } from "../task_pool/task.js"; +import { ZarrEncodeTask, type ZarrDecodeTask } from "../zarr_decode_worker.js"; +import { TaskPool } from "../task_pool/TaskPool.js"; const PLACEHOLDER_NAME = "zarr source"; const PLACEHOLDER_SCENE_INDEX = 0; @@ -170,3 +173,31 @@ export class OMEZarrSource implements ChunkSource { return result; } } + +let codecCounter = 0; + +const encodeTask = taskHandle("zarrEncode", (data) => [data.buffer]); +const decodeTask = taskHandle("zarrDecode", (data) => [data.buffer]); + +export const augmentCodec = (name: string, pool: TaskPool) => { + const codec = zarr.registry.get(name); + if (codec === undefined) { + console.warn(`No such codec: ${name}`); + return; + } + + zarr.registry.set(name, async () => ({ + fromConfig: (config, meta) => { + const descriptor = { config, meta, name, id: codecCounter }; + codecCounter++; + + return { + encode: (data: Uint8Array) => pool.runTask(encodeTask, data, descriptor), + decode: (data: Uint8Array) => pool.runTask(decodeTask, data, descriptor), + }; + }, + })); +}; + +const pool = new TaskPool(); +augmentCodec("numcodecs.blosc", pool); diff --git a/src/data_manager/task_pool/runner.ts b/src/data_manager/task_pool/runner.ts index f8b24fc2..a8336159 100644 --- a/src/data_manager/task_pool/runner.ts +++ b/src/data_manager/task_pool/runner.ts @@ -2,7 +2,7 @@ import { serializeError } from "serialize-error"; import type { TypedArray, NumberType } from "../../types.js"; import { getBorrowed } from "./borrow_guard.js"; -import type { Task, WorkerRequest, WorkerResponse } from "./types.js"; +import type { Task, TaskResult, WorkerRequest, WorkerResponse } from "./types.js"; /** * Combines a `string` task `id` and a `handler` into a `Task` that can be registered and called with a `TaskPool`. @@ -35,7 +35,7 @@ import type { Task, WorkerRequest, WorkerResponse } from "./types.js"; */ export const task = ( id: Id, - handler: (...args: In) => { result: Out; transfer: Transferable[] } + handler: (...args: In) => TaskResult ): Task => { const task = handler as Task; task.taskId = id; @@ -45,7 +45,7 @@ export const task = ( export const initWorker = (tasks: Task[]) => { const handlers = new Map(tasks.map((task) => [task.taskId, task])); - const runTask = ({ id, task, args }: WorkerRequest): [WorkerResponse, Transferable[]] => { + const runTask = async ({ id, task, args }: WorkerRequest): Promise<[WorkerResponse, Transferable[]]> => { const handler = handlers.get(task); // Extract borrowed `TypedArray`s, which must be returned to the main thread on completion @@ -67,7 +67,7 @@ export const initWorker = (tasks: Task[]) => { } try { - const { result, transfer } = handler(...processedArgs); + const { result, transfer } = await handler(...processedArgs); return [{ id, task, borrows, result, error: false }, [...borrowedBuffers, ...transfer]]; } catch (e) { const result = serializeError(e); @@ -75,8 +75,8 @@ export const initWorker = (tasks: Task[]) => { } }; - self.onmessage = ({ data }: MessageEvent) => { - const [result, transfer] = runTask(data); + self.onmessage = async ({ data }: MessageEvent) => { + const [result, transfer] = await runTask(data); self.postMessage(result, { transfer }); }; diff --git a/src/data_manager/task_pool/task_pool.ts b/src/data_manager/task_pool/task_pool.ts index da250702..04d19c33 100644 --- a/src/data_manager/task_pool/task_pool.ts +++ b/src/data_manager/task_pool/task_pool.ts @@ -6,7 +6,7 @@ import { BorrowGuard, markBorrowed } from "./borrow_guard.js"; export const taskHandle = ( id: T["taskId"], transfer?: (...args: Parameters) => Transferable[] -): TaskHandle, ReturnType["result"]> => { +): TaskHandle, Awaited>["result"]> => { return { id, transfer }; }; diff --git a/src/data_manager/task_pool/types.ts b/src/data_manager/task_pool/types.ts index 7448226a..a35c7b94 100644 --- a/src/data_manager/task_pool/types.ts +++ b/src/data_manager/task_pool/types.ts @@ -3,10 +3,14 @@ import type { ErrorObject } from "serialize-error"; import type { TypedArray, NumberType } from "../../types.js"; import type { BorrowGuard } from "./borrow_guard.js"; +export type TaskResult = + | { result: Out; transfer: Transferable[] } + | Promise<{ result: Out; transfer: Transferable[] }>; + // eslint-disable-next-line @typescript-eslint/no-explicit-any export type Task = { taskId: Id; - (...args: In): { result: Out; transfer: Transferable[] }; + (...args: In): TaskResult; }; export type TaskHandle = { @@ -60,7 +64,7 @@ export type WorkerResponse = { } & ( | { error: false; - result: ReturnType["result"]; + result: Awaited>["result"]; } | { error: true; diff --git a/src/data_manager/zarr_decode_worker.ts b/src/data_manager/zarr_decode_worker.ts new file mode 100644 index 00000000..bf1ba838 --- /dev/null +++ b/src/data_manager/zarr_decode_worker.ts @@ -0,0 +1,54 @@ +import { registry, UnknownCodecError } from "zarrita"; +import { registerTask, task } from "./task_pool/registry.js"; + +type Codec = { + encode: (data: Uint8Array) => Uint8Array | Promise; + decode: (data: Uint8Array) => Uint8Array | Promise; +}; + +// hack the private `ChunkMetadata` type out of zarrita +type CodecEntry = typeof registry extends Map Promise> ? V : never; +type ChunkMetadata = Parameters[1]; + +type CodecDescriptor = { + id: number; + name: string; + config: unknown; + meta: ChunkMetadata; +}; + +const codecs: Map = new Map(); + +const getCodec = async (descriptor: CodecDescriptor): Promise => { + const savedCodec = codecs.get(descriptor.id); + if (savedCodec !== undefined) { + return savedCodec; + } + + const entry = await registry.get(descriptor.name)?.(); + if (entry === undefined) { + throw new UnknownCodecError(descriptor.name); + } + + const newCodec = entry.fromConfig(descriptor.config, descriptor.meta); + codecs.set(descriptor.id, newCodec); + return newCodec; +}; + +const encodeTask = task("zarrEncode", async (data: Uint8Array, descriptor: CodecDescriptor) => { + const codec = await getCodec(descriptor); + const result = await codec.encode(data); + return { result, transfer: [result.buffer] }; +}); + +const decodeTask = task("zarrDecode", async (data: Uint8Array, descriptor: CodecDescriptor) => { + const codec = await getCodec(descriptor); + const result = await codec.decode(data); + return { result, transfer: [result.buffer] }; +}); + +export type ZarrEncodeTask = typeof encodeTask; +export type ZarrDecodeTask = typeof decodeTask; + +registerTask(encodeTask); +registerTask(decodeTask); From 8b2d4aab55721f62d5f8f2ab3391451030094a33 Mon Sep 17 00:00:00 2001 From: Cameron Fraser Date: Tue, 8 Sep 2026 14:10:13 -0700 Subject: [PATCH 02/14] rename decode worker -> codec worker --- src/data_manager/sources/ome_zarr_source.ts | 2 +- src/data_manager/task_pool/borrow_guard.ts | 2 +- src/data_manager/task_pool/task_pool.ts | 4 +--- .../{zarr_decode_worker.ts => zarr_codec_worker.ts} | 2 +- 4 files changed, 4 insertions(+), 6 deletions(-) rename src/data_manager/{zarr_decode_worker.ts => zarr_codec_worker.ts} (97%) diff --git a/src/data_manager/sources/ome_zarr_source.ts b/src/data_manager/sources/ome_zarr_source.ts index 37122c33..558fc0fc 100644 --- a/src/data_manager/sources/ome_zarr_source.ts +++ b/src/data_manager/sources/ome_zarr_source.ts @@ -18,7 +18,7 @@ import { unitNameToSymbol } from "../../loaders/VolumeLoaderUtils.js"; import type { Chunk, LocalChunkId } from "../types.js"; import type { NumberType, TypedArray } from "../../types.js"; import { taskHandle } from "../task_pool/task.js"; -import { ZarrEncodeTask, type ZarrDecodeTask } from "../zarr_decode_worker.js"; +import type { ZarrEncodeTask, ZarrDecodeTask } from "../zarr_codec_worker.js"; import { TaskPool } from "../task_pool/TaskPool.js"; const PLACEHOLDER_NAME = "zarr source"; diff --git a/src/data_manager/task_pool/borrow_guard.ts b/src/data_manager/task_pool/borrow_guard.ts index fa191541..4ca2004b 100644 --- a/src/data_manager/task_pool/borrow_guard.ts +++ b/src/data_manager/task_pool/borrow_guard.ts @@ -10,7 +10,7 @@ type BorrowGuardEvents = { restored: void }; * Wraps a `TypedArray` that may be temporarily transferred to a worker for processing. * * A `TypedArray` guarded by this class must be accessed through the `get` method, which returns `undefined` when the - * buffer is on a worker. The guard will trigger the `"restored"` event when it is returned to the main thread. + * buffer is on a worker. The guard will trigger the `restored` event when it is returned to the main thread. */ export class BorrowGuard extends EventDispatcher { private [guardMarker] = true as const; diff --git a/src/data_manager/task_pool/task_pool.ts b/src/data_manager/task_pool/task_pool.ts index 04d19c33..06e50ff5 100644 --- a/src/data_manager/task_pool/task_pool.ts +++ b/src/data_manager/task_pool/task_pool.ts @@ -6,9 +6,7 @@ import { BorrowGuard, markBorrowed } from "./borrow_guard.js"; export const taskHandle = ( id: T["taskId"], transfer?: (...args: Parameters) => Transferable[] -): TaskHandle, Awaited>["result"]> => { - return { id, transfer }; -}; +): TaskHandle, Awaited>["result"]> => ({ id, transfer }); type StoredPromise = { resolve: (value: unknown) => void; diff --git a/src/data_manager/zarr_decode_worker.ts b/src/data_manager/zarr_codec_worker.ts similarity index 97% rename from src/data_manager/zarr_decode_worker.ts rename to src/data_manager/zarr_codec_worker.ts index bf1ba838..b299c206 100644 --- a/src/data_manager/zarr_decode_worker.ts +++ b/src/data_manager/zarr_codec_worker.ts @@ -17,7 +17,7 @@ type CodecDescriptor = { meta: ChunkMetadata; }; -const codecs: Map = new Map(); +const codecs = new Map(); const getCodec = async (descriptor: CodecDescriptor): Promise => { const savedCodec = codecs.get(descriptor.id); From 75a1971316f66a6e26c875b8686cee5834baf155 Mon Sep 17 00:00:00 2001 From: Cameron Fraser Date: Tue, 8 Sep 2026 15:08:12 -0700 Subject: [PATCH 03/14] move codec worker --- src/data_manager/sources/ome_zarr_source.ts | 2 +- src/data_manager/{ => sources}/zarr_codec_worker.ts | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) rename src/data_manager/{ => sources}/zarr_codec_worker.ts (96%) diff --git a/src/data_manager/sources/ome_zarr_source.ts b/src/data_manager/sources/ome_zarr_source.ts index 558fc0fc..9c61c30f 100644 --- a/src/data_manager/sources/ome_zarr_source.ts +++ b/src/data_manager/sources/ome_zarr_source.ts @@ -18,7 +18,7 @@ import { unitNameToSymbol } from "../../loaders/VolumeLoaderUtils.js"; import type { Chunk, LocalChunkId } from "../types.js"; import type { NumberType, TypedArray } from "../../types.js"; import { taskHandle } from "../task_pool/task.js"; -import type { ZarrEncodeTask, ZarrDecodeTask } from "../zarr_codec_worker.js"; +import type { ZarrEncodeTask, ZarrDecodeTask } from "./zarr_codec_worker.js"; import { TaskPool } from "../task_pool/TaskPool.js"; const PLACEHOLDER_NAME = "zarr source"; diff --git a/src/data_manager/zarr_codec_worker.ts b/src/data_manager/sources/zarr_codec_worker.ts similarity index 96% rename from src/data_manager/zarr_codec_worker.ts rename to src/data_manager/sources/zarr_codec_worker.ts index b299c206..f2cc68d5 100644 --- a/src/data_manager/zarr_codec_worker.ts +++ b/src/data_manager/sources/zarr_codec_worker.ts @@ -1,5 +1,5 @@ import { registry, UnknownCodecError } from "zarrita"; -import { registerTask, task } from "./task_pool/registry.js"; +import { registerTask, task } from "../task_pool/registry.js"; type Codec = { encode: (data: Uint8Array) => Uint8Array | Promise; From 824522b97f5918a59159b854de5e0a6368bf280a Mon Sep 17 00:00:00 2001 From: Cameron Fraser Date: Tue, 6 Oct 2026 16:23:34 -0700 Subject: [PATCH 04/14] fix up some imports --- src/data_manager/sources/ome_zarr_source.ts | 12 ++++-------- src/data_manager/sources/zarr_codec_worker.ts | 12 +++--------- 2 files changed, 7 insertions(+), 17 deletions(-) diff --git a/src/data_manager/sources/ome_zarr_source.ts b/src/data_manager/sources/ome_zarr_source.ts index 9c61c30f..3c499805 100644 --- a/src/data_manager/sources/ome_zarr_source.ts +++ b/src/data_manager/sources/ome_zarr_source.ts @@ -17,9 +17,8 @@ import { import { unitNameToSymbol } from "../../loaders/VolumeLoaderUtils.js"; import type { Chunk, LocalChunkId } from "../types.js"; import type { NumberType, TypedArray } from "../../types.js"; -import { taskHandle } from "../task_pool/task.js"; -import type { ZarrEncodeTask, ZarrDecodeTask } from "./zarr_codec_worker.js"; -import { TaskPool } from "../task_pool/TaskPool.js"; +import { taskHandle, TaskPool } from "../task_pool/index.js"; +import type { encodeTask as encodeTaskType, decodeTask as decodeTaskType } from "./zarr_codec_worker.js"; const PLACEHOLDER_NAME = "zarr source"; const PLACEHOLDER_SCENE_INDEX = 0; @@ -176,8 +175,8 @@ export class OMEZarrSource implements ChunkSource { let codecCounter = 0; -const encodeTask = taskHandle("zarrEncode", (data) => [data.buffer]); -const decodeTask = taskHandle("zarrDecode", (data) => [data.buffer]); +const encodeTask = taskHandle("zarrEncode", (data) => [data.buffer]); +const decodeTask = taskHandle("zarrDecode", (data) => [data.buffer]); export const augmentCodec = (name: string, pool: TaskPool) => { const codec = zarr.registry.get(name); @@ -198,6 +197,3 @@ export const augmentCodec = (name: string, pool: TaskPool) => { }, })); }; - -const pool = new TaskPool(); -augmentCodec("numcodecs.blosc", pool); diff --git a/src/data_manager/sources/zarr_codec_worker.ts b/src/data_manager/sources/zarr_codec_worker.ts index f2cc68d5..e0fa7e38 100644 --- a/src/data_manager/sources/zarr_codec_worker.ts +++ b/src/data_manager/sources/zarr_codec_worker.ts @@ -1,5 +1,5 @@ import { registry, UnknownCodecError } from "zarrita"; -import { registerTask, task } from "../task_pool/registry.js"; +import { task } from "../task_pool/index.js"; type Codec = { encode: (data: Uint8Array) => Uint8Array | Promise; @@ -35,20 +35,14 @@ const getCodec = async (descriptor: CodecDescriptor): Promise => { return newCodec; }; -const encodeTask = task("zarrEncode", async (data: Uint8Array, descriptor: CodecDescriptor) => { +export const encodeTask = task("zarrEncode", async (data: Uint8Array, descriptor: CodecDescriptor) => { const codec = await getCodec(descriptor); const result = await codec.encode(data); return { result, transfer: [result.buffer] }; }); -const decodeTask = task("zarrDecode", async (data: Uint8Array, descriptor: CodecDescriptor) => { +export const decodeTask = task("zarrDecode", async (data: Uint8Array, descriptor: CodecDescriptor) => { const codec = await getCodec(descriptor); const result = await codec.decode(data); return { result, transfer: [result.buffer] }; }); - -export type ZarrEncodeTask = typeof encodeTask; -export type ZarrDecodeTask = typeof decodeTask; - -registerTask(encodeTask); -registerTask(decodeTask); From d3f460c798c30b99b99bfcccf8026d71561367fd Mon Sep 17 00:00:00 2001 From: Cameron Fraser Date: Thu, 8 Oct 2026 16:18:41 -0700 Subject: [PATCH 05/14] make a proper "utils" module --- src/data_manager/data_manager.ts | 29 +++-------------------------- src/data_manager/utils.ts | 25 +++++++++++++++++++++++++ 2 files changed, 28 insertions(+), 26 deletions(-) diff --git a/src/data_manager/data_manager.ts b/src/data_manager/data_manager.ts index 52594a35..bce1d8af 100644 --- a/src/data_manager/data_manager.ts +++ b/src/data_manager/data_manager.ts @@ -13,6 +13,7 @@ import { deviceSizeLimitForPriority, } from "./types.js"; import type { NumberType, TypedArray } from "../types.js"; +import { defineGlobalId, swapRemove } from "./utils.js"; import PriorityQueue from "./priority_queue.js"; import type { DeviceInterface } from "./device_interface.js"; import { dataTypeToSize } from "./utils.js"; @@ -20,27 +21,15 @@ import { ChunkSource, type VolumeMetadata, type ExtVolumeDims } from "./sources/ import SlotMap from "./slot_map.js"; import { BorrowGuard } from "./task_pool/borrow_guard.js"; -const SUBSCRIBER_ID = Symbol("DataManager.subscriberId"); +// eslint-disable-next-line @typescript-eslint/no-explicit-any +const getIdForSubscriber = defineGlobalId>(Symbol("DataManager.subscriberId")); export interface IDataSubscriber { - [SUBSCRIBER_ID]?: number; onChunkLoaded?: (id: ChunkId, chunk: Chunk) => void; onChunkOnGpu?: (id: ChunkId, texture: Tex) => void; // TODO events for when chunks are evicted? } -// Subscriber IDs increment globally, like timeout IDs -let subscriberCount = 0; - -// eslint-disable-next-line @typescript-eslint/no-explicit-any -const getIdForSubscriber = (subscriber: IDataSubscriber): number => { - if (subscriber[SUBSCRIBER_ID] === undefined) { - subscriber[SUBSCRIBER_ID] = subscriberCount; - subscriberCount += 1; - } - return subscriber[SUBSCRIBER_ID]; -}; - type SourceEntry = { source: ChunkSource; subscribers: IDataSubscriber[]; @@ -51,18 +40,6 @@ type RequestEntry = { controller: AbortController; }; -const swapRemove = (arr: T[], index: number) => { - if (index < 0) { - return; - } - - const { length } = arr; - const replace = arr.pop(); - if (replace !== undefined && index < length) { - arr[index] = replace; - } -}; - type ChunkQueue = PriorityQueue; class ChunkQueues { diff --git a/src/data_manager/utils.ts b/src/data_manager/utils.ts index 122931ae..69a48670 100644 --- a/src/data_manager/utils.ts +++ b/src/data_manager/utils.ts @@ -10,3 +10,28 @@ export const dataTypeToSize: { [T in NumberType]: number } = { float32: 4, float64: 8, }; + +/** Creates a function that returns a stable and globally unique number for any object passed to it. */ +export const defineGlobalId = (key: symbol): ((obj: T) => number) => { + let id = 0; + return (obj: T): number => { + if (obj[key] === undefined) { + obj[key] = id; + id += 1; + } + return obj[key]; + }; +}; + +/** Removes an element from an array by replacing it with the element at the end of the array. */ +export const swapRemove = (arr: T[], index: number) => { + if (index < 0 || index >= arr.length) { + return; + } + + const { length } = arr; + const replace = arr.pop(); + if (replace !== undefined && index < length) { + arr[index] = replace; + } +}; From 96ed6e9b1739afbf48697d262ffb7cdfb395c4c1 Mon Sep 17 00:00:00 2001 From: Cameron Fraser Date: Fri, 16 Oct 2026 15:48:01 -0700 Subject: [PATCH 06/14] move codec stuff into a new folder --- .../sources/ome_zarr/with_threaded_codecs.ts | 29 +++++++++++++++++++ .../{ => ome_zarr}/zarr_codec_worker.ts | 2 +- src/data_manager/sources/ome_zarr_source.ts | 27 ----------------- 3 files changed, 30 insertions(+), 28 deletions(-) create mode 100644 src/data_manager/sources/ome_zarr/with_threaded_codecs.ts rename src/data_manager/sources/{ => ome_zarr}/zarr_codec_worker.ts (96%) diff --git a/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts b/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts new file mode 100644 index 00000000..4fbc5177 --- /dev/null +++ b/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts @@ -0,0 +1,29 @@ +import * as zarr from "zarrita"; + +import { taskHandle, TaskPool } from "../../task_pool/index.js"; +import type { encodeTask as encodeTaskType, decodeTask as decodeTaskType } from "./zarr_codec_worker.js"; + +let codecCounter = 0; + +const encodeTask = taskHandle("zarrEncode", (data) => [data.buffer]); +const decodeTask = taskHandle("zarrDecode", (data) => [data.buffer]); + +export const augmentCodec = (name: string, pool: TaskPool) => { + const codec = zarr.registry.get(name); + if (codec === undefined) { + console.warn(`No such codec: ${name}`); + return; + } + + zarr.registry.set(name, async () => ({ + fromConfig: (config, meta) => { + const descriptor = { config, meta, name, id: codecCounter }; + codecCounter++; + + return { + encode: (data: Uint8Array) => pool.runTask(encodeTask, data, descriptor), + decode: (data: Uint8Array) => pool.runTask(decodeTask, data, descriptor), + }; + }, + })); +}; diff --git a/src/data_manager/sources/zarr_codec_worker.ts b/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts similarity index 96% rename from src/data_manager/sources/zarr_codec_worker.ts rename to src/data_manager/sources/ome_zarr/zarr_codec_worker.ts index e0fa7e38..e2e64049 100644 --- a/src/data_manager/sources/zarr_codec_worker.ts +++ b/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts @@ -1,5 +1,5 @@ import { registry, UnknownCodecError } from "zarrita"; -import { task } from "../task_pool/index.js"; +import { task } from "../../task_pool/index.js"; type Codec = { encode: (data: Uint8Array) => Uint8Array | Promise; diff --git a/src/data_manager/sources/ome_zarr_source.ts b/src/data_manager/sources/ome_zarr_source.ts index 3c499805..d81545d2 100644 --- a/src/data_manager/sources/ome_zarr_source.ts +++ b/src/data_manager/sources/ome_zarr_source.ts @@ -17,8 +17,6 @@ import { import { unitNameToSymbol } from "../../loaders/VolumeLoaderUtils.js"; import type { Chunk, LocalChunkId } from "../types.js"; import type { NumberType, TypedArray } from "../../types.js"; -import { taskHandle, TaskPool } from "../task_pool/index.js"; -import type { encodeTask as encodeTaskType, decodeTask as decodeTaskType } from "./zarr_codec_worker.js"; const PLACEHOLDER_NAME = "zarr source"; const PLACEHOLDER_SCENE_INDEX = 0; @@ -172,28 +170,3 @@ export class OMEZarrSource implements ChunkSource { return result; } } - -let codecCounter = 0; - -const encodeTask = taskHandle("zarrEncode", (data) => [data.buffer]); -const decodeTask = taskHandle("zarrDecode", (data) => [data.buffer]); - -export const augmentCodec = (name: string, pool: TaskPool) => { - const codec = zarr.registry.get(name); - if (codec === undefined) { - console.warn(`No such codec: ${name}`); - return; - } - - zarr.registry.set(name, async () => ({ - fromConfig: (config, meta) => { - const descriptor = { config, meta, name, id: codecCounter }; - codecCounter++; - - return { - encode: (data: Uint8Array) => pool.runTask(encodeTask, data, descriptor), - decode: (data: Uint8Array) => pool.runTask(decodeTask, data, descriptor), - }; - }, - })); -}; From d7de0533dd95227dd3241d35886fb4eb1aaf091f Mon Sep 17 00:00:00 2001 From: Cameron Fraser Date: Sat, 17 Oct 2026 10:26:41 -0700 Subject: [PATCH 07/14] start setting up single registry entry for threading --- .../sources/ome_zarr/with_threaded_codecs.ts | 57 ++++++++++++++----- .../sources/ome_zarr/zarr_codec_worker.ts | 33 ++++++++--- 2 files changed, 66 insertions(+), 24 deletions(-) diff --git a/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts b/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts index 4fbc5177..9a9568e3 100644 --- a/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts +++ b/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts @@ -8,22 +8,49 @@ let codecCounter = 0; const encodeTask = taskHandle("zarrEncode", (data) => [data.buffer]); const decodeTask = taskHandle("zarrDecode", (data) => [data.buffer]); -export const augmentCodec = (name: string, pool: TaskPool) => { - const codec = zarr.registry.get(name); - if (codec === undefined) { - console.warn(`No such codec: ${name}`); - return; +type WorkerConfig = { + pool: TaskPool; + codecs: zarr.CodecMetadata[]; +}; + +zarr.registry.set("vole-worker", async () => ({ + fromConfig: (config, meta) => { + const { pool, codecs } = config as WorkerConfig; + const descriptor = { codecs, meta, id: codecCounter }; + codecCounter++; + + return { + encode: (data: Uint8Array) => pool.runTask(encodeTask, data, descriptor), + decode: (data: Uint8Array) => pool.runTask(decodeTask, data, descriptor), + kind: "bytes_to_bytes", // TODO always true? + }; + }, +})); + +const isMetaKey = (key: string): boolean => + key.endsWith(".zarray") || key.endsWith(".zgroup") || key.endsWith(".zattrs") || key.endsWith("zarr.json"); + +/** zarrita `Store` extension that intercepts metadata to replace as many codecs as possible with workerized variants. */ +export const withThreadedCodecs = zarr.defineStoreExtension((store, opts: { pool: TaskPool }) => { + if (opts?.pool === undefined) { + return {}; } + const { pool } = opts; + + return { + get: async (key: zarr.AbsolutePath, options?: zarr.GetOptions): Promise => { + const bytes = await store.get(key, options); - zarr.registry.set(name, async () => ({ - fromConfig: (config, meta) => { - const descriptor = { config, meta, name, id: codecCounter }; - codecCounter++; + if (!isMetaKey(key)) { + return bytes; + } - return { - encode: (data: Uint8Array) => pool.runTask(encodeTask, data, descriptor), - decode: (data: Uint8Array) => pool.runTask(decodeTask, data, descriptor), - }; + const str = new TextDecoder().decode(bytes); + try { + JSON.parse(str); + } catch (cause) { + throw new zarr.InvalidMetadataError("Failed to decode JSON", { cause }); + } }, - })); -}; + }; +}); diff --git a/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts b/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts index e2e64049..884a6396 100644 --- a/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts +++ b/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts @@ -1,4 +1,4 @@ -import { registry, UnknownCodecError } from "zarrita"; +import { registry, UnknownCodecError, type CodecMetadata } from "zarrita"; import { task } from "../../task_pool/index.js"; type Codec = { @@ -12,25 +12,40 @@ type ChunkMetadata = Parameters[1]; type CodecDescriptor = { id: number; - name: string; - config: unknown; + codecs: CodecMetadata[]; meta: ChunkMetadata; }; const codecs = new Map(); +const joinCodecs = (codecs: Codec[]): Codec => { + // TODO +}; + +const loadCodec = async (descriptor: CodecDescriptor): Promise => { + const codecPromises = descriptor.codecs.map(async ({ name, configuration }): Promise => { + const entry = await registry.get(name)?.(); + if (entry === undefined) { + throw new UnknownCodecError(name); + } + return entry.fromConfig(configuration, descriptor.meta); + }); + + const codecs = await Promise.all(codecPromises); + + if (codecs.length === 1) { + return codecs[0]; + } + return joinCodecs(codecs); +}; + const getCodec = async (descriptor: CodecDescriptor): Promise => { const savedCodec = codecs.get(descriptor.id); if (savedCodec !== undefined) { return savedCodec; } - const entry = await registry.get(descriptor.name)?.(); - if (entry === undefined) { - throw new UnknownCodecError(descriptor.name); - } - - const newCodec = entry.fromConfig(descriptor.config, descriptor.meta); + const newCodec = await loadCodec(descriptor); codecs.set(descriptor.id, newCodec); return newCodec; }; From fb793214a95158ddfa247b6ef4e08b439425f946 Mon Sep 17 00:00:00 2001 From: Cameron Fraser Date: Sun, 18 Oct 2026 23:03:27 -0700 Subject: [PATCH 08/14] implement `joinCodecs` --- .../sources/ome_zarr/zarr_codec_worker.ts | 36 +++++++++++++++---- tsconfig.json | 2 +- 2 files changed, 31 insertions(+), 7 deletions(-) diff --git a/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts b/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts index 884a6396..dd27c182 100644 --- a/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts +++ b/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts @@ -1,4 +1,4 @@ -import { registry, UnknownCodecError, type CodecMetadata } from "zarrita"; +import { CodecPipelineError, registry, UnknownCodecError, type CodecMetadata } from "zarrita"; import { task } from "../../task_pool/index.js"; type Codec = { @@ -18,23 +18,47 @@ type CodecDescriptor = { const codecs = new Map(); -const joinCodecs = (codecs: Codec[]): Codec => { - // TODO +const joinCodecs = (codecs: { codec: Codec; name: string }[]): Codec => { + // Simplified version of zarrita's `createCodecPipeline` + const runStep = async (direction: "encode" | "decode", codec: string, fn: () => Promise | T): Promise => { + try { + return await fn(); + } catch (cause) { + throw new CodecPipelineError({ direction, codec, cause }); + } + }; + const codecsReversed = codecs.toReversed(); + + return { + encode: async (data: Uint8Array): Promise => { + for (const { name, codec } of codecs) { + data = await runStep("encode", name, () => codec.encode(data)); + } + return data; + }, + decode: async (data: Uint8Array): Promise => { + for (const { name, codec } of codecsReversed) { + data = await runStep("decode", name, () => codec.decode(data)); + } + return data; + }, + }; }; const loadCodec = async (descriptor: CodecDescriptor): Promise => { - const codecPromises = descriptor.codecs.map(async ({ name, configuration }): Promise => { + const codecPromises = descriptor.codecs.map(async ({ name, configuration }) => { const entry = await registry.get(name)?.(); if (entry === undefined) { throw new UnknownCodecError(name); } - return entry.fromConfig(configuration, descriptor.meta); + const codec = entry.fromConfig(configuration, descriptor.meta); + return { codec, name }; }); const codecs = await Promise.all(codecPromises); if (codecs.length === 1) { - return codecs[0]; + return codecs[0].codec; } return joinCodecs(codecs); }; diff --git a/tsconfig.json b/tsconfig.json index abd0857c..fd26722f 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -4,7 +4,7 @@ "allowSyntheticDefaultImports": true, "declaration": true, "declarationDir": "es/types", - "lib": ["DOM", "ES2022", "ES6", "DOM.Iterable", "ScriptHost", "ES2016.Array.Include", "ES2024.ArrayBuffer"], + "lib": ["DOM", "ES2023", "ES6", "DOM.Iterable", "ScriptHost", "ES2016.Array.Include", "ES2024.ArrayBuffer"], "module": "Node16", "moduleResolution": "node16", "noImplicitAny": false, From e7e3824ea65f5a8a98a1fc711602183ceb4fbb14 Mon Sep 17 00:00:00 2001 From: Cameron Fraser Date: Mon, 19 Oct 2026 15:14:30 -0700 Subject: [PATCH 09/14] futzing --- .../sources/ome_zarr/with_threaded_codecs.ts | 13 ++++++++++--- .../sources/ome_zarr/zarr_codec_worker.ts | 13 +++++++------ 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts b/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts index 9a9568e3..368bf59a 100644 --- a/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts +++ b/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts @@ -5,33 +5,40 @@ import type { encodeTask as encodeTaskType, decodeTask as decodeTaskType } from let codecCounter = 0; +type CodecKind = "array_to_array" | "array_to_bytes" | "bytes_to_bytes"; + +// MARK: Codec interface + const encodeTask = taskHandle("zarrEncode", (data) => [data.buffer]); const decodeTask = taskHandle("zarrDecode", (data) => [data.buffer]); type WorkerConfig = { pool: TaskPool; + kind: CodecKind; codecs: zarr.CodecMetadata[]; }; zarr.registry.set("vole-worker", async () => ({ fromConfig: (config, meta) => { - const { pool, codecs } = config as WorkerConfig; + const { pool, codecs, kind } = config as WorkerConfig; const descriptor = { codecs, meta, id: codecCounter }; codecCounter++; return { encode: (data: Uint8Array) => pool.runTask(encodeTask, data, descriptor), decode: (data: Uint8Array) => pool.runTask(decodeTask, data, descriptor), - kind: "bytes_to_bytes", // TODO always true? + kind, }; }, })); +// MARK: Metadata parsing + const isMetaKey = (key: string): boolean => key.endsWith(".zarray") || key.endsWith(".zgroup") || key.endsWith(".zattrs") || key.endsWith("zarr.json"); /** zarrita `Store` extension that intercepts metadata to replace as many codecs as possible with workerized variants. */ -export const withThreadedCodecs = zarr.defineStoreExtension((store, opts: { pool: TaskPool }) => { +export const withWorkerCodec = zarr.defineStoreExtension((store, opts: { pool?: TaskPool }) => { if (opts?.pool === undefined) { return {}; } diff --git a/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts b/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts index dd27c182..f91265b9 100644 --- a/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts +++ b/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts @@ -27,18 +27,19 @@ const joinCodecs = (codecs: { codec: Codec; name: string }[]): Codec => { throw new CodecPipelineError({ direction, codec, cause }); } }; - const codecsReversed = codecs.toReversed(); + // Put codecs in bytes -> array order to prioritize decoding over encoding + codecs.reverse(); return { - encode: async (data: Uint8Array): Promise => { + decode: async (data: Uint8Array): Promise => { for (const { name, codec } of codecs) { - data = await runStep("encode", name, () => codec.encode(data)); + data = await runStep("decode", name, () => codec.decode(data)); } return data; }, - decode: async (data: Uint8Array): Promise => { - for (const { name, codec } of codecsReversed) { - data = await runStep("decode", name, () => codec.decode(data)); + encode: async (data: Uint8Array): Promise => { + for (const { name, codec } of codecs.toReversed()) { + data = await runStep("encode", name, () => codec.encode(data)); } return data; }, From dd540e89cf94d0922d56681d312342835c962a14 Mon Sep 17 00:00:00 2001 From: Cameron Fraser Date: Mon, 19 Oct 2026 15:41:48 -0700 Subject: [PATCH 10/14] back to one new codec per pool --- .../sources/ome_zarr/with_threaded_codecs.ts | 50 +++++++++++++------ 1 file changed, 34 insertions(+), 16 deletions(-) diff --git a/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts b/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts index 368bf59a..920964b0 100644 --- a/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts +++ b/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts @@ -1,36 +1,53 @@ import * as zarr from "zarrita"; import { taskHandle, TaskPool } from "../../task_pool/index.js"; +import { defineGlobalId } from "../../utils.js"; import type { encodeTask as encodeTaskType, decodeTask as decodeTaskType } from "./zarr_codec_worker.js"; -let codecCounter = 0; - type CodecKind = "array_to_array" | "array_to_bytes" | "bytes_to_bytes"; -// MARK: Codec interface +// MARK: Codec wrapper const encodeTask = taskHandle("zarrEncode", (data) => [data.buffer]); const decodeTask = taskHandle("zarrDecode", (data) => [data.buffer]); type WorkerConfig = { - pool: TaskPool; kind: CodecKind; codecs: zarr.CodecMetadata[]; }; -zarr.registry.set("vole-worker", async () => ({ - fromConfig: (config, meta) => { - const { pool, codecs, kind } = config as WorkerConfig; - const descriptor = { codecs, meta, id: codecCounter }; - codecCounter++; +const getPoolId = defineGlobalId(Symbol("OMEZarrSource.codecPoolId")); + +/** + * Returns the name of the `zarrita` wrapper codec that runs codec pipelines in the given `pool`. + * + * If no such codec exists, this function creates one and registers it with `zarrita`. + */ +const getPoolCodecName = (pool: TaskPool): string => { + const id = getPoolId(pool); + const name = `vole-worker-${id}`; + + if (zarr.registry.get(name) === undefined) { + // Assign every codec pipeline created in this pool a serial number for easy lookup worker-side + let codecCounter = 0; + + zarr.registry.set(name, async () => ({ + fromConfig: (config, meta) => { + const { codecs, kind } = config as WorkerConfig; + const descriptor = { codecs, meta, id: codecCounter }; + codecCounter++; - return { - encode: (data: Uint8Array) => pool.runTask(encodeTask, data, descriptor), - decode: (data: Uint8Array) => pool.runTask(decodeTask, data, descriptor), - kind, - }; - }, -})); + return { + encode: (data: Uint8Array) => pool.runTask(encodeTask, data, descriptor), + decode: (data: Uint8Array) => pool.runTask(decodeTask, data, descriptor), + kind, + }; + }, + })); + } + + return name; +}; // MARK: Metadata parsing @@ -43,6 +60,7 @@ export const withWorkerCodec = zarr.defineStoreExtension((store, opts: { pool?: return {}; } const { pool } = opts; + const name = getPoolCodecName(pool); return { get: async (key: zarr.AbsolutePath, options?: zarr.GetOptions): Promise => { From a4bb60eccce449359d31e6583afa1467501d615d Mon Sep 17 00:00:00 2001 From: Cameron Fraser Date: Tue, 20 Oct 2026 13:23:18 -0700 Subject: [PATCH 11/14] finish basic metadata parsing --- .../sources/ome_zarr/with_threaded_codecs.ts | 35 ++++++++++++++++--- 1 file changed, 31 insertions(+), 4 deletions(-) diff --git a/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts b/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts index 920964b0..0b4009ec 100644 --- a/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts +++ b/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts @@ -51,8 +51,15 @@ const getPoolCodecName = (pool: TaskPool): string => { // MARK: Metadata parsing -const isMetaKey = (key: string): boolean => - key.endsWith(".zarray") || key.endsWith(".zgroup") || key.endsWith(".zattrs") || key.endsWith("zarr.json"); +const wrapCodecMetadata = (name: string, meta: zarr.CodecMetadata[]): zarr.CodecMetadata[] => { + const shardIndex = meta.findIndex(({ name }) => name === "sharding_indexed"); + if (shardIndex !== undefined) { + const codecs = meta[shardIndex].configuration.codecs as zarr.CodecMetadata[]; + meta[shardIndex].configuration.codecs = wrapCodecMetadata(name, codecs); + return meta; + } + return [{ name, configuration: { kind: "array_to_array", codecs: meta } satisfies WorkerConfig }]; +}; /** zarrita `Store` extension that intercepts metadata to replace as many codecs as possible with workerized variants. */ export const withWorkerCodec = zarr.defineStoreExtension((store, opts: { pool?: TaskPool }) => { @@ -66,16 +73,36 @@ export const withWorkerCodec = zarr.defineStoreExtension((store, opts: { pool?: get: async (key: zarr.AbsolutePath, options?: zarr.GetOptions): Promise => { const bytes = await store.get(key, options); - if (!isMetaKey(key)) { + if (!key.endsWith("zarr.json")) { return bytes; } const str = new TextDecoder().decode(bytes); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + let meta: any; try { - JSON.parse(str); + meta = JSON.parse(str); } catch (cause) { throw new zarr.InvalidMetadataError("Failed to decode JSON", { cause }); } + + if (meta.zarr_format === 3 && Array.isArray(meta.codecs)) { + meta = wrapCodecMetadata(name, meta); + } + + const result = JSON.stringify( + meta, + (_, value) => { + if (typeof value === "number") { + if (Number.isNaN(value)) return "NaN"; + if (value === Infinity) return "Infinity"; + if (value === -Infinity) return "-Infinity"; + } + return value; + }, + 2 + ); + return new TextEncoder().encode(result); }, }; }); From e119ed2297e1345d5ed0b91217670b8afd674ffb Mon Sep 17 00:00:00 2001 From: Cameron Fraser Date: Tue, 20 Oct 2026 14:13:19 -0700 Subject: [PATCH 12/14] set up test with task pool --- public/index.ts | 5 ++++- public/test_worker.ts | 5 +++++ 2 files changed, 9 insertions(+), 1 deletion(-) create mode 100644 public/test_worker.ts diff --git a/public/index.ts b/public/index.ts index 0c6eb25a..f43afe93 100644 --- a/public/index.ts +++ b/public/index.ts @@ -30,6 +30,7 @@ import DataManager, { IDataSubscriber } from "../src/data_manager/data_manager.j import { OMEZarrSource } from "../src/data_manager/sources/index.js"; import { ChunkId, ChunkPriority, ChunkPriorityLevel } from "../src/data_manager/types.js"; import { ThreeInterface } from "../src/data_manager/device_interface.js"; +import { TaskPool } from "../src/data_manager/task_pool/task_pool.js"; const CACHE_MAX_SIZE = 1_000_000_000; const CONCURRENCY_LIMIT = 8; @@ -925,10 +926,12 @@ function shuffle(arr: T[]): T[] { } async function testDataManager() { + const pool = new TaskPool(new URL("./test_worker.js", import.meta.url)); const dataManager = new DataManager(new ThreeInterface()); const zarrSource = await OMEZarrSource.new( // "https://animatedcell-test-data.s3.us-west-2.amazonaws.com/variance/1.zarr" - "https://uk1s3.embassy.ebi.ac.uk/idr/zarr/v0.4/idr0048A/9846152.zarr/" + "https://vast-files.int.allencell.org/users/alex.khang/repos/bioio_conversion/H2B_640_CAAX_488_20250904_flat_colony_1.ome.zarr/", + { pool } ); const sourceId = dataManager.addSource(zarrSource); const mockSubscriber: IDataSubscriber = { diff --git a/public/test_worker.ts b/public/test_worker.ts new file mode 100644 index 00000000..962404a8 --- /dev/null +++ b/public/test_worker.ts @@ -0,0 +1,5 @@ +import { copyChunkTask, reorderChunkTask } from "../src/data_manager/indexing_task.js"; +import { decodeTask, encodeTask } from "../src/data_manager/sources/ome_zarr/zarr_codec_worker.js"; +import { initWorker } from "../src/data_manager/task_pool/runner.js"; + +initWorker([encodeTask, decodeTask, copyChunkTask, reorderChunkTask]); From f83b3a2a8762eb7b8f4ee2d6a9fc7c914d3ad35e Mon Sep 17 00:00:00 2001 From: Cameron Fraser Date: Tue, 20 Oct 2026 14:13:36 -0700 Subject: [PATCH 13/14] properly use task pool in zarr source --- src/data_manager/sources/ome_zarr_source.ts | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/data_manager/sources/ome_zarr_source.ts b/src/data_manager/sources/ome_zarr_source.ts index d81545d2..830e2c9e 100644 --- a/src/data_manager/sources/ome_zarr_source.ts +++ b/src/data_manager/sources/ome_zarr_source.ts @@ -17,6 +17,8 @@ import { import { unitNameToSymbol } from "../../loaders/VolumeLoaderUtils.js"; import type { Chunk, LocalChunkId } from "../types.js"; import type { NumberType, TypedArray } from "../../types.js"; +import type { TaskPool } from "../task_pool/task_pool.js"; +import { withWorkerCodec } from "./ome_zarr/with_threaded_codecs.js"; const PLACEHOLDER_NAME = "zarr source"; const PLACEHOLDER_SCENE_INDEX = 0; @@ -34,8 +36,8 @@ export class OMEZarrSource implements ChunkSource { public readonly axesTCZYX: [number, number, number, number, number] ) {} - static async new(url: string): Promise { - const store = new zarr.FetchStore(url); + static async new(url: string, options?: { pool?: TaskPool }): Promise { + const store = zarr.extendStore(new zarr.FetchStore(url), (s) => withWorkerCodec(s, { pool: options?.pool })); const root = zarr.root(store); const group = await zarr.open(root, { kind: "group" }); From 6957706cdf88725004839ee2b3c8d7004e82a3e5 Mon Sep 17 00:00:00 2001 From: Cameron Fraser Date: Tue, 20 Oct 2026 14:15:48 -0700 Subject: [PATCH 14/14] fix everything up and get it working! --- .../sources/ome_zarr/with_threaded_codecs.ts | 24 +++-------- .../sources/ome_zarr/zarr_codec_worker.ts | 43 +++++++++++-------- 2 files changed, 33 insertions(+), 34 deletions(-) diff --git a/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts b/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts index 0b4009ec..f18aa1d6 100644 --- a/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts +++ b/src/data_manager/sources/ome_zarr/with_threaded_codecs.ts @@ -3,12 +3,13 @@ import * as zarr from "zarrita"; import { taskHandle, TaskPool } from "../../task_pool/index.js"; import { defineGlobalId } from "../../utils.js"; import type { encodeTask as encodeTaskType, decodeTask as decodeTaskType } from "./zarr_codec_worker.js"; +import type { NumberType } from "../../../types.js"; type CodecKind = "array_to_array" | "array_to_bytes" | "bytes_to_bytes"; // MARK: Codec wrapper -const encodeTask = taskHandle("zarrEncode", (data) => [data.buffer]); +const encodeTask = taskHandle("zarrEncode", (chunk) => [chunk.data.buffer]); const decodeTask = taskHandle("zarrDecode", (data) => [data.buffer]); type WorkerConfig = { @@ -38,7 +39,7 @@ const getPoolCodecName = (pool: TaskPool): string => { codecCounter++; return { - encode: (data: Uint8Array) => pool.runTask(encodeTask, data, descriptor), + encode: (data: zarr.Chunk) => pool.runTask(encodeTask, data, descriptor), decode: (data: Uint8Array) => pool.runTask(decodeTask, data, descriptor), kind, }; @@ -53,12 +54,12 @@ const getPoolCodecName = (pool: TaskPool): string => { const wrapCodecMetadata = (name: string, meta: zarr.CodecMetadata[]): zarr.CodecMetadata[] => { const shardIndex = meta.findIndex(({ name }) => name === "sharding_indexed"); - if (shardIndex !== undefined) { + if (shardIndex > -1) { const codecs = meta[shardIndex].configuration.codecs as zarr.CodecMetadata[]; meta[shardIndex].configuration.codecs = wrapCodecMetadata(name, codecs); return meta; } - return [{ name, configuration: { kind: "array_to_array", codecs: meta } satisfies WorkerConfig }]; + return [{ name, configuration: { kind: "array_to_bytes", codecs: meta } satisfies WorkerConfig }]; }; /** zarrita `Store` extension that intercepts metadata to replace as many codecs as possible with workerized variants. */ @@ -87,21 +88,10 @@ export const withWorkerCodec = zarr.defineStoreExtension((store, opts: { pool?: } if (meta.zarr_format === 3 && Array.isArray(meta.codecs)) { - meta = wrapCodecMetadata(name, meta); + meta.codecs = wrapCodecMetadata(name, meta.codecs); } - const result = JSON.stringify( - meta, - (_, value) => { - if (typeof value === "number") { - if (Number.isNaN(value)) return "NaN"; - if (value === Infinity) return "Infinity"; - if (value === -Infinity) return "-Infinity"; - } - return value; - }, - 2 - ); + const result = JSON.stringify(meta); return new TextEncoder().encode(result); }, }; diff --git a/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts b/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts index f91265b9..6765b26f 100644 --- a/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts +++ b/src/data_manager/sources/ome_zarr/zarr_codec_worker.ts @@ -1,10 +1,15 @@ -import { CodecPipelineError, registry, UnknownCodecError, type CodecMetadata } from "zarrita"; +import { type Chunk, CodecPipelineError, registry, UnknownCodecError, type CodecMetadata } from "zarrita"; import { task } from "../../task_pool/index.js"; +import type { NumberType } from "../../../types.js"; -type Codec = { - encode: (data: Uint8Array) => Uint8Array | Promise; - decode: (data: Uint8Array) => Uint8Array | Promise; +type Codec = { + encode: (data: In) => Out | Promise; + decode: (data: Out) => In | Promise; }; +type ArrayToArrayCodec = Codec, Chunk>; +type ArrayToBytesCodec = Codec, Uint8Array>; +type BytesToBytesCodec = Codec; +type AnyCodec = ArrayToArrayCodec | ArrayToBytesCodec | BytesToBytesCodec; // hack the private `ChunkMetadata` type out of zarrita type CodecEntry = typeof registry extends Map Promise> ? V : never; @@ -16,9 +21,9 @@ type CodecDescriptor = { meta: ChunkMetadata; }; -const codecs = new Map(); +const codecs = new Map(); -const joinCodecs = (codecs: { codec: Codec; name: string }[]): Codec => { +const joinCodecs = (codecs: { codec: AnyCodec; name: string }[]): ArrayToBytesCodec => { // Simplified version of zarrita's `createCodecPipeline` const runStep = async (direction: "encode" | "decode", codec: string, fn: () => Promise | T): Promise => { try { @@ -31,22 +36,26 @@ const joinCodecs = (codecs: { codec: Codec; name: string }[]): Codec => { codecs.reverse(); return { - decode: async (data: Uint8Array): Promise => { + decode: async (data: Uint8Array): Promise> => { + let result: Uint8Array | Chunk = data; for (const { name, codec } of codecs) { - data = await runStep("decode", name, () => codec.decode(data)); + // evil cast + result = await runStep("decode", name, () => codec.decode(result as Chunk & Uint8Array)); } - return data; + return result as Chunk; }, - encode: async (data: Uint8Array): Promise => { + encode: async (data: Chunk): Promise => { + let result: Uint8Array | Chunk = data; for (const { name, codec } of codecs.toReversed()) { - data = await runStep("encode", name, () => codec.encode(data)); + // evil cast pt. 2 + result = await runStep("encode", name, () => codec.encode(result as Chunk & Uint8Array)); } - return data; + return result as Uint8Array; }, }; }; -const loadCodec = async (descriptor: CodecDescriptor): Promise => { +const loadCodec = async (descriptor: CodecDescriptor): Promise => { const codecPromises = descriptor.codecs.map(async ({ name, configuration }) => { const entry = await registry.get(name)?.(); if (entry === undefined) { @@ -64,7 +73,7 @@ const loadCodec = async (descriptor: CodecDescriptor): Promise => { return joinCodecs(codecs); }; -const getCodec = async (descriptor: CodecDescriptor): Promise => { +const getCodec = async (descriptor: CodecDescriptor): Promise => { const savedCodec = codecs.get(descriptor.id); if (savedCodec !== undefined) { return savedCodec; @@ -75,14 +84,14 @@ const getCodec = async (descriptor: CodecDescriptor): Promise => { return newCodec; }; -export const encodeTask = task("zarrEncode", async (data: Uint8Array, descriptor: CodecDescriptor) => { +export const encodeTask = task("zarrEncode", async (chunk: Chunk, descriptor: CodecDescriptor) => { const codec = await getCodec(descriptor); - const result = await codec.encode(data); + const result = await codec.encode(chunk); return { result, transfer: [result.buffer] }; }); export const decodeTask = task("zarrDecode", async (data: Uint8Array, descriptor: CodecDescriptor) => { const codec = await getCodec(descriptor); const result = await codec.decode(data); - return { result, transfer: [result.buffer] }; + return { result, transfer: [result.data.buffer] }; });