1import { Blosc } from "numcodecs";
2import pako from "pako";
3import { qfcDecompress } from "./qfc";
5/* eslint-disable @typescript-eslint/no-explicit-any */
6const zarrDecodeChunkArray = async (
7 chunk: ArrayBuffer,
8 dtype?: string,
9 compressor?: any,
10 filters?: any[],
11 shape?: number[],
12): Promise<any> => {
13 let ret: any = chunk;
14 if (compressor) {
15 if (compressor.id === "blosc") {
16 ret = await new Blosc().decode(chunk);
17 } else if (compressor.id === "gzip") {
18 ret = pako.inflate(chunk);
19 } else if (compressor.id === "neurosift.mp4") {
20 // ret = await decodeMp4(chunk, shape![0], shape![1], shape![2]);
21 throw Error("neurosift.mp4 decoder not yet implemented");
22 } else if (compressor.id === "qfc") {
23 if (!shape) {
24 throw Error("No shape for qfc");
25 }
26 ret = await qfcDecompress(chunk, shape, compressor);
27 } else {
28 throw Error("Unhandled compressor " + compressor.id);
29 }
30 }
31 // check if Uint8Array
32 if (ret instanceof Uint8Array) {
33 ret = ret.buffer;
34 }
35 if (dtype === "|O") {
36 if (!shape) throw Error("No shape for |O");
37 if (!filters) {
38 throw Error("No filters for |O");
39 }
40 if (filters.length === 0) {
41 throw Error("No filters for |O");
42 }
43 // work backwards through the filters, besides the first one which should be json2
44 for (let i = filters.length - 1; i > 0; i--) {
45 ret = await applyFilter(ret, filters[i]);
46 }
47 const filter0 = filters[0];
48 if (filter0.id !== "json2") {
49 throw Error("First filter for |O should be json2");
50 }
51 ret = await applyFilterToOType(ret, filter0, shape);
52 } else {
53 // work our way backward through the filters
54 if (filters) {
55 for (let i = filters.length - 1; i >= 0; i--) {
56 ret = await applyFilter(ret, filters[i]);
57 }
58 }
59 if (!dtype) {
60 // pass
61 } else if (dtype === "<f4") {
62 ret = new Float32Array(ret);
63 } else if (dtype === "<f8") {
64 ret = new Float64Array(ret);
65 } else if (dtype === "<i1" || dtype === "|i1") {
66 ret = new Int8Array(ret);
67 } else if (dtype === "<i2") {
68 ret = new Int16Array(ret);
69 } else if (dtype === "<i4") {
70 ret = new Int32Array(ret);
71 } else if (dtype === "<i8") {
72 const ret0 = new BigInt64Array(ret);
73 // convert to Int32Array because javascript has trouble mixing BigInt64Array with other types
74 ret = new Int32Array(ret0.length);
75 for (let i = 0; i < ret0.length; i++) {
76 ret[i] = Number(ret0[i]);
77 }
78 } else if (dtype === "<u1" || dtype === "|u1") {
79 ret = new Uint8Array(ret);
80 } else if (dtype === "<u2") {
81 ret = new Uint16Array(ret);
82 } else if (dtype === "<u4") {
83 ret = new Uint32Array(ret);
84 } else if (dtype === "<u8") {
85 const ret0 = new BigUint64Array(ret);
86 // convert to Uint32Array because javascript has trouble mixing BigUint64Array with other types
87 ret = new Uint32Array(ret0.length);
88 for (let i = 0; i < ret0.length; i++) {
89 ret[i] = Number(ret0[i]);
90 }
91 } else if (dtype === "|b1") {
92 ret = new Uint8Array(ret);
93 } else if (dtype.startsWith("<U")) {
94 const fixedLength = parseInt(dtype.slice(2));
95 const nn = ret.byteLength / (fixedLength * 4);
96 const ret2 = [];
97 for (let i = 0; i < nn; i++) {
98 const ret1 = new Uint32Array(ret, i * fixedLength * 4, fixedLength);
99 const ret3 = [];
100 for (let j = 0; j < fixedLength; j++) {
101 if (ret1[j] === 0) break; // null terminates the string, matching NumPy behavior
102 ret3.push(String.fromCodePoint(ret1[j]));
103 }
104 ret2.push(ret3.join(""));
105 }
106 ret = ret2;
107 } else if (dtype.startsWith("|S")) {
108 const fixedLength = parseInt(dtype.slice(2));
109 const nn = ret.byteLength / fixedLength;
110 const ret2 = [];
111 for (let i = 0; i < nn; i++) {
112 const ret1 = new Uint8Array(ret, i * fixedLength, fixedLength);
113 // Find first null byte to trim, matching NumPy behavior
114 let len = fixedLength;
115 for (let j = 0; j < fixedLength; j++) {
116 if (ret1[j] === 0) {
117 len = j;
118 break;
119 }
120 }
121 ret2.push(new TextDecoder().decode(ret1.subarray(0, len)));
122 }
123 ret = ret2;
124 } else {
125 throw Error("Unhandled dtype " + dtype);
126 }
127 }
128 return ret;
129};
131const applyFilterToOType = async (
132 chunk: ArrayBuffer,
133 filter: any,
134 shape: number[],
135) => {
136 if (filter.id === "vlen-utf8") {
137 const view = new DataView(chunk);
138 const ret = [];
139 let i = 4;
140 while (i < chunk.byteLength) {
141 const byte1 = view.getUint32(i, true);
142 const byte2 = view.getUint32(i + 1, true);
143 const byte3 = view.getUint32(i + 2, true);
144 const byte4 = view.getUint32(i + 3, true);
145 const len = byte1 + (byte2 << 8) + (byte3 << 16) + (byte4 << 24);
146 i += 4;
147 ret.push(new TextDecoder().decode(chunk.slice(i, i + len)));
148 i += len;
149 }
150 return ret;
151 } else if (filter.id === "vlen-bytes") {
152 const view = new DataView(chunk);
153 const ret = [];
154 let i = 4;
155 while (i < chunk.byteLength) {
156 const byte1 = view.getUint32(i, true);
157 const byte2 = view.getUint32(i + 1, true);
158 const byte3 = view.getUint32(i + 2, true);
159 const byte4 = view.getUint32(i + 3, true);
160 const len = byte1 + (byte2 << 8) + (byte3 << 16) + (byte4 << 24);
161 i += 4;
162 ret.push(chunk.slice(i, i + len));
163 i += len;
164 }
165 return ret;
166 } else if (filter.id === "json2") {
167 const aa = JSON.parse(new TextDecoder().decode(chunk));
168 // aa has the form [item1, item2, ..., itemN, '|O', shape]
169 if (aa.length <= 2) {
170 console.warn("Unexpected json2", aa);
171 return new TextDecoder().decode(chunk);
172 }
173 if (!sameShape(aa[aa.length - 1], shape)) {
174 throw Error(
175 `Unexpected shape for json2 filter: ${aa[aa.length - 1]} !== ${shape}`,
176 );
177 }
178 if (!(aa[aa.length - 2] === "|O")) {
179 throw Error(
180 `Unexpected dtype for json2 filter: ${aa[aa.length - 2]} !== |O`,
181 );
182 }
183 if (!(aa.length - 2 === shape[0])) {
184 throw Error(
185 `Unexpected length for json2 filter: ${aa.length - 2} !== ${shape[0]}`,
186 );
187 }
188 if (shape.length > 1) {
189 return flattenArray(aa.slice(0, aa.length - 2), shape);
190 }
191 return aa.slice(0, aa.length - 2);
192 } else {
193 throw Error("Unhandled filter for |O " + filter.id);
194 }
195};
197const applyFilter = async (chunk: ArrayBuffer, filter: any) => {
198 if (filter.id === "zlib") {
199 const a = pako.inflate(chunk);
200 return a.buffer;
201 } else if (filter.id === "blosc") {
202 return new Blosc().decode(chunk);
203 } else if (filter.id === "shuffle") {
204 const elementSize = filter.elementsize;
205 const view = new DataView(chunk);
206 const ret = new Uint8Array(chunk.byteLength);
207 const a = chunk.byteLength / elementSize;
208 for (let i = 0; i < chunk.byteLength; i++) {
209 const b = i % elementSize;
210 const c = Math.floor(i / elementSize) * elementSize;
211 const j = b * a + c;
212 ret[j] = view.getUint8(i);
213 }
214 return ret.buffer;
215 }
216 console.warn("Filter not yet implemented", filter);
217 throw Error("Filter not yet implemented");
218};
220const sameShape = (a: number[], b: number[]): boolean => {
221 if (a.length !== b.length) return false;
222 for (let i = 0; i < a.length; i++) {
223 if (a[i] !== b[i]) return false;
224 }
225 return true;
226};
228const flattenArray = (aa: any[], shape: number[]): any[] => {
229 if (shape.length === 1) return aa;
230 const ret = [];
231 for (let i = 0; i < shape[0]; i++) {
232 const x = flattenArray(aa[i], shape.slice(1));
233 for (const xx of x) {
234 ret.push(xx);
235 }
236 }
237 return ret;
238};
240export default zarrDecodeChunkArray;