Fixes #884: replace boolean fancy indexing with dask-safe percentile path

quantile() and percentiles() used data[module.isfinite(data)] on dask
arrays, which creates unknown chunk sizes that degrade scheduling and
can force unexpected materialisations.  Replace with dedicated dask
functions that use da.where to clean inf→nan (preserving known chunks),
compute to numpy, then use np.nanpercentile + np.unique.

Author: brendancol

xrspatial/classify.py

@@ -405,17 +405,31 @@ def _run_quantile(data, k, module):
     return q
 
 
+def _run_dask_quantile(data, k):
+    # Avoid boolean fancy indexing (data[da.isfinite(data)]) which creates
+    # unknown dask chunk sizes.  Instead, replace inf with nan (preserves
+    # known chunks), compute to numpy, then use np.nanpercentile (#884).
+    w = 100.0 / k
+    p = np.arange(w, 100 + w, w)
+    if p[-1] > 100.0:
+        p[-1] = 100.0
+    clean = da.where(da.isinf(data), np.nan, data)
+    values = clean.ravel().compute()
+    q = np.nanpercentile(values, p)
+    q = np.unique(q)
+    return q
+
+
 def _run_dask_cupy_quantile(data, k):
-    # Convert dask+cupy chunks to numpy one at a time via map_blocks,
-    # then use dask's streaming approximate percentile (no full materialization).
+    # Convert dask+cupy chunks to numpy, then same safe path as dask (#884).
     data_cpu = data.map_blocks(cupy.asnumpy, dtype=data.dtype, meta=np.array(()))
-    return _run_quantile(data_cpu, k, da)
+    return _run_dask_quantile(data_cpu, k)
 
 
 def _quantile(agg, k):
     mapper = ArrayTypeFunctionMapping(
         numpy_func=lambda *args: _run_quantile(*args, module=np),
-        dask_func=lambda *args: _run_quantile(*args, module=da),
+        dask_func=_run_dask_quantile,
         cupy_func=lambda *args: _run_quantile(*args, module=cupy),
         dask_cupy_func=_run_dask_cupy_quantile
     )
@@ -1105,9 +1119,21 @@ def _run_percentiles(data, pct, module):
     return q
 
 
+def _run_dask_percentiles(data, pct):
+    # Avoid boolean fancy indexing (data[da.isfinite(data)]) which creates
+    # unknown dask chunk sizes.  Replace inf with nan, compute to numpy,
+    # then use np.nanpercentile (#884).
+    clean = da.where(da.isinf(data), np.nan, data)
+    values = clean.ravel().compute()
+    q = np.nanpercentile(values, pct)
+    q = np.unique(q)
+    return q
+
+
 def _run_dask_cupy_percentiles(data, pct):
+    # Convert dask+cupy chunks to numpy, then same safe path as dask (#884).
     data_cpu = data.map_blocks(cupy.asnumpy, dtype=data.dtype, meta=np.array(()))
-    return _run_percentiles(data_cpu, pct, da)
+    return _run_dask_percentiles(data_cpu, pct)
 
 
 @supports_dataset
@@ -1144,7 +1170,7 @@ def percentiles(agg: xr.DataArray,
 
     mapper = ArrayTypeFunctionMapping(
         numpy_func=lambda *args: _run_percentiles(*args, module=np),
-        dask_func=lambda *args: _run_percentiles(*args, module=da),
+        dask_func=_run_dask_percentiles,
         cupy_func=lambda *args: _run_percentiles(*args, module=cupy),
         dask_cupy_func=_run_dask_cupy_percentiles,
     )

xrspatial/tests/test_classify.py

@@ -945,3 +945,63 @@ def test_maximum_breaks_dask_num_sample():
     assert result1.shape == elevation.shape
     unique_vals = np.unique(result1.data.compute())
     assert len(unique_vals) <= 3 + 1  # at most k classes + possible nan
+
+
+# ===================================================================
+# Regression tests: dask paths must not use boolean fancy indexing
+# ===================================================================
+
+@dask_array_available
+def test_quantile_dask_no_unknown_chunks():
+    """quantile on dask must not create unknown chunk sizes (#884)."""
+    elevation = np.arange(100, dtype=np.float64).reshape(10, 10)
+    numpy_agg = xr.DataArray(elevation)
+    dask_agg = xr.DataArray(da.from_array(elevation, chunks=(5, 5)))
+
+    numpy_result = quantile(numpy_agg, k=5)
+    dask_result = quantile(dask_agg, k=5)
+
+    # Dask percentile is approximate, so just check same shape and k classes
+    assert dask_result.shape == numpy_result.shape
+    dask_vals = dask_result.data.compute()
+    unique_vals = np.unique(dask_vals[np.isfinite(dask_vals)])
+    assert len(unique_vals) == 5
+
+
+@dask_array_available
+def test_quantile_dask_with_nan_inf():
+    """quantile on dask handles NaN and inf without unknown chunks (#884)."""
+    elevation = np.array([
+        [-np.inf, 2., 3., 4., np.nan],
+        [5., 6., 7., 8., 9.],
+        [10., 11., 12., 13., 14.],
+        [15., 16., 17., 18., np.inf],
+    ])
+    dask_agg = xr.DataArray(da.from_array(elevation, chunks=(2, 5)))
+    result = quantile(dask_agg, k=5)
+    result_data = result.data.compute()
+    # NaN and inf inputs should produce NaN in the output
+    assert np.isnan(result_data[0, 0])   # was -inf
+    assert np.isnan(result_data[0, 4])   # was nan
+    assert np.isnan(result_data[3, 4])   # was inf
+    # Finite values should be classified
+    finite_result = result_data[np.isfinite(result_data)]
+    assert len(np.unique(finite_result)) == 5
+
+
+@dask_array_available
+def test_percentiles_dask_no_unknown_chunks():
+    """percentiles on dask must not create unknown chunk sizes (#884)."""
+    from xrspatial import percentiles as percentiles_fn
+    elevation = np.arange(100, dtype=np.float64).reshape(10, 10)
+    numpy_agg = xr.DataArray(elevation)
+    dask_agg = xr.DataArray(da.from_array(elevation, chunks=(5, 5)))
+
+    numpy_result = percentiles_fn(numpy_agg)
+    dask_result = percentiles_fn(dask_agg)
+
+    np.testing.assert_allclose(
+        numpy_result.data,
+        dask_result.data.compute(),
+        equal_nan=True,
+    )