Fix dask benchmark setup for proximity and zonal

- Proximity: compute unique values from numpy data before slicing,
  since dask unique() returns unknown chunk sizes
- Zonal: skip custom stats benchmark for dask (API only supports
  default stats); fix else branch to not override None for dask

Author: brendancol

benchmarks/benchmarks/proximity.py

@@ -17,7 +17,10 @@ class Base:
     def setup(self, nx, n_target_values, distance_metric, type):
         ny = nx // 2
         self.agg = get_xr_dataarray((ny, nx), type, is_int=True)
-        unique_values = np.unique(self.agg.data)
+        data = self.agg.data
+        if type == "dask":
+            data = data.compute()
+        unique_values = np.unique(data)
         self.target_values = unique_values[:n_target_values]
 
 

benchmarks/benchmarks/zonal.py

@@ -64,22 +64,9 @@ def setup(self, raster_dim, zone_dim, backend):
         self.zones = create_arr(zones, backend=backend)
 
         # Now setup the custom stat funcs
-        if backend == 'dask':
-            from xrspatial.utils import ngjit
-
-            @ngjit
-            def l2normKernel(arr):
-                acc = 0
-                for x in arr:
-                    acc += x * x
-                return np.sqrt(acc)
-
-            self.custom_stats = {
-                'double_sum': lambda val: val.sum()*2,
-                'l2norm': lambda val: np.sqrt(np.sum(val * val)),
-                'l2normKernel': lambda val: l2normKernel(val)
-            }
-        elif backend == 'cupy':
+        # Dask backend only supports default stats, so custom_stats is None.
+        self.custom_stats = None
+        if backend == 'cupy':
             import cupy
             l2normKernel = cupy.ReductionKernel(
                 in_params='T x', out_params='float64 y',
@@ -92,7 +79,7 @@ def l2normKernel(arr):
                 'l2norm': lambda val: np.sqrt(cupy.sum(val * val)),
                 'l2normKernel': lambda val: l2normKernel(val)
             }
-        else:
+        elif backend == 'numpy':
             from xrspatial.utils import ngjit
 
             @ngjit
@@ -112,5 +99,7 @@ def time_zonal_stats_default(self, raster_dim, zone_dim, backend):
         zonal.stats(zones=self.zones, values=self.values)
 
     def time_zonal_stats_custom(self, raster_dim, zone_dim, backend):
+        if self.custom_stats is None:
+            raise NotImplementedError()
         zonal.stats(zones=self.zones, values=self.values,
                     stats_funcs=self.custom_stats)