combine fragments per chunk

seung-lab · Feb 18, 2025 · 605f3cb · 605f3cb
1 parent f13c8b4
commit 605f3cb
Showing 1 changed file with 30 additions and 14 deletions.
diff --git a/src/datasource/graphene/backend.ts b/src/datasource/graphene/backend.ts
@@ -202,6 +202,7 @@ interface ShardInfo {
 }
 
 interface GrapheneMultiscaleManifestChunk extends MultiscaleManifestChunk {
+  fragments: Array<FragmentId[]> | null;
   fragmentIds: FragmentId[] | null;
   shardInfo?: ShardInfo;
 }
@@ -217,7 +218,7 @@ function decodeMultiscaleManifestChunk(chunk: GrapheneMultiscaleManifestChunk, r
     clipLowerBound: vec3.clone(response.clipLowerBound),
     clipUpperBound: vec3.clone(response.clipUpperBound),
   };
-  chunk.fragmentIds = response.fragments;
+  chunk.fragments = response.fragments;
   chunk.manifest.clipLowerBound.fill(0);
   chunk.manifest.clipUpperBound.fill(10000000);
   chunk.manifest.octree[5 * (response.fragments.length - 1) + 4] &= 0x7fffffff;
@@ -260,6 +261,19 @@ async function decodeMultiscaleFragmentChunk(
   );
 }
 
+function mergeArrayBuffers(buffers: ArrayBuffer[]): ArrayBuffer {
+  const totalLength = buffers.reduce((sum, buf) => sum + buf.byteLength, 0);
+  const mergedBuffer = new ArrayBuffer(totalLength);
+  const mergedView = new Uint8Array(mergedBuffer);
+
+  let offset = 0;
+  for (const buffer of buffers) {
+      mergedView.set(new Uint8Array(buffer), offset);
+      offset += buffer.byteLength;
+  }
+  return mergedBuffer;
+}
+
 @registerSharedObject()
 export class GrapheneMultiscaleMeshSource extends WithParameters(
   WithSharedKvStoreContextCounterpart(MultiscaleMeshSource),
@@ -295,7 +309,7 @@ export class GrapheneMultiscaleMeshSource extends WithParameters(
       return decodeManifestChunk(chunk, { fragments: [] });
     }
     const { fetchOkImpl, baseUrl } = this.manifestHttpSource;
-    const manifestPath = `/manifest/multiscale/${chunk.objectId}?verify=1&prepend_seg_ids=1`;
+    const manifestPath = `/manifest/multiscale/${chunk.objectId}?verify=1`;
     const response = await (
       await fetchOkImpl(baseUrl + manifestPath, { signal })
     ).json();
@@ -322,21 +336,23 @@ export class GrapheneMultiscaleMeshSource extends WithParameters(
     const { parameters } = this;
     const manifestChunk = chunk.manifestChunk! as GrapheneMultiscaleManifestChunk;
     const chunkIndex = chunk.chunkIndex;
-    const { fragmentIds } = manifestChunk;
+    const { fragments } = manifestChunk;
 
     try {
-      let fragmentId = null;
-      if (fragmentIds !== null) {
-        fragmentId = fragmentIds[chunkIndex];
-        fragmentId = fragmentId.substring(fragmentId.indexOf(":") + 1);
+      if (fragments !== null) {
+        const fragmentsIds = fragments[chunkIndex];
+        let result: ArrayBuffer[] = new Array();
+        for (const fragmentId of fragmentsIds) {
+          let { response } = await downloadFragment(
+            this.fragmentKvStore,
+            fragmentId,
+            parameters,
+            signal,
+          );
+          result.push(await response.arrayBuffer())
+        }
+        await decodeMultiscaleFragmentChunk(chunk, mergeArrayBuffers(result));
       }
-      const { response } = await downloadFragment(
-        this.fragmentKvStore,
-        fragmentId,
-        parameters,
-        signal,
-      );
-      await decodeMultiscaleFragmentChunk(chunk, await response.arrayBuffer());
     } catch (e) {
       if (isNotFoundError(e)) {
         chunk.source!.removeChunk(chunk);