1/* eslint-disable @typescript-eslint/no-explicit-any */
2import ReferenceFileSystemClient from "./ReferenceFileSystemClient";
3import { ZarrFileSystemClient, ZMetaDataZArray } from "./RemoteH5FileLindi";
5const lindiDatasetDataLoader = async (o: {
6 client: ReferenceFileSystemClient | ZarrFileSystemClient;
7 path: string;
8 zarray: ZMetaDataZArray;
9 slice: [number, number][];
10 assertSingleChunkInFirstTwoDimensions?: boolean;
11 disableCache?: boolean;
12}) => {
13 const { client, zarray, path, slice, assertSingleChunkInFirstTwoDimensions } =
14 o;
16 const chunkShape = zarray.chunks;
17 const shape = zarray.shape;
18 const dtype = zarray.dtype;
19 if (!chunkShape) throw Error("No chunks shape for " + path);
20 if (!shape) throw Error("No shape for " + path);
21 if (!dtype) throw Error("No dtype for " + path);
22 const ndims = shape.length;
23 if (ndims !== chunkShape.length)
24 throw Error("Mismatched ndims and chunk shape for " + path);
26 if (o.slice.length === 3) {
27 // in this case we slice by two and then return the result of slicing by the third
28 const slice1 = slice.slice(0, 2);
29 const sN1 = slice1[0][1] - slice1[0][0];
30 const sN2 = slice1[1][1] - slice1[1][0];
31 const sN3 = o.slice[2][1] - o.slice[2][0];
32 const sNother = shape.slice(3).reduce((a, b) => a * b, 1);
33 const N3 = shape[2];
34 const xx = await lindiDatasetDataLoader({
35 client,
36 path,
37 zarray,
38 slice: slice1,
39 assertSingleChunkInFirstTwoDimensions,
40 disableCache: o.disableCache,
41 });
42 const xxRet = allocateArrayWithDtype(sN1 * sN2 * sN3 * sNother, dtype);
43 let iRet = 0;
44 for (let i1 = 0; i1 < sN1; i1++) {
45 for (let i2 = 0; i2 < sN2; i2++) {
46 for (let i3 = o.slice[2][0]; i3 < o.slice[2][1]; i3++) {
47 for (let i4 = 0; i4 < sNother; i4++) {
48 xxRet[iRet] = xx[i4 + sNother * (i3 + N3 * (i2 + sN2 * i1))];
49 iRet++;
50 }
51 }
52 }
53 }
54 return xxRet;
55 }
56 if (o.slice.length > 3) {
57 throw Error(
58 `For now, you can't slice more than three dimensions at a time. You tried to slice ${o.slice.length} dimensions for ${path}.`,
59 );
60 }
62 const macroChunkShape = chunkShape.map((cs, i) => Math.ceil(shape[i] / cs));
64 // check if we have a single chunk with no filters or compression (single contiguous block of data)
65 // It's important to handle this case specially because in this situation we don't need to download
66 // the entire chunk, we can just download the slice we need.
67 const singleChunk = macroChunkShape.reduce((a, b) => a * b, 1) === 1;
68 const noFiltersOrCompression =
69 !zarray.compressor && (!zarray.filters || zarray.filters.length === 0);
70 if (singleChunk && noFiltersOrCompression && slice && slice.length > 0) {
71 if (slice.length > 2) {
72 throw Error(
73 "For now, you can only slice two dimensions at a time for single chunk contiguous data",
74 );
75 }
76 const dtypeByteSize = getDtypeByteSize(dtype);
77 const startByte =
78 slice[0][0] * shape.slice(1).reduce((a, b) => a * b, 1) * dtypeByteSize;
79 const endByte =
80 slice[0][1] * shape.slice(1).reduce((a, b) => a * b, 1) * dtypeByteSize;
81 let singleChunkPath = path + "/0";
82 for (let i = 1; i < ndims; i++) {
83 singleChunkPath += ".0";
84 }
85 const dd = await client.readBinary(singleChunkPath, {
86 decodeArray: false,
87 startByte,
88 endByte,
89 });
90 let a = createDataView(dd, dtype);
91 if (slice.length === 2) {
92 if (shape.length === 1) {
93 if (slice[1][0] !== 0 || slice[1][1] !== 1) {
94 throw Error(
95 `For now, you can't slice the second dimension for single chunk contiguous data`,
96 );
97 }
98 return a;
99 }
100 const ss = shape.slice(2).reduce((a, b) => a * b, 1);
101 const newRet = allocateArrayWithDtype(
102 (slice[0][1] - slice[0][0]) * (slice[1][1] - slice[1][0]) * ss,
103 dtype,
104 );
105 let iRet = 0;
106 for (let i = 0; i < slice[0][1] - slice[0][0]; i++) {
107 for (let j = slice[1][0]; j < slice[1][1]; j++) {
108 for (let k = 0; k < ss; k++) {
109 newRet[iRet] = a[(i * shape[1] + j) * ss + k];
110 iRet++;
111 }
112 }
113 }
114 a = newRet as any;
115 }
116 return a;
117 }
119 const prodChunkSizeOfAllButFirstDimension = chunkShape
120 .slice(1)
121 .reduce((a, b) => a * b, 1);
122 const prodChunkSizeOfAllButFirstTwoDimensions = chunkShape
123 .slice(2)
124 .reduce((a, b) => a * b, 1);
125 const prodShapeSizeOfAllButFirstTwoDimensions = shape
126 .slice(2)
127 .reduce((a, b) => a * b, 1);
128 const prodMacroChunkShapeAllButFirstTwoDimensions = macroChunkShape
129 .slice(2)
130 .reduce((a, b) => a * b, 1);
132 let i1Start = 0;
133 let i1End = shape[0];
134 let i2Start = 0;
135 let i2End = ndims > 1 ? shape[1] : 1;
136 if (slice) {
137 if (slice.length >= 1) {
138 i1Start = slice[0][0];
139 i1End = slice[0][1];
140 }
141 if (slice.length >= 2) {
142 i2Start = slice[1][0];
143 i2End = slice[1][1];
144 }
145 if (slice.length > 2) {
146 throw Error(
147 `For now, you can't slice more than two dimensions at a time. You tried to slice ${slice.length} dimensions for ${path}.`,
148 );
149 }
150 }
151 if (i1End > shape[0]) i1End = shape[0];
152 if (i2End > shape[1]) i2End = shape[1];
154 const shape2 = ndims > 1 ? shape[1] : 1;
155 const chunkShape2 = ndims > 1 ? chunkShape[1] : 1;
157 if (i1Start < 0) {
158 throw Error(`Problem slicing ${path}: i1Start < 0: ${i1Start}`);
159 }
160 if (i1End > shape[0]) {
161 throw Error(
162 `Problem slicing ${path}: i1End > shape[0]: ${i1End} > ${shape[0]}`,
163 );
164 }
165 if (i2Start < 0) {
166 throw Error(`Problem slicing ${path}: i2Start < 0: ${i2Start}`);
167 }
168 if (i2End > shape2) {
169 throw Error(
170 `Problem slicing ${path}: i2End > shape[1]: ${i2End} > ${chunkShape2}`,
171 );
172 }
174 const i1StartChunk = Math.floor(i1Start / chunkShape[0]);
175 const i1EndChunk = Math.floor((i1End - 1) / chunkShape[0]);
176 const i2StartChunk = ndims > 1 ? Math.floor(i2Start / chunkShape[1]) : 0;
177 const i2EndChunk = ndims > 1 ? Math.floor((i2End - 1) / chunkShape[1]) : 0;
178 if (i1StartChunk === i1EndChunk && i2StartChunk === i2EndChunk) {
179 // With respect to the first two dimensions,
180 // we are entirely within a single chunk.
182 if (prodMacroChunkShapeAllButFirstTwoDimensions === 1) {
183 // in this case we are truly in a single chunk because there is only one chunk in the other dimensions
184 let chunkPath = path + "/" + i1StartChunk;
185 if (ndims > 1) {
186 chunkPath += "." + i2StartChunk;
187 }
188 for (let d = 2; d < ndims; d++) {
189 chunkPath += ".0";
190 }
191 const x = await client.readBinary(chunkPath, {
192 decodeArray: true,
193 disableCache: o.disableCache,
194 });
195 if (!x) {
196 console.log({
197 i1StartChunk,
198 i1EndChunk,
199 i2StartChunk,
200 i2EndChunk,
201 i1Start,
202 i1End,
203 i2Start,
204 i2End,
205 shape,
206 chunkShape,
207 });
208 throw Error("Unable to read chunk: " + chunkPath);
209 }
210 const j1Start = i1Start - i1StartChunk * chunkShape[0];
211 const j1End = i1End - i1StartChunk * chunkShape[0];
212 const j2Start = i2Start - i2StartChunk * chunkShape2;
213 const j2End = i2End - i2StartChunk * chunkShape2;
214 const slicingInSecondDimension =
215 ndims > 1 && (j2Start > 0 || j2End < chunkShape2);
216 if (!slicingInSecondDimension) {
217 // we are not slicing in second dimension. In this case we don't need to make a copy of the data
218 const ret = x.slice(
219 j1Start * prodChunkSizeOfAllButFirstDimension,
220 j1End * prodChunkSizeOfAllButFirstDimension,
221 );
222 return ret;
223 } else {
224 // we are slicing in second dimension, so we need to make a copy of the data
225 const ret = allocateArrayWithDtype(
226 (i1End - i1Start) *
227 (i2End - i2Start) *
228 prodShapeSizeOfAllButFirstTwoDimensions,
229 dtype,
230 );
231 let iRet = 0;
232 for (let j1 = j1Start; j1 < j1End; j1++) {
233 for (let j2 = j2Start; j2 < j2End; j2++) {
234 for (
235 let j3 = 0;
236 j3 < prodShapeSizeOfAllButFirstTwoDimensions;
237 j3++
238 ) {
239 ret[iRet] =
240 x[
241 (j1 * chunkShape2 + j2) *
242 prodChunkSizeOfAllButFirstTwoDimensions +
243 j3
244 ];
245 iRet++;
246 }
247 }
248 }
249 return ret;
250 }
251 } else {
252 // there is more than one chunk in the other dimensions, and we need to concatenate them
253 if (ndims > 4) {
254 throw Error("Case not yet supported: C2");
255 }
256 while (macroChunkShape.length < 4) {
257 macroChunkShape.push(1);
258 }
259 const retList = [];
260 for (let iii3 = 0; iii3 < macroChunkShape[2]; iii3++) {
261 for (let iii4 = 0; iii4 < macroChunkShape[3]; iii4++) {
262 let chunkPath = path + "/" + i1StartChunk;
263 chunkPath += "." + i2StartChunk;
264 chunkPath += "." + iii3;
265 if (ndims === 4) {
266 chunkPath += "." + iii4;
267 }
268 const x = await client.readBinary(chunkPath, {
269 decodeArray: true,
270 disableCache: o.disableCache,
271 });
272 if (!x) {
273 console.log({
274 i1StartChunk,
275 i1EndChunk,
276 i2StartChunk,
277 i2EndChunk,
278 i1Start,
279 i1End,
280 i2Start,
281 i2End,
282 shape,
283 chunkShape,
284 });
285 throw Error("Unable to read chunk: " + chunkPath);
286 }
287 const j1Start = i1Start - i1StartChunk * chunkShape[0];
288 const j1End = i1End - i1StartChunk * chunkShape[0];
289 const j2Start = i2Start - i2StartChunk * chunkShape2;
290 const j2End = i2End - i2StartChunk * chunkShape2;
291 const slicingInSecondDimension =
292 ndims > 1 && (j2Start > 0 || j2End < chunkShape2);
293 if (!slicingInSecondDimension) {
294 // we are not slicing in second dimension. In this case we don't need to make a copy of the data
295 const ret0 = x.slice(
296 j1Start * prodChunkSizeOfAllButFirstDimension,
297 j1End * prodChunkSizeOfAllButFirstDimension,
298 );
299 retList.push(ret0);
300 } else {
301 // we are slicing in second dimension, so we need to make a copy of the data
302 const ret0 = allocateArrayWithDtype(
303 (i1End - i1Start) *
304 (i2End - i2Start) *
305 prodChunkSizeOfAllButFirstTwoDimensions,
306 dtype,
307 );
308 let iRet0 = 0;
309 for (let j1 = j1Start; j1 < j1End; j1++) {
310 for (let j2 = j2Start; j2 < j2End; j2++) {
311 for (
312 let j3 = 0;
313 j3 < prodChunkSizeOfAllButFirstTwoDimensions;
314 j3++
315 ) {
316 ret0[iRet0] =
317 x[
318 (j1 * chunkShape2 + j2) *
319 prodChunkSizeOfAllButFirstTwoDimensions +
320 j3
321 ];
322 iRet0++;
323 }
324 }
325 }
326 retList.push(ret0);
327 }
328 }
329 }
330 // now concatenate the ret0s
331 const ret = allocateArrayWithDtype(
332 (i1End - i1Start) *
333 (i2End - i2Start) *
334 prodShapeSizeOfAllButFirstTwoDimensions,
335 dtype,
336 );
337 let iRet = 0;
338 for (let i1 = 0; i1 < i1End - i1Start; i1++) {
339 for (let i2 = 0; i2 < i2End - i2Start; i2++) {
340 for (let i = 0; i < retList.length; i++) {
341 for (
342 let i3 = 0;
343 i3 < prodChunkSizeOfAllButFirstTwoDimensions;
344 i3++
345 ) {
346 ret[iRet] =
347 retList[i][
348 i1 *
349 (i2End - i2Start) *
350 prodChunkSizeOfAllButFirstTwoDimensions +
351 i2 * prodChunkSizeOfAllButFirstTwoDimensions +
352 i3
353 ];
354 iRet++;
355 }
356 }
357 }
358 }
359 return ret;
360 }
361 }
363 if (assertSingleChunkInFirstTwoDimensions)
364 throw Error(
365 "Unexpected case. We should have handled all cases by now (assertSingleChunkInFirstTwoDimensions)",
366 );
368 const ret = allocateArrayWithDtype(
369 (i1End - i1Start) *
370 (i2End - i2Start) *
371 prodShapeSizeOfAllButFirstTwoDimensions,
372 dtype,
373 );
375 const handleChunk = async (o2: {
376 slice1: [number, number];
377 slice2: [number, number];
378 }) => {
379 const { slice1, slice2 } = o2;
380 const sliceA = [slice1, slice2];
381 const xx = await lindiDatasetDataLoader({
382 client,
383 path,
384 zarray,
385 slice: sliceA,
386 assertSingleChunkInFirstTwoDimensions: true, // avoid infinite recursion by accident
387 disableCache: o.disableCache,
388 });
389 let iXX = 0;
390 for (let ii1 = slice1[0]; ii1 < slice1[1]; ii1++) {
391 for (let ii2 = slice2[0]; ii2 < slice2[1]; ii2++) {
392 let iRet =
393 (ii1 - i1Start) *
394 (i2End - i2Start) *
395 prodShapeSizeOfAllButFirstTwoDimensions +
396 (ii2 - i2Start) * prodShapeSizeOfAllButFirstTwoDimensions;
397 for (
398 let ii3 = 0;
399 ii3 < prodShapeSizeOfAllButFirstTwoDimensions;
400 ii3++
401 ) {
402 ret[iRet] = xx[iXX];
403 iRet++;
404 iXX++;
405 }
406 }
407 }
408 };
410 const promises: Promise<void>[] = [];
411 for (let i1Chunk = i1StartChunk; i1Chunk <= i1EndChunk; i1Chunk++) {
412 let slice1: [number, number];
413 if (i1Chunk === i1StartChunk && i1Chunk === i1EndChunk) {
414 slice1 = [i1Start, i1End];
415 } else if (i1Chunk === i1StartChunk) {
416 slice1 = [i1Start, (i1Chunk + 1) * chunkShape[0]];
417 } else if (i1Chunk === i1EndChunk) {
418 slice1 = [i1Chunk * chunkShape[0], i1End];
419 } else {
420 slice1 = [i1Chunk * chunkShape[0], (i1Chunk + 1) * chunkShape[0]];
421 }
422 for (let i2Chunk = i2StartChunk; i2Chunk <= i2EndChunk; i2Chunk++) {
423 let slice2: [number, number];
424 if (i2Chunk === i2StartChunk && i2Chunk === i2EndChunk) {
425 slice2 = [i2Start, i2End];
426 } else if (i2Chunk === i2StartChunk) {
427 slice2 = [i2Start, (i2Chunk + 1) * chunkShape2];
428 } else if (i2Chunk === i2EndChunk) {
429 slice2 = [i2Chunk * chunkShape2, i2End];
430 } else {
431 slice2 = [i2Chunk * chunkShape2, (i2Chunk + 1) * chunkShape2];
432 }
434 promises.push(handleChunk({ slice1, slice2 }));
435 }
436 }
437 await Promise.all(promises);
438 return ret;
439};
441const allocateArrayWithDtype = (size: number, dtype: string) => {
442 if (dtype === "<f4") return new Float32Array(size);
443 if (dtype === "<f8") return new Float64Array(size);
444 if (dtype === "<i1" || dtype === "|i1") return new Int8Array(size);
445 if (dtype === "<i2") return new Int16Array(size);
446 if (dtype === "<i4") return new Int32Array(size);
447 if (dtype === "<i8") {
448 const a = new BigInt64Array(size);
449 // convert to regular Int32Array because js has trouble mixing BigInt64Array with other numbers
450 const ret = new Int32Array(size);
451 for (let i = 0; i < size; i++) {
452 ret[i] = Number(a[i]);
453 }
454 return ret;
455 }
456 if (dtype === "<u1" || dtype === "|u1") return new Uint8Array(size);
457 if (dtype === "<u2") return new Uint16Array(size);
458 if (dtype === "<u4") return new Uint32Array(size);
459 if (dtype === "<u8") {
460 const a = new BigUint64Array(size);
461 // convert to regular Uint32Array because js has trouble mixing BigUint64Array with other numbers
462 const ret = new Uint32Array(size);
463 for (let i = 0; i < size; i++) {
464 ret[i] = Number(a[i]);
465 }
466 return ret;
467 }
468 if (dtype === "|O") return new Array(size);
469 throw Error(`Unsupported dtype: ${dtype}`);
470};
472const createDataView = (dd: ArrayBuffer, dtype: string) => {
473 if (dtype === "<f4") return new Float32Array(dd);
474 if (dtype === "<f8") return new Float64Array(dd);
475 if (dtype === "<i1" || dtype === "|i1") return new Int8Array(dd);
476 if (dtype === "<i2") return new Int16Array(dd);
477 if (dtype === "<i4") return new Int32Array(dd);
478 if (dtype === "<i8") return new BigInt64Array(dd);
479 if (dtype === "<u1" || dtype === "|u1") return new Uint8Array(dd);
480 if (dtype === "<u2") return new Uint16Array(dd);
481 if (dtype === "<u4") return new Uint32Array(dd);
482 if (dtype === "<u8") return new BigUint64Array(dd);
483 throw Error(`Unsupported dtype: ${dtype}`);
484};
486const getDtypeByteSize = (dtype: string) => {
487 if (dtype === "<f4") return 4;
488 if (dtype === "<f8") return 8;
489 if (dtype === "<i1" || dtype === "|i1") return 1;
490 if (dtype === "<i2") return 2;
491 if (dtype === "<i4") return 4;
492 if (dtype === "<i8") return 8;
493 if (dtype === "<u1" || dtype === "|u1") return 1;
494 if (dtype === "<u2") return 2;
495 if (dtype === "<u4") return 4;
496 if (dtype === "<u8") return 8;
497 throw Error(`Unsupported dtype: ${dtype}`);
498};
500export default lindiDatasetDataLoader;