Improve initialization performance for pyramids with many levels.

src/raster.js

@@ -1,10 +1,11 @@
-import React, { useRef, useEffect, useState } from 'react'
+import React, { useRef, useEffect, useState, useMemo } from 'react'
 import { useRegl } from './regl'
 import { useMap } from './map-provider'
 import { useControls } from './use-controls'
 import { createTiles } from './tiles'
 import { useRegion } from './region/context'
 import { useSetLoading } from './loading'
+import ZarrStore from './zarr-store'
 
 const Raster = (props) => {
   const {
@@ -32,6 +33,17 @@ const Raster = (props) => {
 
   camera.current = { center: center, zoom: zoom }
 
+  const store = useMemo(
+    () =>
+      new ZarrStore({
+        source: props.source,
+        version: props.version,
+        variable: props.variable,
+        coordinateKeys: Object.keys(selector),
+      }),
+    [props.source, props.version, props.variable]
+  )
+
   const queryRegion = async (r, s) => {
     const queryStart = new Date().getTime()
     lastQueried.current = queryStart
@@ -46,19 +58,26 @@ const Raster = (props) => {
     }
   }
 
+  useEffect(() => {
+    return () => {
+      store.cleanup()
+    }
+  }, [store])
+
   useEffect(() => {
     tiles.current = createTiles(regl, {
       ...props,
       setLoading,
       clearLoading,
+      store,
       invalidate: () => {
         map.triggerRepaint()
       },
       invalidateRegion: () => {
         setRegionDataInvalidated(new Date().getTime())
       },
     })
-  }, [])
+  }, [store])
 
   useEffect(() => {
     if (props.setLoading) {

src/tile.js

@@ -46,25 +46,27 @@ class Tile {
 
   async loadChunks(chunks) {
     const updated = await Promise.all(
-      chunks.map(
-        (chunk) =>
-          new Promise((resolve) => {
-            const key = chunk.join('.')
-            if (this.chunkedData[key]) {
-              resolve(false)
-            } 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)
-                })
-              })
-            }
-          })
-      )
+      chunks.map((chunk) => {
+        const key = chunk.join('.')
+        if (this.chunkedData[key]) {
+          return false
+        } else {
+          this._loading[key] = true
+          this._ready[key] = this._loader(chunk)
+            .then((data) => {
+              this.chunkedData[key] = data
+              this._loading[key] = false
+              return true
+            })
+            .catch((err) => {
+              console.error('Error loading chunk', key, err)
+              this._loading[key] = false
+              return false
+            })
+
+          return this._ready[key]
+        }
+      })
     )
 
     return updated.some(Boolean)

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'],