Improve initialization performance for pyramids with many levels.

src/initialize-store.js

@@ -2,6 +2,14 @@ import zarr from 'zarr-js'
 
 import { getPyramidMetadata } from './utils'
 
+// return promises for all for consistency
+const wrapGet = (getFn) => {
+  return (chunkIndices) =>
+    new Promise((resolve, reject) => {
+      getFn(chunkIndices, (err, out) => (err ? reject(err) : resolve(out)))
+    })
+}
+
 const initializeStore = async (source, version, variable, coordinateKeys) => {
   let metadata
   let loaders
@@ -14,36 +22,106 @@ const initializeStore = async (source, version, variable, coordinateKeys) => {
   const coordinates = {}
   switch (version) {
     case 'v2':
-      await new Promise((resolve) =>
-        zarr(window.fetch, version).openGroup(source, (err, l, m) => {
-          loaders = l
-          metadata = m
-          resolve()
+      try {
+        // Fetch consolidated metadata directly
+        const zmetadata = await fetch(`${source}/.zmetadata`).then((res) =>
+          res.json()
+        )
+        metadata = { metadata: zmetadata.metadata }
+        const rootAttrs = zmetadata.metadata['.zattrs']
+        ;({ levels, maxZoom, tileSize, crs } = getPyramidMetadata(
+          rootAttrs.multiscales
+        ))
+
+        const zattrs = metadata.metadata[`${levels[0]}/${variable}/.zattrs`]
+        const zarray = metadata.metadata[`${levels[0]}/${variable}/.zarray`]
+        dimensions = zattrs['_ARRAY_DIMENSIONS']
+        shape = zarray.shape
+        chunks = zarray.chunks
+        fill_value = zarray.fill_value
+        dtype = zarray.dtype
+
+        const getCache = new Map()
+
+        const callGet = async (key, chunkIndices) => {
+          let getPromise = getCache.get(key)
+
+          if (!getPromise) {
+            const arrayMeta = metadata.metadata[`${key}/.zarray`]
+            getPromise = new Promise((resolve, reject) => {
+              zarr(window.fetch, version).open(
+                `${source}/${key}`,
+                (err, get) => (err ? reject(err) : resolve(get)),
+                arrayMeta
+              )
+            })
+            getCache.set(key, getPromise)
+          }
+
+          const get = await getPromise
+          return new Promise((resolve, reject) => {
+            get(chunkIndices, (err, out) => (err ? reject(err) : resolve(out)))
+          })
+        }
+
+        await Promise.all(
+          coordinateKeys.map(async (key) => {
+            const coordKey = `${levels[0]}/${key}`
+            const chunk = await callGet(coordKey, [0])
+            coordinates[key] = Array.from(chunk.data)
+          })
+        )
+
+        loaders = {}
+        levels.forEach((level) => {
+          const key = `${level}/${variable}`
+          loaders[key] = (chunkIndices) => callGet(key, chunkIndices)
         })
-      )
-      ;({ levels, maxZoom, tileSize, crs } = getPyramidMetadata(
-        metadata.metadata['.zattrs'].multiscales
-      ))
 
-      const zattrs = metadata.metadata[`${levels[0]}/${variable}/.zattrs`]
-      const zarray = metadata.metadata[`${levels[0]}/${variable}/.zarray`]
-      dimensions = zattrs['_ARRAY_DIMENSIONS']
-      shape = zarray.shape
-      chunks = zarray.chunks
-      fill_value = zarray.fill_value
-      dtype = zarray.dtype
+        coordinateKeys.forEach((key) => {
+          const coordKey = `${levels[0]}/${key}`
+          loaders[coordKey] = (chunkIndices) => callGet(coordKey, chunkIndices)
+        })
+      } catch (e) {
+        // Fallback to openGroup
+        let rawLoaders
+        await new Promise((resolve) =>
+          zarr(window.fetch, version).openGroup(source, (err, l, m) => {
+            rawLoaders = l
+            metadata = m
+            resolve()
+          })
+        )
+        ;({ levels, maxZoom, tileSize, crs } = getPyramidMetadata(
+          metadata.metadata['.zattrs'].multiscales
+        ))
 
-      await Promise.all(
-        coordinateKeys.map(
-          (key) =>
-            new Promise((resolve) => {
-              loaders[`${levels[0]}/${key}`]([0], (err, chunk) => {
+        const zattrs = metadata.metadata[`${levels[0]}/${variable}/.zattrs`]
+        const zarray = metadata.metadata[`${levels[0]}/${variable}/.zarray`]
+        dimensions = zattrs['_ARRAY_DIMENSIONS']
+        shape = zarray.shape
+        chunks = zarray.chunks
+        fill_value = zarray.fill_value
+        dtype = zarray.dtype
+
+        await Promise.all(
+          coordinateKeys.map((key) => {
+            const coordKey = `${levels[0]}/${key}`
+            return new Promise((resolve, reject) => {
+              rawLoaders[coordKey]([0], (err, chunk) => {
+                if (err) return reject(err)
                 coordinates[key] = Array.from(chunk.data)
                 resolve()
               })
             })
+          })
         )
-      )
+
+        loaders = {}
+        Object.keys(rawLoaders).forEach((key) => {
+          loaders[key] = wrapGet(rawLoaders[key])
+        })
+      }
 
       break
     case 'v3':
@@ -65,40 +143,51 @@ const initializeStore = async (source, version, variable, coordinateKeys) => {
       fill_value = arrayMetadata.fill_value
       // dtype = arrayMetadata.data_type
 
+      const getCache = new Map()
+
+      const callGet = async (key, chunkIndices, meta = null) => {
+        let getPromise = getCache.get(key)
+
+        if (!getPromise) {
+          getPromise = new Promise((resolve, reject) => {
+            zarr(window.fetch, version).open(
+              `${source}/${key}`,
+              (err, get) => (err ? reject(err) : resolve(get)),
+              meta
+            )
+          })
+          getCache.set(key, getPromise)
+        }
+
+        const get = await getPromise
+        return new Promise((resolve, reject) => {
+          get(chunkIndices, (err, out) => (err ? reject(err) : resolve(out)))
+        })
+      }
+
+      await Promise.all(
+        coordinateKeys.map(async (key) => {
+          const coordKey = `${levels[0]}/${key}`
+          const chunk = await callGet(coordKey, [0])
+          coordinates[key] = Array.from(chunk.data)
+        })
+      )
+
       loaders = {}
-      await Promise.all([
-        ...levels.map(
-          (level) =>
-            new Promise((resolve) => {
-              zarr(window.fetch, version).open(
-                `${source}/${level}/${variable}`,
-                (err, get) => {
-                  loaders[`${level}/${variable}`] = get
-                  resolve()
-                },
-                level === 0 ? arrayMetadata : null
-              )
-            })
-        ),
-        ...coordinateKeys.map(
-          (key) =>
-            new Promise((resolve) => {
-              zarr(window.fetch, version).open(
-                `${source}/${levels[0]}/${key}`,
-                (err, get) => {
-                  get([0], (err, chunk) => {
-                    coordinates[key] = Array.from(chunk.data)
-                    resolve()
-                  })
-                }
-              )
-            })
-        ),
-      ])
+      levels.forEach((level) => {
+        const key = `${level}/${variable}`
+        const meta = level === 0 ? arrayMetadata : null
+        loaders[key] = (chunkIndices) => callGet(key, chunkIndices, meta)
+      })
+
+      coordinateKeys.forEach((key) => {
+        const coordKey = `${levels[0]}/${key}`
+        loaders[coordKey] = (chunkIndices) => callGet(coordKey, chunkIndices)
+      })
       break
     default:
       throw new Error(
-        `Unexpected Zarr version: ${version}. Must be one of 'v1', 'v2'.`
+        `Unexpected Zarr version: ${version}. Must be one of 'v2', 'v3'.`
       )
   }
 

src/tile.js

@@ -55,12 +55,19 @@ class Tile {
             } else {
               this._loading[key] = true
               this._ready[key] = new Promise((innerResolve) => {
-                this._loader(chunk, (err, data) => {
-                  this.chunkedData[key] = data
-                  this._loading[key] = false
-                  innerResolve(true)
-                  resolve(true)
-                })
+                this._loader(chunk)
+                  .then((data) => {
+                    this.chunkedData[key] = data
+                    this._loading[key] = false
+                    innerResolve(true)
+                    resolve(true)
+                  })
+                  .catch((err) => {
+                    console.error('Error loading chunk', key, err)
+                    this._loading[key] = false
+                    innerResolve(false)
+                    resolve(false)
+                  })
               })
             }
           })

src/tile.test.js

@@ -28,12 +28,9 @@ describe('Tile', () => {
     buffer = jest.fn()
     defaults = {
       key: '0,0,0',
-      loader: jest.fn().mockImplementation((chunk, cb) =>
-        cb(
-          null, // error
-          createMockChunk(chunk)
-        )
-      ),
+      loader: jest
+        .fn()
+        .mockImplementation((chunk) => Promise.resolve(createMockChunk(chunk))),
       shape: [10, 1, 1],
       chunks: [5, 1, 1],
       dimensions: ['year', 'y', 'x'],
@@ -148,14 +145,8 @@ describe('Tile', () => {
         ])
 
         expect(defaults.loader).toHaveBeenCalledTimes(2)
-        expect(defaults.loader).toHaveBeenCalledWith(
-          [0, 0, 0],
-          expect.anything()
-        )
-        expect(defaults.loader).toHaveBeenCalledWith(
-          [1, 0, 0],
-          expect.anything()
-        )
+        expect(defaults.loader).toHaveBeenCalledWith([0, 0, 0])
+        expect(defaults.loader).toHaveBeenCalledWith([1, 0, 0])
       })
 
       it('does not repeat loading for any chunks have been loaded', async () => {
@@ -267,15 +258,11 @@ describe('Tile', () => {
 
       beforeEach(() => {
         resolvers = []
-        const loader = jest.fn().mockImplementation((chunk, cb) =>
-          new Promise((resolve) => {
-            resolvers.push(resolve)
-          }).then(() => {
-            cb(
-              null, // error
-              createMockChunk(chunk)
-            )
-          })
+        const loader = jest.fn().mockImplementation(
+          (chunk) =>
+            new Promise((resolve) => {
+              resolvers.push(() => resolve(createMockChunk(chunk)))
+            })
         )
         tile = new Tile({ ...defaults, loader })
       })
@@ -441,12 +428,11 @@ describe('Tile', () => {
         const selector = {}
         const tile = new Tile({
           ...defaults,
-          loader: jest.fn().mockImplementation((chunk, cb) =>
-            cb(
-              null, // error
-              ndarray([1, 2, 3, 4], [4, 1, 1])
-            )
-          ),
+          loader: jest
+            .fn()
+            .mockImplementation((chunk) =>
+              Promise.resolve(ndarray([1, 2, 3, 4], [4, 1, 1]))
+            ),
           shape: [4, 1, 1],
           chunks: [4, 1, 1],
           dimensions: ['band', 'y', 'x'],
@@ -491,12 +477,11 @@ describe('Tile', () => {
         const selector = {}
         const tile = new Tile({
           ...defaults,
-          loader: jest.fn().mockImplementation((chunk, cb) =>
-            cb(
-              null, // error
-              ndarray([1, 2, 3, 4], [2, 2])
-            )
-          ),
+          loader: jest
+            .fn()
+            .mockImplementation((chunk) =>
+              Promise.resolve(ndarray([1, 2, 3, 4], [2, 2]))
+            ),
           shape: [2, 2],
           chunks: [2, 2],
           dimensions: ['y', 'x'],