From f0b583c512de0deb7c7ba7e6b4dbd717d7c88b41 Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Mon, 27 Apr 2026 10:11:02 -0700 Subject: [PATCH] Reject oversize focal kernels before allocation (#1284) focal.apply(), focal.focal_stats(), and focal.hotspots() accept any 2-D NumPy array as the kernel argument, validated only by custom_kernel() (which checks ndarray-ness and odd shape parity). The kernel-size memory guard added in #1241 only runs inside circle_kernel/annulus_kernel, so a caller building a kernel directly from NumPy bypasses it. With a small raster and a big kernel, the downstream paths allocate gigabytes before any work: - _apply_numpy allocates np.zeros_like(kernel) per call. - _apply_numpy_boundary and _convolve_2d_numpy_boundary pad the raster by kernel.shape // 2 on each side via _pad_array. - The dask paths set map_overlap depth to kernel.shape // 2 per chunk. Same pattern as the bilateral fix in #1236, where one float from the caller drove an unbounded radius. Add a _check_kernel_vs_raster_memory helper that budgets kernel bytes plus padded raster bytes (4 bytes per cell, matching the float32 internal dtype) and raises MemoryError when the combined total exceeds half of available memory. Wire it into the three public APIs after custom_kernel() validation, before dispatch. Adds four regression tests: apply, focal_stats, and hotspots all reject a (50001, 50001) kernel on a 10x10 raster, and a small 3x3 circle kernel on a 50x50 raster still works. --- xrspatial/focal.py | 53 +++++++++++++++++++++++++++++++++++ xrspatial/tests/test_focal.py | 37 ++++++++++++++++++++++++ 2 files changed, 90 insertions(+) diff --git a/xrspatial/focal.py b/xrspatial/focal.py index ba5740e8..8f1dd971 100644 --- a/xrspatial/focal.py +++ b/xrspatial/focal.py @@ -29,6 +29,7 @@ class cupy(object): from xrspatial.convolution import ( convolve_2d, custom_kernel, _convolve_2d_numpy, _convolve_2d_cupy, + _available_memory_bytes, ) from xrspatial.utils import (ArrayTypeFunctionMapping, _boundary_to_dask, _pad_array, _validate_boundary, _validate_raster, _validate_scalar, @@ -36,6 +37,49 @@ class cupy(object): from xrspatial.dataset_support import supports_dataset +def _check_kernel_vs_raster_memory(kernel, rows, cols, func_name): + """Reject kernel + raster combinations that would OOM the host. + + The focal public APIs (apply, focal_stats, hotspots) accept any 2-D + kernel passed through ``custom_kernel``, which only validates shape + parity. Several downstream allocations are driven by ``kernel.shape`` + on top of the raster shape: + + - ``_apply_numpy`` allocates ``np.zeros_like(kernel)`` per call. + - The non-NaN boundary path pads the raster by ``kernel.shape // 2`` + on each side via ``_pad_array`` -> ``np.pad``. + - The dask paths use ``map_overlap`` with depth ``kernel.shape // 2`` + per chunk. + + Without a guard, a small raster paired with an oversized kernel can + OOM the host (e.g. kernel of shape (50001, 50001) is ~10 GB on its + own; the padded raster is larger). Budget 4 bytes per kernel cell + (float32 internal dtype) plus the padded raster footprint, and raise + ``MemoryError`` when the total exceeds half of available memory. + """ + krows, kcols = kernel.shape + pad_h = krows // 2 + pad_w = kcols // 2 + + # 4 bytes per cell -- focal internals cast to float32. + kernel_bytes = krows * kcols * 4 + padded_rows = rows + 2 * pad_h + padded_cols = cols + 2 * pad_w + padded_bytes = padded_rows * padded_cols * 4 + + required = kernel_bytes + padded_bytes + available = _available_memory_bytes() + + if required > 0.5 * available: + raise MemoryError( + f"{func_name}(): kernel of shape {kernel.shape} on a " + f"{rows}x{cols} raster would need ~{required / 1e9:.1f} GB " + f"(kernel + padded raster), but only " + f"{available / 1e9:.1f} GB is available. " + f"Use a smaller kernel or a coarser cellsize." + ) + + def _apply_per_band(band_func, agg, *args, **kwargs): """Apply a 2D focal function independently to each band of a 3D array. @@ -565,6 +609,9 @@ def apply(raster, kernel, func=_calc_mean, name='focal_apply', boundary='nan'): _validate_boundary(boundary) + rows, cols = raster.shape[-2], raster.shape[-1] + _check_kernel_vs_raster_memory(kernel, rows, cols, func_name='apply') + # apply kernel to raster values # if raster is a numpy or dask with numpy backed data array, # the function func must be a @ngjit @@ -1073,6 +1120,9 @@ def focal_stats(agg, _validate_boundary(boundary) + rows, cols = agg.shape[-2], agg.shape[-1] + _check_kernel_vs_raster_memory(kernel, rows, cols, func_name='focal_stats') + mapper = ArrayTypeFunctionMapping( numpy_func=partial(_focal_stats_cpu, boundary=boundary), cupy_func=_focal_stats_cupy, @@ -1354,6 +1404,9 @@ def hotspots(raster, kernel, boundary='nan'): _validate_boundary(boundary) + rows, cols = raster.shape[-2], raster.shape[-1] + _check_kernel_vs_raster_memory(kernel, rows, cols, func_name='hotspots') + mapper = ArrayTypeFunctionMapping( numpy_func=partial(_hotspots_numpy, boundary=boundary), cupy_func=partial(_hotspots_cupy, boundary=boundary), diff --git a/xrspatial/tests/test_focal.py b/xrspatial/tests/test_focal.py index 40ddaa65..1f93c21d 100644 --- a/xrspatial/tests/test_focal.py +++ b/xrspatial/tests/test_focal.py @@ -224,6 +224,43 @@ def test_circle_kernel_small_radius_not_rejected_1241(): assert kernel.shape == (201, 201) +def test_apply_rejects_oversize_kernel_1284(): + # Regression for #1284: focal.apply must reject a user-supplied + # kernel that would OOM on the padded raster + kernel allocation. + # custom_kernel only checks shape parity, so a tiny raster paired + # with a giant kernel used to allocate many GB before any work. + raster = xr.DataArray(np.zeros((10, 10), dtype=np.float32)) + big_kernel = np.ones((50001, 50001), dtype=np.float32) + with pytest.raises(MemoryError, match=r"apply\(\): kernel of shape"): + apply(raster, big_kernel) + + +def test_focal_stats_rejects_oversize_kernel_1284(): + # Regression for #1284: focal_stats must apply the same kernel + # vs raster guard before dispatching to any backend. + raster = xr.DataArray(np.zeros((10, 10), dtype=np.float32)) + big_kernel = np.ones((50001, 50001), dtype=np.float32) + with pytest.raises(MemoryError, match=r"focal_stats\(\): kernel of shape"): + focal_stats(raster, big_kernel, stats_funcs=['mean']) + + +def test_hotspots_rejects_oversize_kernel_1284(): + # Regression for #1284: hotspots calls convolve_2d under the hood, + # which inherits the same padded-allocation footprint. + raster = xr.DataArray(np.zeros((10, 10), dtype=np.float32)) + big_kernel = np.ones((50001, 50001), dtype=np.float32) + with pytest.raises(MemoryError, match=r"hotspots\(\): kernel of shape"): + hotspots(raster, big_kernel) + + +def test_apply_small_kernel_not_rejected_1284(): + # The guard must not fire for realistic kernel + raster combos. + raster = xr.DataArray(np.ones((50, 50), dtype=np.float32)) + kernel = circle_kernel(1, 1, 3) + out = apply(raster, kernel) + assert out.shape == (50, 50) + + def test_convolution_numpy( convolve_2d_data, convolution_custom_kernel,