/ concept-collection / remote-hdf5-lazy-read
Sign in
concept-collection / remote-hdf5-lazy-read
remote-hdf5-lazy-read / src / remote-h5-file / lib / lindi / RemoteH5FileLindi.ts
598 lines · 18.1 KBBlameHistoryRaw
1/* eslint-disable @typescript-eslint/no-explicit-any */
2import {
3 DatasetDataType,
4 RemoteH5Dataset,
5 RemoteH5Group,
6 RemoteH5Subdataset,
7 RemoteH5Subgroup,
8 // getRemoteH5File,
9 globalRemoteH5FileStats,
10} from "../RemoteH5File";
11// import { Canceler } from "../helpers";
13type Canceler = {
14 onCancel: (() => void)[];
15};
17import ReferenceFileSystemClient, {
18 ReferenceFileSystemObject,
19 RemoteTarInterface,
20 isReferenceFileSystemObject,
21} from "./ReferenceFileSystemClient";
22import lindiDatasetDataLoader from "./lindiDatasetDataLoader";
23import zarrDecodeChunkArray from "./zarrDecodeChunkArray";
25type ZMetaDataZAttrs = { [key: string]: any };
27type ZMetaDataZGroup = {
28 zarr_format: number;
29};
31export type ZMetaDataZArray = {
32 chunks?: number[];
33 compressor?: any;
34 dtype?: string;
35 fill_value?: any;
36 filters?: any[];
37 order?: "C" | "F";
38 shape?: number[];
39 zarr_format?: 2;
40};
42export class ZarrFileSystemClient {
43 #fileContentCache: {
44 [key: string]: { content: any | undefined; found: boolean };
45 } = {};
46 #inProgressReads: { [key: string]: boolean } = {};
47 constructor(
48 private url: string,
49 private zmetadata: any,
50 ) {}
51 async readJson(path: string): Promise<{ [key: string]: any } | undefined> {
52 if (path in this.zmetadata.metadata) {
53 return this.zmetadata.metadata[path];
54 }
55 const lastPartOfPath = path.split("/").slice(-1)[0];
56 if (lastPartOfPath.startsWith(".")) {
57 // if it's not in the metadata, we assume it's not there
58 return undefined;
59 }
60 const buf = await this.readBinary(path, { decodeArray: false });
61 if (!buf) return undefined;
62 const text = new TextDecoder().decode(buf);
63 try {
64 return JSON.parse(text, (_key, value) => {
65 if (value === "___NaN___") return NaN;
66 return value;
67 });
68 } catch (e) {
69 console.warn(text);
70 throw Error("Failed to parse JSON for " + path + ": " + e);
71 }
72 }
73 async readBinary(
74 path: string,
75 o: {
76 decodeArray?: boolean;
77 startByte?: number;
78 endByte?: number;
79 disableCache?: boolean;
80 },
81 ): Promise<any | undefined> {
82 if (o.startByte !== undefined) {
83 if (o.decodeArray)
84 throw Error("Cannot decode array and read a slice at the same time");
85 if (o.endByte === undefined)
86 throw Error("If you specify startByte, you must also specify endByte");
87 } else if (o.endByte !== undefined) {
88 throw Error("If you specify endByte, you must also specify startByte");
89 }
90 if (
91 o.endByte !== undefined &&
92 o.startByte !== undefined &&
93 o.endByte < o.startByte
94 ) {
95 throw Error(
96 `endByte must be greater than or equal to startByte: ${o.startByte} ${o.endByte} for ${path}`,
97 );
98 }
99 if (
100 o.endByte !== undefined &&
101 o.startByte !== undefined &&
102 o.endByte === o.startByte
103 ) {
104 return new ArrayBuffer(0);
105 }
106 const kk =
107 path +
108 "|" +
109 (o.decodeArray ? "decode" : "") +
110 "|" +
111 o.startByte +
112 "|" +
113 o.endByte;
114 while (this.#inProgressReads[kk]) {
115 await new Promise((resolve) => setTimeout(resolve, 100));
116 }
117 this.#inProgressReads[kk] = true;
118 try {
119 if (path.startsWith("/")) path = path.slice(1);
120 if (this.#fileContentCache[kk]) {
121 if (this.#fileContentCache[kk].found) {
122 return this.#fileContentCache[kk].content;
123 }
124 return undefined;
125 }
126 const url = this.url + "/" + path;
127 let buf: ArrayBuffer | undefined;
128 if (o.startByte !== undefined && o.endByte !== undefined) {
129 buf = await fetchByteRange(url, o.startByte, o.endByte - o.startByte);
130 } else {
131 const r = await fetch(url);
132 if (!r.ok) {
133 if (r.status === 404) {
134 this.#fileContentCache[kk] = { content: undefined, found: false };
135 return undefined; // file not found
136 }
137 throw Error(`Failed to fetch ${url}: ${r.statusText}`);
138 }
139 buf = await r.arrayBuffer();
140 }
141 if (o.decodeArray) {
142 const parentPath = path.split("/").slice(0, -1).join("/");
143 const zarray = (await this.readJson(parentPath + "/.zarray")) as
144 | ZMetaDataZArray
145 | undefined;
146 if (!zarray) throw Error("Failed to read .zarray for " + path);
147 try {
148 buf = await zarrDecodeChunkArray(
149 buf,
150 zarray.dtype,
151 zarray.compressor,
152 zarray.filters,
153 zarray.chunks,
154 );
155 } catch (e) {
156 throw Error(`Failed to decode chunk array for ${path}: ${e}`);
157 }
158 }
159 if (buf) {
160 this.#fileContentCache[kk] = { content: buf, found: true };
161 } else {
162 this.#fileContentCache[kk] = { content: undefined, found: false };
163 }
164 return buf;
165 } catch (e) {
166 this.#fileContentCache[kk] = { content: undefined, found: false }; // important to do this so we don't keep trying to read the same file
167 throw e;
168 } finally {
169 this.#inProgressReads[kk] = false;
170 }
171 }
174class RemoteH5FileLindi {
175 #cacheDisabled = false; // just for benchmarking
176 #sourceUrls: string[] | undefined = undefined;
177 constructor(
178 public url: string,
179 private lindiFileSystemClient:
180 | ReferenceFileSystemClient
181 | ZarrFileSystemClient,
182 private pathsByParentPath: { [key: string]: string[] },
183 ) {}
184 static async create(url: string) {
185 const { rfs: obj, remoteTar } = await fetchRfsFromRemoteLindi(url);
186 // console.info(`reference file system for ${url}`, obj);
187 // console.info(`Meta only`, metaOnly(obj));
188 const pathsByParentPath: { [key: string]: string[] } = {};
189 for (const path in obj.refs) {
190 if (path === ".zattrs" || path === ".zgroup") continue;
191 const parts = path.split("/");
192 if (parts.length <= 1) continue;
193 const lastPart = parts[parts.length - 1];
194 if (
195 lastPart === ".zattrs" ||
196 lastPart === ".zgroup" ||
197 lastPart === ".zarray"
198 ) {
199 const thePath = parts.slice(0, parts.length - 1).join("/");
200 const theParentPath = parts.slice(0, parts.length - 2).join("/");
201 if (!pathsByParentPath[theParentPath])
202 pathsByParentPath[theParentPath] = [];
203 if (!pathsByParentPath[theParentPath].includes(thePath)) {
204 pathsByParentPath[theParentPath].push(thePath);
205 }
206 }
207 }
208 return new RemoteH5FileLindi(
209 url,
210 new ReferenceFileSystemClient(obj, remoteTar),
211 pathsByParentPath,
212 );
213 }
214 static async createFromZarr(url: string) {
215 const zmetadataUrl = `${url}/.zmetadata`;
216 const zmetadataResponse = await fetch(zmetadataUrl);
217 if (!zmetadataResponse.ok) {
218 throw new Error(`Failed to fetch Zarr metadata from ${zmetadataUrl}`);
219 }
220 const zmetadata = await zmetadataResponse.json();
221 const zarrFileSystemClient = new ZarrFileSystemClient(url, zmetadata);
222 return new RemoteH5FileLindi(url, zarrFileSystemClient, {});
223 }
224 get dataIsRemote() {
225 return !this.url.startsWith("http://localhost");
226 }
227 async getGroup(path: string): Promise<RemoteH5Group | undefined> {
228 if (path === "") path = "/";
229 let group: RemoteH5Group | undefined;
230 const pathWithoutBeginningSlash = path.startsWith("/")
231 ? path.slice(1)
232 : path;
233 let zgroup: ZMetaDataZGroup | undefined;
234 let zattrs: ZMetaDataZAttrs | undefined;
235 if (path === "/") {
236 zgroup = (await this.lindiFileSystemClient.readJson(".zgroup")) as
237 | ZMetaDataZGroup
238 | undefined;
239 zattrs = (await this.lindiFileSystemClient.readJson(".zattrs")) as
240 | ZMetaDataZAttrs
241 | undefined;
242 } else {
243 zgroup = (await this.lindiFileSystemClient.readJson(
244 pathWithoutBeginningSlash + "/.zgroup",
245 )) as ZMetaDataZGroup | undefined;
246 zattrs = (await this.lindiFileSystemClient.readJson(
247 pathWithoutBeginningSlash + "/.zattrs",
248 )) as ZMetaDataZAttrs | undefined;
249 }
250 if (zgroup) {
251 const subgroups: RemoteH5Subgroup[] = [];
252 const subdatasets: RemoteH5Subdataset[] = [];
253 const childPaths: string[] =
254 this.pathsByParentPath[pathWithoutBeginningSlash] || [];
255 for (const childPath of childPaths) {
256 const childZgroup = await this.lindiFileSystemClient.readJson(
257 childPath + "/.zgroup",
258 );
259 const childZarray = await this.lindiFileSystemClient.readJson(
260 childPath + "/.zarray",
261 );
262 const childZattrs = await this.lindiFileSystemClient.readJson(
263 childPath + "/.zattrs",
264 );
265 if (childZgroup) {
266 subgroups.push({
267 name: getNameFromPath(childPath),
268 path: "/" + childPath,
269 attrs: childZattrs || {},
270 });
271 } else if (childZarray) {
272 const shape = childZarray.shape;
273 const dtype = childZarray.dtype;
274 if (shape && dtype) {
275 subdatasets.push({
276 name: getNameFromPath(childPath),
277 path: "/" + childPath,
278 shape,
279 dtype,
280 attrs: childZattrs || {},
281 chunks: childZarray.chunks,
282 compressor: formatCompressor(childZarray.compressor),
283 filters: formatFilters(childZarray.filters),
284 });
285 } else {
286 console.warn("Unexpected .zarray item", childPath, childZarray);
287 }
288 }
289 }
290 group = {
291 path: path,
292 subgroups,
293 datasets: subdatasets,
294 attrs: zattrs || {},
295 };
296 }
297 globalRemoteH5FileStats.getGroupCount++;
298 return group;
299 }
300 async getDataset(path: string): Promise<RemoteH5Dataset | undefined> {
301 const pathWithoutBeginningSlash = path.startsWith("/")
302 ? path.slice(1)
303 : path;
304 const zarray = (await this.lindiFileSystemClient.readJson(
305 pathWithoutBeginningSlash + "/.zarray",
306 )) as ZMetaDataZArray;
307 const zattrs = (await this.lindiFileSystemClient.readJson(
308 pathWithoutBeginningSlash + "/.zattrs",
309 )) as ZMetaDataZAttrs;
310 let dataset: RemoteH5Dataset | undefined;
311 if (zarray) {
312 dataset = {
313 name: getNameFromPath(path),
314 path,
315 shape: zarray.shape || [],
316 dtype: zarray.dtype || "",
317 attrs: zattrs || {},
318 chunks: zarray.chunks,
319 compressor: formatCompressor(zarray.compressor),
320 filters: formatFilters(zarray.filters),
321 };
322 } else {
323 dataset = undefined;
324 }
325 globalRemoteH5FileStats.getDatasetCount++;
326 return dataset;
327 }
328 async getDatasetData(
329 path: string,
330 o: {
331 slice?: [number, number][];
332 allowBigInt?: boolean;
333 canceler?: Canceler;
334 },
335 ): Promise<DatasetDataType | undefined> {
336 // check for invalid slice
337 if (o.slice) {
338 for (const ss of o.slice) {
339 if (isNaN(ss[0]) || isNaN(ss[1])) {
340 console.warn("Invalid slice", path, o.slice);
341 throw Error("Invalid slice");
342 }
343 }
344 }
345 if (o.slice && o.slice.length > 3) {
346 console.warn(
347 "Tried to slice more than three dimensions at a time",
348 path,
349 o.slice,
350 );
351 throw Error(
352 `For now, you can't slice more than three dimensions at a time. You tried to slice ${o.slice.length} dimensions for ${path}.`,
353 );
354 }
356 const pathWithoutBeginningSlash = path.startsWith("/")
357 ? path.slice(1)
358 : path;
359 const zarray = (await this.lindiFileSystemClient.readJson(
360 pathWithoutBeginningSlash + "/.zarray",
361 )) as ZMetaDataZArray | undefined;
362 if (!zarray) {
363 console.warn("No .zarray for", path);
364 return undefined;
365 }
367 // const { slice, allowBigInt, canceler } = o;
369 globalRemoteH5FileStats.getDatasetDataCount++;
371 // old system (not used by lindi)
372 const externalHdf5 = await this.lindiFileSystemClient.readJson(
373 pathWithoutBeginningSlash + "/.external_hdf5",
374 );
375 if (externalHdf5) {
376 throw Error("External hdf5 not supported on server side");
377 // const a = await getRemoteH5File(externalHdf5.url);
378 // return a.getDatasetData(externalHdf5.name, o);
379 }
381 const zattrs = (await this.lindiFileSystemClient.readJson(
382 pathWithoutBeginningSlash + "/.zattrs",
383 )) as ZMetaDataZAttrs;
384 if (zattrs && zattrs["_EXTERNAL_ARRAY_LINK"]) {
385 throw Error("External array link not supported on server side");
386 // const externalArrayLink = zattrs["_EXTERNAL_ARRAY_LINK"];
387 // let url0 = externalArrayLink.url;
388 // if (this.#cacheDisabled) {
389 // url0 += `?cacheBust=${Date.now()}`;
390 // }
391 // const a = await getRemoteH5File(url0);
392 // return a.getDatasetData(externalArrayLink.name, o);
393 }
395 const ret = await lindiDatasetDataLoader({
396 client: this.lindiFileSystemClient,
397 path: pathWithoutBeginningSlash,
398 zarray,
399 slice: o.slice || [],
400 disableCache: this.#cacheDisabled,
401 });
402 if (ret.length === 1) {
403 // candidate for scalar, need to check for _SCALAR attribute
404 const ds = await this.getDataset(path);
405 if (ds && ds.attrs["_SCALAR"]) {
406 return ret[0];
407 }
408 }
409 return ret;
410 }
411 get _lindiFileSystemClient() {
412 return this.lindiFileSystemClient;
413 }
414 async getLindiZarray(path: string): Promise<ZMetaDataZArray | undefined> {
415 const pathWithoutBeginningSlash = path.startsWith("/")
416 ? path.slice(1)
417 : path;
418 return (await this.lindiFileSystemClient.readJson(
419 pathWithoutBeginningSlash + "/.zarray",
420 )) as ZMetaDataZArray | undefined;
421 }
422 getUrls() {
423 return [this.url];
424 }
425 get sourceUrls(): string[] | undefined {
426 return this.#sourceUrls;
427 }
428 set sourceUrls(v: string[] | undefined) {
429 this.#sourceUrls = v;
430 }
431 _disableCache() {
432 this.#cacheDisabled = true;
433 }
436const fetchRfsFromRemoteLindi = async (
437 url: string,
438): Promise<{
439 rfs: ReferenceFileSystemObject;
440 remoteTar: RemoteTarInterface | undefined;
441}> => {
442 const buf: ArrayBuffer = await fetchByteRange(url, 0, 512 * 3);
443 if (isTarHeader(buf.slice(0, 512))) {
444 const tarEntryBuf = buf.slice(512, 512 + 1024);
445 const tarEntryJson = new TextDecoder().decode(tarEntryBuf);
446 const tarEntry = JSON.parse(tarEntryJson);
447 const indexInfo = tarEntry["index"];
448 const entryDataStartByte = indexInfo["d"];
449 const entryDataSize = indexInfo["s"];
451 const indexBuf = await fetchByteRange(
452 url,
453 entryDataStartByte,
454 entryDataSize,
455 );
456 const indexStr = new TextDecoder().decode(indexBuf);
457 const index = JSON.parse(indexStr);
458 const remoteTar: RemoteTarInterface = {
459 url: url,
460 getByteRangeForFile: async (fileName: string) => {
461 const f = index.files.find((ff: any) => ff.n === fileName);
462 if (!f) {
463 throw Error(`File ${fileName} not found in tar`);
464 }
465 return {
466 startByte: f.d as number,
467 endByte: (f.d + f.s) as number,
468 };
469 },
470 };
471 const { startByte: rfsStartByte, endByte: rfsEndByte } =
472 await remoteTar.getByteRangeForFile("lindi.json");
473 const rfsBuf = await fetchByteRange(
474 url,
475 rfsStartByte,
476 rfsEndByte - rfsStartByte,
477 );
478 const rfs = JSON.parse(new TextDecoder().decode(rfsBuf));
479 if (!isReferenceFileSystemObject(rfs)) {
480 console.warn(rfs);
481 throw Error("Invalid rfs from tar");
482 }
483 return {
484 rfs,
485 remoteTar,
486 };
487 } else {
488 const r = await fetch(url);
489 if (!r.ok) throw Error("Failed to fetch LINDI file" + url);
490 const rfs = await r.json();
491 if (!isReferenceFileSystemObject(rfs)) {
492 console.warn(rfs);
493 throw Error("Invalid rfs");
494 }
495 return {
496 rfs,
497 remoteTar: undefined,
498 };
499 }
500};
502const fetchByteRange = async (url: string, startByte: number, size: number) => {
503 const r = await fetch(url, {
504 headers: {
505 Range: `bytes=${startByte}-${startByte + size - 1}`,
506 },
507 });
508 if (!r.ok)
509 throw Error(
510 `Failed to fetch byte range ${startByte}-${startByte + size - 1} of ${url}`,
511 );
512 return await r.arrayBuffer();
513};
515const isTarHeader = (buf: ArrayBuffer) => {
516 if (buf.byteLength < 512) {
517 return false;
518 }
520 // We're only going to support ustar format
521 // get the ustar indicator at bytes 257-262
522 const ustarIndicator = buf.slice(257, 262);
523 const ustarIndicatorStr = new TextDecoder().decode(
524 ustarIndicator.slice(0, 5),
525 );
526 const bb = new Uint8Array(buf);
527 if (ustarIndicatorStr === "ustar" && bb[257 + 5] == 0) {
528 return true;
529 }
531 // Check for any 0 bytes in the header
532 const bb2 = new Uint8Array(buf);
533 if (bb2.includes(0)) {
534 console.warn(ustarIndicatorStr);
535 throw Error(
536 "Problem with lindi file: 0 byte found in header, but not ustar tar format",
537 );
538 }
540 return false;
541};
543const getNameFromPath = (path: string) => {
544 const parts = path.split("/");
545 if (parts.length === 0) return "";
546 return parts[parts.length - 1];
547};
549const lock1: { locked: boolean } = { locked: false };
550const globalLindiRemoteH5Files: { [url: string]: RemoteH5FileLindi } = {};
551export const getRemoteH5FileLindi = async (url: string) => {
552 while (lock1.locked) await new Promise((resolve) => setTimeout(resolve, 100));
553 try {
554 lock1.locked = true;
555 const kk = url;
556 if (!globalLindiRemoteH5Files[kk]) {
557 globalLindiRemoteH5Files[kk] = await RemoteH5FileLindi.create(url);
558 }
559 return globalLindiRemoteH5Files[kk];
560 } finally {
561 lock1.locked = false;
562 }
563};
565// const metaOnly = (obj: ReferenceFileSystemObject) => {
566// const ret = {
567// refs: {} as any,
568// version: obj.version,
569// };
570// for (const k in obj.refs) {
571// if (
572// k.endsWith(".zattrs") ||
573// k.endsWith(".zgroup") ||
574// k.endsWith(".zarray")
575// ) {
576// ret.refs[k] = obj.refs[k];
577// }
578// }
579// return ret;
580// };
582const formatCompressor = (compressor: any): string | undefined => {
583 if (!compressor) return undefined;
584 if (typeof compressor === "string") return compressor;
585 if (compressor.id) return compressor.id;
586 return JSON.stringify(compressor);
587};
589const formatFilters = (filters: any[] | undefined): string[] | undefined => {
590 if (!filters || filters.length === 0) return undefined;
591 return filters.map((f) => {
592 if (typeof f === "string") return f;
593 if (f.id) return f.id;
594 return JSON.stringify(f);
595 });
596};
598export default RemoteH5FileLindi;
moveopenescclose