Add BoundaryStore: memmap-backed boundary strip storage for dask tile sweeps

Author: brendancol

xrspatial/_boundary_store.py

@@ -0,0 +1,104 @@
+"""Memmap-backed boundary strip storage for dask tile sweeps.
+
+Stores top/bottom/left/right boundary strips for a 2D tile grid in
+flat numpy memory-mapped files, avoiding O(N_tiles) in-memory nested
+lists that can OOM on very large inputs.
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+import tempfile
+
+import numpy as np
+
+
+class BoundaryStore:
+    """Disk-backed storage for boundary strips of a tiled 2D grid.
+
+    Each of the four sides (top, bottom, left, right) is stored as a
+    single contiguous memmap file.  Strip lookup is O(1) via
+    precomputed cumulative offsets, and ``get`` returns a zero-copy
+    memmap view.
+
+    Parameters
+    ----------
+    chunks_y : sequence of int
+        Tile heights (one per tile row).
+    chunks_x : sequence of int
+        Tile widths (one per tile column).
+    fill_value : float, optional
+        Initial fill value for all strips (default 0.0).
+    """
+
+    def __init__(self, chunks_y, chunks_x, fill_value=0.0):
+        self._tmpdir = tempfile.mkdtemp(prefix='xrs_bdry_')
+        self._closed = False
+        self._chunks_y = tuple(chunks_y)
+        self._chunks_x = tuple(chunks_x)
+        n_ty = len(chunks_y)
+        n_tx = len(chunks_x)
+        total_h = sum(chunks_y)
+        total_w = sum(chunks_x)
+
+        # Cumulative offsets for O(1) strip lookup
+        self._cum_x = np.zeros(n_tx + 1, dtype=np.int64)
+        np.cumsum(chunks_x, out=self._cum_x[1:])
+        self._cum_y = np.zeros(n_ty + 1, dtype=np.int64)
+        np.cumsum(chunks_y, out=self._cum_y[1:])
+
+        # top/bottom: strip length = chunks_x[ix], indexed by (iy, ix)
+        #   shape (n_ty, total_w) — row iy holds all top/bottom strips
+        # left/right: strip length = chunks_y[iy], indexed by (iy, ix)
+        #   shape (n_tx, total_h) — row ix holds all left/right strips
+        for name, shape in [('top', (n_ty, total_w)),
+                            ('bottom', (n_ty, total_w)),
+                            ('left', (n_tx, total_h)),
+                            ('right', (n_tx, total_h))]:
+            path = os.path.join(self._tmpdir, f'{name}.dat')
+            mm = np.memmap(path, dtype=np.float64, mode='w+', shape=shape)
+            mm[:] = fill_value
+            mm.flush()
+            setattr(self, f'_{name}', mm)
+
+    def get(self, side, iy, ix):
+        """Return a memmap view of the boundary strip for tile (iy, ix)."""
+        if side == 'top':
+            return self._top[iy, self._cum_x[ix]:self._cum_x[ix + 1]]
+        elif side == 'bottom':
+            return self._bottom[iy, self._cum_x[ix]:self._cum_x[ix + 1]]
+        elif side == 'left':
+            return self._left[ix, self._cum_y[iy]:self._cum_y[iy + 1]]
+        elif side == 'right':
+            return self._right[ix, self._cum_y[iy]:self._cum_y[iy + 1]]
+        else:
+            raise ValueError(f"Unknown side: {side!r}")
+
+    def set(self, side, iy, ix, data):
+        """Write *data* into the boundary strip for tile (iy, ix)."""
+        self.get(side, iy, ix)[:] = data
+
+    def close(self):
+        """Flush memmaps and remove temporary files."""
+        if self._closed:
+            return
+        self._closed = True
+        for name in ('top', 'bottom', 'left', 'right'):
+            mm = getattr(self, f'_{name}', None)
+            if mm is not None:
+                del mm
+                setattr(self, f'_{name}', None)
+        try:
+            shutil.rmtree(self._tmpdir)
+        except OSError:
+            pass
+
+    def __del__(self):
+        self.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        self.close()