Add multi-band write, integer nodata, PackBits, dask reads, BigTIFF write

Six features filling the main gaps for real-world use:

1. Multi-band write: 3D arrays (height, width, bands) now write as
   multi-band GeoTIFFs with correct BitsPerSample, SampleFormat, and
   PhotometricInterpretation (RGB for 3+ bands). Overviews work for
   multi-band too. read_geotiff returns all bands by default (band=None)
   with a 'band' dimension.

2. Integer nodata masking: uint8/uint16/int16 arrays with nodata values
   are promoted to float64 and masked with NaN on read, matching
   rioxarray behavior. Previously only float arrays were masked.

3. PackBits compression (tag 32773): simple RLE codec, both read and
   write. Common in older TIFF files.

4. JPEG decompression (tag 7): read support via Pillow for
   JPEG-compressed tiles/strips. Import is optional and lazy.

5. BigTIFF write: auto-detects when output exceeds ~4GB and switches
   to BigTIFF format (16-byte header, 20-byte IFD entries, 8-byte
   offsets). Prevents silent offset overflow corruption on large files.

6. Dask lazy reads: read_geotiff_dask() returns a dask-backed
   DataArray using windowed reads per chunk. Works for single-band
   and multi-band files with nodata masking per chunk.

178 tests passing.

Author: brendancol

xrspatial/geotiff/__init__.py

@@ -20,7 +20,7 @@
 from ._reader import read_to_array
 from ._writer import write
 
-__all__ = ['read_geotiff', 'write_geotiff', 'open_cog']
+__all__ = ['read_geotiff', 'write_geotiff', 'open_cog', 'read_geotiff_dask']
 
 
 def _geo_to_coords(geo_info, height: int, width: int) -> dict:
@@ -86,7 +86,7 @@ def _coords_to_transform(da: xr.DataArray) -> GeoTransform | None:
 
 def read_geotiff(source: str, *, window=None,
                  overview_level: int | None = None,
-                 band: int = 0,
+                 band: int | None = None,
                  name: str | None = None) -> xr.DataArray:
     """Read a GeoTIFF file into an xarray.DataArray.
 
@@ -139,13 +139,27 @@ def read_geotiff(source: str, *, window=None,
     nodata = geo_info.nodata
     if nodata is not None:
         attrs['nodata'] = nodata
-        if arr.dtype.kind == 'f' and not np.isnan(nodata):
-            arr = arr.copy()
-            arr[arr == np.float32(nodata)] = np.nan
+        if arr.dtype.kind == 'f':
+            if not np.isnan(nodata):
+                arr = arr.copy()
+                arr[arr == arr.dtype.type(nodata)] = np.nan
+        elif arr.dtype.kind in ('u', 'i'):
+            # Integer arrays: convert to float to represent NaN
+            nodata_int = int(nodata)
+            mask = arr == arr.dtype.type(nodata_int)
+            if mask.any():
+                arr = arr.astype(np.float64)
+                arr[mask] = np.nan
+
+    if arr.ndim == 3:
+        dims = ['y', 'x', 'band']
+        coords['band'] = np.arange(arr.shape[2])
+    else:
+        dims = ['y', 'x']
 
     da = xr.DataArray(
         arr,
-        dims=['y', 'x'],
+        dims=dims,
         coords=coords,
         name=name,
         attrs=attrs,
@@ -204,8 +218,8 @@ def write_geotiff(data: xr.DataArray | np.ndarray, path: str, *,
     else:
         arr = np.asarray(data)
 
-    if arr.ndim != 2:
-        raise ValueError(f"Expected 2D array, got {arr.ndim}D")
+    if arr.ndim not in (2, 3):
+        raise ValueError(f"Expected 2D or 3D array, got {arr.ndim}D")
 
     write(
         arr, path,
@@ -240,3 +254,115 @@ def open_cog(url: str, *,
     xr.DataArray
     """
     return read_geotiff(url, overview_level=overview_level)
+
+
+def read_geotiff_dask(source: str, *, chunks: int | tuple = 512,
+                      overview_level: int | None = None,
+                      name: str | None = None) -> xr.DataArray:
+    """Read a GeoTIFF as a dask-backed DataArray for out-of-core processing.
+
+    Each chunk is loaded lazily via windowed reads.
+
+    Parameters
+    ----------
+    source : str
+        File path.
+    chunks : int or (row_chunk, col_chunk) tuple
+        Chunk size in pixels. Default 512.
+    overview_level : int or None
+        Overview level (0 = full resolution).
+    name : str or None
+        Name for the DataArray.
+
+    Returns
+    -------
+    xr.DataArray
+        Dask-backed DataArray with y/x coordinates.
+    """
+    import dask.array as da
+
+    # First, do a metadata-only read to get shape, dtype, coords, attrs
+    arr, geo_info = read_to_array(source, overview_level=overview_level)
+    full_h, full_w = arr.shape[:2]
+    n_bands = arr.shape[2] if arr.ndim == 3 else 0
+    dtype = arr.dtype
+
+    coords = _geo_to_coords(geo_info, full_h, full_w)
+
+    if name is None:
+        import os
+        name = os.path.splitext(os.path.basename(source))[0]
+
+    attrs = {}
+    if geo_info.crs_epsg is not None:
+        attrs['crs'] = geo_info.crs_epsg
+    if geo_info.raster_type == RASTER_PIXEL_IS_POINT:
+        attrs['raster_type'] = 'point'
+    if geo_info.nodata is not None:
+        attrs['nodata'] = geo_info.nodata
+
+    if isinstance(chunks, int):
+        ch_h = ch_w = chunks
+    else:
+        ch_h, ch_w = chunks
+
+    # Build dask array from delayed windowed reads
+    rows = list(range(0, full_h, ch_h))
+    cols = list(range(0, full_w, ch_w))
+
+    # For multi-band, each window read returns (h, w, bands); for single-band (h, w)
+    # read_to_array with band=0 extracts a single band, band=None returns all
+    band_arg = None  # return all bands (or 2D if single-band)
+
+    dask_rows = []
+    for r0 in rows:
+        r1 = min(r0 + ch_h, full_h)
+        dask_cols = []
+        for c0 in cols:
+            c1 = min(c0 + ch_w, full_w)
+            if n_bands > 0:
+                block_shape = (r1 - r0, c1 - c0, n_bands)
+            else:
+                block_shape = (r1 - r0, c1 - c0)
+            block = da.from_delayed(
+                _delayed_read_window(source, r0, c0, r1, c1,
+                                     overview_level, geo_info.nodata,
+                                     dtype, band_arg),
+                shape=block_shape,
+                dtype=dtype,
+            )
+            dask_cols.append(block)
+        dask_rows.append(da.concatenate(dask_cols, axis=1))
+
+    dask_arr = da.concatenate(dask_rows, axis=0)
+
+    if n_bands > 0:
+        dims = ['y', 'x', 'band']
+        coords['band'] = np.arange(n_bands)
+    else:
+        dims = ['y', 'x']
+
+    return xr.DataArray(
+        dask_arr, dims=dims, coords=coords, name=name, attrs=attrs,
+    )
+
+
+def _delayed_read_window(source, r0, c0, r1, c1, overview_level, nodata,
+                         dtype, band):
+    """Dask-delayed function to read a single window."""
+    import dask
+    @dask.delayed
+    def _read():
+        arr, _ = read_to_array(source, window=(r0, c0, r1, c1),
+                               overview_level=overview_level, band=band)
+        if nodata is not None:
+            if arr.dtype.kind == 'f' and not np.isnan(nodata):
+                arr = arr.copy()
+                arr[arr == arr.dtype.type(nodata)] = np.nan
+            elif arr.dtype.kind in ('u', 'i'):
+                mask = arr == arr.dtype.type(int(nodata))
+                if mask.any():
+                    arr = arr.astype(np.float64)
+                    arr[mask] = np.nan
+        return arr
+    return _read()

xrspatial/geotiff/_compression.py

@@ -522,16 +522,126 @@ def fp_predictor_encode(data: np.ndarray, width: int, height: int,
     return buf
 
 
+# -- PackBits (simple RLE) ----------------------------------------------------
+
+def packbits_decompress(data: bytes) -> bytes:
+    """Decompress PackBits (TIFF compression tag 32773).
+
+    Simple RLE: read a header byte n.
+    - 0 <= n <= 127: copy the next n+1 bytes literally.
+    - -127 <= n <= -1: repeat the next byte 1-n times.
+    - n == -128: no-op.
+    """
+    src = data if isinstance(data, (bytes, bytearray)) else bytes(data)
+    out = bytearray()
+    i = 0
+    length = len(src)
+    while i < length:
+        n = src[i]
+        if n > 127:
+            n = n - 256  # interpret as signed
+        i += 1
+        if 0 <= n <= 127:
+            count = n + 1
+            out.extend(src[i:i + count])
+            i += count
+        elif -127 <= n <= -1:
+            if i < length:
+                out.extend(bytes([src[i]]) * (1 - n))
+                i += 1
+        # n == -128: skip
+    return bytes(out)
+
+
+def packbits_compress(data: bytes) -> bytes:
+    """Compress data using PackBits."""
+    src = data if isinstance(data, (bytes, bytearray)) else bytes(data)
+    out = bytearray()
+    i = 0
+    length = len(src)
+    while i < length:
+        # Check for a run of identical bytes
+        j = i + 1
+        while j < length and j - i < 128 and src[j] == src[i]:
+            j += 1
+        run_len = j - i
+
+        if run_len >= 3:
+            # Encode as run
+            out.append((256 - (run_len - 1)) & 0xFF)
+            out.append(src[i])
+            i = j
+        else:
+            # Literal run: accumulate non-repeating bytes
+            lit_start = i
+            i = j
+            while i < length and i - lit_start < 128:
+                # Check if a run starts here
+                if i + 2 < length and src[i] == src[i + 1] == src[i + 2]:
+                    break
+                i += 1
+            lit_len = i - lit_start
+            out.append(lit_len - 1)
+            out.extend(src[lit_start:lit_start + lit_len])
+    return bytes(out)
+
+
+# -- JPEG codec (via Pillow) --------------------------------------------------
+
+JPEG_AVAILABLE = False
+try:
+    from PIL import Image
+    JPEG_AVAILABLE = True
+except ImportError:
+    pass
+
+
+def jpeg_decompress(data: bytes, width: int = 0, height: int = 0,
+                    samples: int = 1) -> bytes:
+    """Decompress JPEG tile/strip data. Requires Pillow."""
+    if not JPEG_AVAILABLE:
+        raise ImportError(
+            "Pillow is required to read JPEG-compressed TIFFs. "
+            "Install it with: pip install Pillow")
+    import io
+    img = Image.open(io.BytesIO(data))
+    return np.asarray(img).tobytes()
+
+
+def jpeg_compress(data: bytes, width: int, height: int,
+                  samples: int = 1, quality: int = 75) -> bytes:
+    """Compress raw pixel data as JPEG. Requires Pillow."""
+    if not JPEG_AVAILABLE:
+        raise ImportError(
+            "Pillow is required to write JPEG-compressed TIFFs. "
+            "Install it with: pip install Pillow")
+    import io
+    if samples == 1:
+        arr = np.frombuffer(data, dtype=np.uint8).reshape(height, width)
+        img = Image.fromarray(arr, mode='L')
+    elif samples == 3:
+        arr = np.frombuffer(data, dtype=np.uint8).reshape(height, width, 3)
+        img = Image.fromarray(arr, mode='RGB')
+    else:
+        raise ValueError(f"JPEG compression requires 1 or 3 bands, got {samples}")
+    buf = io.BytesIO()
+    img.save(buf, format='JPEG', quality=quality)
+    return buf.getvalue()
+
+
 # -- Dispatch helpers ---------------------------------------------------------
 
 # TIFF compression tag values
 COMPRESSION_NONE = 1
 COMPRESSION_LZW = 5
+COMPRESSION_JPEG = 7
 COMPRESSION_DEFLATE = 8
+COMPRESSION_PACKBITS = 32773
 COMPRESSION_ADOBE_DEFLATE = 32946
 
 
-def decompress(data, compression: int, expected_size: int = 0) -> np.ndarray:
+def decompress(data, compression: int, expected_size: int = 0,
+               width: int = 0, height: int = 0, samples: int = 1) -> np.ndarray:
     """Decompress tile/strip data based on TIFF compression tag.
 
     Parameters
@@ -552,11 +662,14 @@ def decompress(data, compression: int, expected_size: int = 0) -> np.ndarray:
     if compression == COMPRESSION_NONE:
         return np.frombuffer(data, dtype=np.uint8)
     elif compression in (COMPRESSION_DEFLATE, COMPRESSION_ADOBE_DEFLATE):
-        # zlib returns bytes; wrap as read-only view (no copy)
         return np.frombuffer(deflate_decompress(data), dtype=np.uint8)
     elif compression == COMPRESSION_LZW:
-        # lzw_decompress already returns a mutable np.ndarray
         return lzw_decompress(data, expected_size)
+    elif compression == COMPRESSION_PACKBITS:
+        return np.frombuffer(packbits_decompress(data), dtype=np.uint8)
+    elif compression == COMPRESSION_JPEG:
+        return np.frombuffer(jpeg_decompress(data, width, height, samples),
+                             dtype=np.uint8)
     else:
         raise ValueError(f"Unsupported compression type: {compression}")
 
@@ -583,5 +696,9 @@ def compress(data: bytes, compression: int, level: int = 6) -> bytes:
         return deflate_compress(data, level)
     elif compression == COMPRESSION_LZW:
         return lzw_compress(data)
+    elif compression == COMPRESSION_PACKBITS:
+        return packbits_compress(data)
+    elif compression == COMPRESSION_JPEG:
+        raise ValueError("Use jpeg_compress() directly with width/height/samples")
     else:
         raise ValueError(f"Unsupported compression type: {compression}")

xrspatial/geotiff/_reader.py

@@ -155,10 +155,10 @@ def _read_strips(data: bytes, ifd: IFD, header: TIFFHeader,
 
         strip_data = data[offsets[strip_idx]:offsets[strip_idx] + byte_counts[strip_idx]]
         expected = strip_rows * width * samples * bytes_per_sample
-        chunk = decompress(strip_data, compression, expected)
+        chunk = decompress(strip_data, compression, expected,
+                           width=width, height=strip_rows, samples=samples)
 
         if pred in (2, 3):
-            # Predictor mutates in-place; copy if the array is read-only
             if not chunk.flags.writeable:
                 chunk = chunk.copy()
             chunk = _apply_predictor(chunk, pred, width, strip_rows, bytes_per_sample * samples)
@@ -266,7 +266,8 @@ def _read_tiles(data: bytes, ifd: IFD, header: TIFFHeader,
 
             tile_data = data[offsets[tile_idx]:offsets[tile_idx] + byte_counts[tile_idx]]
             expected = tw * th * samples * bytes_per_sample
-            chunk = decompress(tile_data, compression, expected)
+            chunk = decompress(tile_data, compression, expected,
+                               width=tw, height=th, samples=samples)
 
             if pred in (2, 3):
                 if not chunk.flags.writeable:
@@ -316,7 +317,7 @@ def _read_tiles(data: bytes, ifd: IFD, header: TIFFHeader,
 # ---------------------------------------------------------------------------
 
 def _read_cog_http(url: str, overview_level: int | None = None,
-                   band: int = 0) -> tuple[np.ndarray, GeoInfo]:
+                   band: int | None = None) -> tuple[np.ndarray, GeoInfo]:
     """Read a COG via HTTP range requests.
 
     Parameters
@@ -401,7 +402,8 @@ def _read_cog_http(url: str, overview_level: int | None = None,
 
             tile_data = source.read_range(off, bc)
             expected = tw * th * samples * bytes_per_sample
-            chunk = decompress(tile_data, compression, expected)
+            chunk = decompress(tile_data, compression, expected,
+                               width=tw, height=th, samples=samples)
 
             if pred in (2, 3):
                 if not chunk.flags.writeable:
@@ -431,7 +433,7 @@ def _read_cog_http(url: str, overview_level: int | None = None,
 # ---------------------------------------------------------------------------
 
 def read_to_array(source: str, *, window=None, overview_level: int | None = None,
-                  band: int = 0) -> tuple[np.ndarray, GeoInfo]:
+                  band: int | None = None) -> tuple[np.ndarray, GeoInfo]:
     """Read a GeoTIFF/COG to a numpy array.
 
     Parameters
@@ -483,7 +485,7 @@ def read_to_array(source: str, *, window=None, overview_level: int | None = None
             arr = _read_strips(data, ifd, header, dtype, window)
 
         # For multi-band with band selection, extract single band
-        if arr.ndim == 3 and ifd.samples_per_pixel > 1:
+        if arr.ndim == 3 and ifd.samples_per_pixel > 1 and band is not None:
             arr = arr[:, :, band]
     finally:
         src.close()