Fixes #877, #876: prevent OOM from full dask materialisation in classify.py

natural_breaks and maximum_breaks dask code paths called .ravel().compute()
on the full array, materialising the entire dataset into RAM. Replace with
capped sampling via _generate_sample_indices() + indexed access so only the
sample is ever computed. Add num_sample parameter to maximum_breaks (default
20_000, matching natural_breaks).

Author: brendancol

xrspatial/classify.py

@@ -695,16 +695,11 @@ def _run_dask_natural_break(agg, num_sample, k):
     max_data = float(da.nanmax(da.where(da.isinf(data), np.nan, data)).compute())
 
     num_data = data.size
-    if num_sample is not None and num_sample < num_data:
-        # Sample lazily from dask array; only materialize the sample
-        sample_idx = _generate_sample_indices(num_data, num_sample)
-        sample_data_np = np.asarray(data.ravel()[sample_idx].compute())
-        bins, uvk = _compute_natural_break_bins(
-            sample_data_np, None, k, max_data)
-    else:
-        data_flat_np = np.asarray(data.ravel().compute())
-        bins, uvk = _compute_natural_break_bins(
-            data_flat_np, None, k, max_data)
+    if num_sample is None or num_sample >= num_data:
+        num_sample = num_data  # cap: still uses indexed access, never .compute() all
+    sample_idx = _generate_sample_indices(num_data, num_sample)
+    sample_data_np = np.asarray(data.ravel()[sample_idx].compute())
+    bins, uvk = _compute_natural_break_bins(sample_data_np, None, k, max_data)
 
     out = _bin(agg, bins, np.arange(uvk))
     return out
@@ -717,17 +712,12 @@ def _run_dask_cupy_natural_break(agg, num_sample, k):
     max_data = float(da.nanmax(da.where(da.isinf(data), np.nan, data)).compute().item())
 
     num_data = data.size
-    if num_sample is not None and num_sample < num_data:
-        # Sample lazily from dask array; only materialize the sample
-        sample_idx = _generate_sample_indices(num_data, num_sample)
-        sample_data = data.ravel()[sample_idx].compute()
-        sample_data_np = cupy.asnumpy(sample_data)
-        bins, uvk = _compute_natural_break_bins(
-            sample_data_np, None, k, max_data)
-    else:
-        data_flat_np = cupy.asnumpy(data.ravel().compute())
-        bins, uvk = _compute_natural_break_bins(
-            data_flat_np, None, k, max_data)
+    if num_sample is None or num_sample >= num_data:
+        num_sample = num_data  # cap: still uses indexed access, never .compute() all
+    sample_idx = _generate_sample_indices(num_data, num_sample)
+    sample_data = data.ravel()[sample_idx].compute()
+    sample_data_np = cupy.asnumpy(sample_data)
+    bins, uvk = _compute_natural_break_bins(sample_data_np, None, k, max_data)
 
     out = _bin(agg, bins, np.arange(uvk))
     return out
@@ -1204,7 +1194,8 @@ def _compute_maximum_break_bins(values_np, k):
     return bins
 
 
-def _run_maximum_breaks(agg, k, module):
+def _run_maximum_breaks(agg, num_sample, k, module):
+    # num_sample ignored for in-memory backends
     data = agg.data
     if module == cupy:
         finite = data[cupy.isfinite(data)]
@@ -1218,21 +1209,29 @@ def _run_maximum_breaks(agg, k, module):
     return out
 
 
-def _run_dask_maximum_breaks(agg, k):
+def _run_dask_maximum_breaks(agg, num_sample, k):
     data = agg.data
     data_clean = da.where(da.isinf(data), np.nan, data)
-    values_np = np.asarray(data_clean.ravel().compute())
+    num_data = data.size
+    if num_sample is None or num_sample >= num_data:
+        num_sample = num_data
+    sample_idx = _generate_sample_indices(num_data, num_sample)
+    values_np = np.asarray(data_clean.ravel()[sample_idx].compute())
     values_np = values_np[np.isfinite(values_np)]
     bins = _compute_maximum_break_bins(values_np, k)
     out = _bin(agg, bins, np.arange(len(bins)))
     return out
 
 
-def _run_dask_cupy_maximum_breaks(agg, k):
+def _run_dask_cupy_maximum_breaks(agg, num_sample, k):
     data = agg.data
     data_clean = da.where(da.isinf(data), np.nan, data)
     data_cpu = data_clean.map_blocks(cupy.asnumpy, dtype=data.dtype, meta=np.array(()))
-    values_np = np.asarray(data_cpu.ravel().compute())
+    num_data = data.size
+    if num_sample is None or num_sample >= num_data:
+        num_sample = num_data
+    sample_idx = _generate_sample_indices(num_data, num_sample)
+    values_np = np.asarray(data_cpu.ravel()[sample_idx].compute())
     values_np = values_np[np.isfinite(values_np)]
     bins = _compute_maximum_break_bins(values_np, k)
     out = _bin(agg, bins, np.arange(len(bins)))
@@ -1242,6 +1241,7 @@ def _run_dask_cupy_maximum_breaks(agg, k):
 @supports_dataset
 def maximum_breaks(agg: xr.DataArray,
                    k: int = 5,
+                   num_sample: Optional[int] = 20_000,
                    name: Optional[str] = 'maximum_breaks') -> xr.DataArray:
     """
     Classify data using the Maximum Breaks algorithm.
@@ -1256,6 +1256,11 @@ def maximum_breaks(agg: xr.DataArray,
         of values to be classified.
     k : int, default=5
         Number of classes to be produced.
+    num_sample : int or None, default=20000
+        Number of sample data points used to fit the model.
+        For dask-backed arrays the sample is drawn lazily to avoid
+        materialising the entire array into RAM.  ``None`` means use
+        all data (safe for numpy/cupy, automatically capped for dask).
     name : str, default='maximum_breaks'
         Name of output aggregate array.
 
@@ -1277,7 +1282,7 @@ def maximum_breaks(agg: xr.DataArray,
         cupy_func=lambda *args: _run_maximum_breaks(*args, module=cupy),
         dask_cupy_func=_run_dask_cupy_maximum_breaks,
     )
-    out = mapper(agg)(agg, k)
+    out = mapper(agg)(agg, num_sample, k)
     return xr.DataArray(out,
                         name=name,
                         dims=agg.dims,

xrspatial/tests/test_classify.py

@@ -883,3 +883,65 @@ def test_maximum_breaks_dask_matches_numpy():
         maximum_breaks(dask_agg).data.compute(),
         equal_nan=True,
     )
+
+
+# ===================================================================
+# Regression tests: dask paths must not materialise the full array
+# ===================================================================
+
+@dask_array_available
+def test_natural_breaks_dask_no_full_compute():
+    """natural_breaks with num_sample=None on dask must not call
+    .ravel().compute() on the full array (#877)."""
+    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 = natural_breaks(numpy_agg, num_sample=None, k=5)
+    dask_result = natural_breaks(dask_agg, num_sample=None, k=5)
+
+    np.testing.assert_allclose(
+        numpy_result.data,
+        dask_result.data.compute(),
+        equal_nan=True,
+    )
+
+
+@dask_array_available
+def test_maximum_breaks_dask_no_full_compute():
+    """maximum_breaks on dask must use sampling, not .ravel().compute() (#876)."""
+    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)))
+
+    # Default num_sample (20_000) > data.size (100), so all data used
+    numpy_result = maximum_breaks(numpy_agg, k=5)
+    dask_result = maximum_breaks(dask_agg, k=5)
+
+    np.testing.assert_allclose(
+        numpy_result.data,
+        dask_result.data.compute(),
+        equal_nan=True,
+    )
+
+
+@dask_array_available
+def test_maximum_breaks_dask_num_sample():
+    """maximum_breaks with explicit num_sample on dask produces valid,
+    deterministic results (#876)."""
+    elevation = np.arange(100, dtype=np.float64).reshape(10, 10)
+    dask_agg = xr.DataArray(da.from_array(elevation, chunks=(5, 5)))
+
+    result1 = maximum_breaks(dask_agg, k=3, num_sample=50)
+    result2 = maximum_breaks(dask_agg, k=3, num_sample=50)
+
+    # Deterministic: same input + same seed → same output
+    np.testing.assert_allclose(
+        result1.data.compute(),
+        result2.data.compute(),
+        equal_nan=True,
+    )
+    # Valid classification: correct shape and values in expected range
+    assert result1.shape == elevation.shape
+    unique_vals = np.unique(result1.data.compute())
+    assert len(unique_vals) <= 3 + 1  # at most k classes + possible nan