feat: raster Python serde and with_bands() support

Add Python-side serialize() for InDbSedonaRaster, enabling Python UDFs
to return raster objects directly instead of the lossy .tolist() +
RS_MakeRaster workaround. Rasters now round-trip as contiguous bytes
preserving native dtypes and all metadata (CRS, nodata, affine, etc.).

Add with_bands() to InDbSedonaRaster for replacing pixel data (NumPy
array) while preserving spatial metadata. Band count and dtype may
differ from the source raster.

Add reconcileColorModel() to DeepCopiedRenderedImage (JVM) to fix
colorModel/sampleModel mismatches at deserialization when Python UDFs
change band count or dtype.

Cherry-picked from wherobots/wherobots-compute@e08bde1da08 with
vectorized UDF wiring excluded.

Author: prantogg

common/src/main/java/org/apache/sedona/common/raster/DeepCopiedRenderedImage.java

@@ -23,13 +23,16 @@
 import com.esotericsoftware.kryo.io.Input;
 import com.esotericsoftware.kryo.io.Output;
 import com.esotericsoftware.kryo.serializers.JavaSerializer;
+import com.sun.media.imageioimpl.common.BogusColorSpace;
 import com.sun.media.jai.rmi.ColorModelState;
 import com.sun.media.jai.util.ImageUtil;
 import it.geosolutions.jaiext.range.NoDataContainer;
 import java.awt.Image;
 import java.awt.Point;
 import java.awt.Rectangle;
+import java.awt.Transparency;
 import java.awt.image.ColorModel;
+import java.awt.image.ComponentColorModel;
 import java.awt.image.DataBuffer;
 import java.awt.image.Raster;
 import java.awt.image.RenderedImage;
@@ -39,6 +42,7 @@
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
 import java.io.Serializable;
+import java.util.Arrays;
 import java.util.Enumeration;
 import java.util.Hashtable;
 import java.util.Vector;
@@ -384,6 +388,11 @@ private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundE
     // The deserialized rendered image contains only one tile (imageRaster). We need to update
     // the sample model and tile properties to reflect this.
     this.sampleModel = this.imageRaster.getSampleModel();
+
+    // Reconcile colorModel with the actual sampleModel. Python UDFs may change the band count
+    // or data type, producing a sampleModel that is incompatible with the serialized colorModel.
+    this.colorModel = reconcileColorModel(this.colorModel, this.sampleModel);
+
     this.tileWidth = this.width;
     this.tileHeight = this.height;
     this.numXTiles = 1;
@@ -421,6 +430,40 @@ private Hashtable<String, Object> getSerializableProperties() {
     return propertyTable;
   }
 
+  /**
+   * Reconcile a deserialized ColorModel with the actual SampleModel. When a Python UDF changes the
+   * band count or data type, the serialized colorModel blob may describe the original raster's
+   * structure (e.g., 4 bands) while the SampleModel reflects the new structure (e.g., 1 band). This
+   * mismatch causes {@link javax.media.jai.PlanarImage#setImageLayout} to throw {@link
+   * IllegalArgumentException} when wrapping this image in a {@link
+   * javax.media.jai.RenderedImageAdapter}.
+   *
+   * <p>This method applies the same self-healing pattern used by {@link MapAlgebra#fetchColorModel}
+   * and {@link RasterUtils#clone}: check compatibility first, then try {@link
+   * PlanarImage#createColorModel}, then fall back to {@link BogusColorSpace}.
+   *
+   * @param colorModel the deserialized ColorModel (may be incompatible)
+   * @param sampleModel the actual SampleModel from the deserialized raster data
+   * @return a compatible ColorModel — either the original if compatible, or a newly constructed one
+   */
+  private static ColorModel reconcileColorModel(ColorModel colorModel, SampleModel sampleModel) {
+    if (colorModel.isCompatibleSampleModel(sampleModel)) {
+      return colorModel;
+    }
+    // Try PlanarImage's heuristic first (handles common cases like grayscale, RGB)
+    ColorModel cm = PlanarImage.createColorModel(sampleModel);
+    if (cm != null) {
+      return cm;
+    }
+    // Fall back to BogusColorSpace — always works for any band count / data type
+    int numBands = sampleModel.getNumBands();
+    int dataType = sampleModel.getDataType();
+    final int[] nBits = new int[numBands];
+    Arrays.fill(nBits, DataBuffer.getDataTypeSize(dataType));
+    return new ComponentColorModel(
+        new BogusColorSpace(numBands), nBits, false, true, Transparency.OPAQUE, dataType);
+  }
+
   public static void registerKryo(Kryo kryo) {
     kryo.register(ColorModelState.class, new JavaSerializer());
   }
@@ -486,6 +529,11 @@ public void read(Kryo kryo, Input input) {
     // The deserialized rendered image contains only one tile (imageRaster). We need to update
     // the sample model and tile properties to reflect this.
     this.sampleModel = this.imageRaster.getSampleModel();
+
+    // Reconcile colorModel with the actual sampleModel. Python UDFs may change the band count
+    // or data type, producing a sampleModel that is incompatible with the serialized colorModel.
+    this.colorModel = reconcileColorModel(this.colorModel, this.sampleModel);
+
     this.tileWidth = this.width;
     this.tileHeight = this.height;
     this.numXTiles = 1;

common/src/test/java/org/apache/sedona/common/raster/serde/SerdeTest.java

@@ -23,10 +23,13 @@
 import java.io.File;
 import java.io.IOException;
 import org.apache.sedona.common.raster.RasterConstructors;
+import org.apache.sedona.common.raster.RasterConstructorsForTesting;
 import org.apache.sedona.common.raster.RasterTestBase;
+import org.apache.sedona.common.utils.RasterUtils;
 import org.geotools.api.referencing.FactoryException;
 import org.geotools.coverage.grid.GridCoverage2D;
 import org.geotools.gce.geotiff.GeoTiffReader;
+import org.junit.Assert;
 import org.junit.Test;
 
 public class SerdeTest extends RasterTestBase {
@@ -85,4 +88,129 @@ private GridCoverage2D testRoundTrip(GridCoverage2D raster, int density)
     assertSameCoverage(raster, roundTripRaster, density);
     return roundTripRaster;
   }
+
+  @Test
+  public void testDeserializeWithMismatchedColorModel() throws Exception {
+    // Simulate what Python UDFs produce: a serialized raster where the colorModel blob
+    // describes 4 bands but the SampleModel/DataBuffer contain only 1 band. Before the
+    // ColorModel reconciliation fix in DeepCopiedRenderedImage.read(), this would fail with:
+    // IllegalArgumentException: "The specified ColorModel is incompatible with the image
+    // SampleModel"
+    // when GridCoverageFactory.create() wraps the image in RenderedImageAdapter → PlanarImage.
+
+    // 1. Create a 4-band raster and a 1-band raster with the same spatial dimensions
+    GridCoverage2D raster4band =
+        RasterConstructorsForTesting.makeRasterForTesting(
+            4, "F", "BandedSampleModel", 4, 3, 100, 100, 10, -10, 0, 0, 3857);
+    GridCoverage2D raster1band =
+        RasterConstructorsForTesting.makeRasterForTesting(
+            1, "F", "BandedSampleModel", 4, 3, 100, 100, 10, -10, 0, 0, 3857);
+
+    // 2. Serialize both
+    byte[] bytes4 = Serde.serialize(raster4band);
+    byte[] bytes1 = Serde.serialize(raster1band);
+
+    // 3. Locate the colorModel region in both byte arrays using UnsafeInput to navigate
+    //    the Kryo-encoded stream (same approach as Serde.extractPixelBanks).
+    int[] cmRegion4 = findColorModelRegion(bytes4);
+    int cmOffset4 = cmRegion4[0];
+    int cmEnd4 = cmRegion4[1];
+
+    int[] cmRegion1 = findColorModelRegion(bytes1);
+    int cmOffset1 = cmRegion1[0];
+    int cmEnd1 = cmRegion1[1];
+
+    // 4. Build spliced bytes: bytes1[0..cmOffset1) + bytes4[cmOffset4..cmEnd4) + bytes1[cmEnd1..)
+    //    This gives us: 1-band metadata, 4-band colorModel blob, 1-band raster data
+    byte[] spliced = new byte[cmOffset1 + (cmEnd4 - cmOffset4) + (bytes1.length - cmEnd1)];
+    System.arraycopy(bytes1, 0, spliced, 0, cmOffset1);
+    System.arraycopy(bytes4, cmOffset4, spliced, cmOffset1, cmEnd4 - cmOffset4);
+    System.arraycopy(
+        bytes1, cmEnd1, spliced, cmOffset1 + (cmEnd4 - cmOffset4), bytes1.length - cmEnd1);
+
+    // 5. Deserialize — this would throw IllegalArgumentException before the fix
+    GridCoverage2D deserialized = Serde.deserialize(spliced);
+    assertNotNull(deserialized);
+
+    // 6. Verify the deserialized raster has correct structure
+    Assert.assertEquals(1, deserialized.getNumSampleDimensions());
+    java.awt.image.Raster deserializedRaster =
+        RasterUtils.getRaster(deserialized.getRenderedImage());
+    Assert.assertEquals(1, deserializedRaster.getNumBands());
+    Assert.assertEquals(4, deserializedRaster.getWidth());
+    Assert.assertEquals(3, deserializedRaster.getHeight());
+
+    // Verify pixel values are from the 1-band raster (band=0, pixel[y][x] = y*4+x)
+    for (int y = 0; y < 3; y++) {
+      for (int x = 0; x < 4; x++) {
+        double expected = y * 4.0 + x;
+        Assert.assertEquals(
+            "Pixel mismatch at x=" + x + " y=" + y,
+            expected,
+            deserializedRaster.getSampleDouble(x, y, 0),
+            1e-6);
+      }
+    }
+
+    // 7. Verify the raster can be re-serialized (proves the colorModel is now valid)
+    byte[] reserialized = Serde.serialize(deserialized);
+    GridCoverage2D reDeserialized = Serde.deserialize(reserialized);
+    assertNotNull(reDeserialized);
+    Assert.assertEquals(1, reDeserialized.getNumSampleDimensions());
+
+    reDeserialized.dispose(true);
+    deserialized.dispose(true);
+    raster4band.dispose(true);
+    raster1band.dispose(true);
+  }
+
+  /**
+   * Find the byte range [startOffset, endOffset) of the colorModel length-prefixed section in a
+   * serialized IN_DB raster. Uses Kryo's UnsafeInput to navigate the stream correctly, mirroring
+   * the approach in {@link Serde#extractPixelBanks(byte[])}.
+   *
+   * @return int[2] where [0]=start offset (at length prefix), [1]=end offset (past data)
+   */
+  private int[] findColorModelRegion(byte[] bytes) {
+    try (com.esotericsoftware.kryo.io.UnsafeInput in =
+        new com.esotericsoftware.kryo.io.UnsafeInput(bytes)) {
+
+      // Skip rasterType byte
+      in.readByte();
+
+      // Skip name (UTF-8 string: int length + bytes)
+      KryoUtil.skipUTF8String(in);
+
+      // Skip gridEnvelope2D (4 ints = 16 bytes)
+      in.skip(16);
+
+      // Skip affine transform (6 doubles = 48 bytes)
+      in.skip(48);
+
+      // Skip CRS (length-prefixed bytes)
+      int crsLength = in.readInt();
+      in.skip(crsLength);
+
+      // Read bandCount and skip GridSampleDimensions
+      int bandCount = in.readInt();
+      for (int i = 0; i < bandCount; i++) {
+        GridSampleDimensionSerializer.skip(in);
+      }
+
+      // Skip DeepCopiedRenderedImage header (minX, minY, width, height = 4 ints)
+      in.skip(16);
+
+      // Skip properties (length-prefixed via writeObjectWithLength)
+      int propsLength = in.readInt();
+      in.skip(propsLength);
+
+      // Now at the colorModel length prefix
+      int cmStart = in.position();
+      int cmDataLength = in.readInt();
+      in.skip(cmDataLength);
+      int cmEnd = in.position();
+
+      return new int[] {cmStart, cmEnd};
+    }
+  }
 }

docs/tutorial/raster.md

@@ -707,8 +707,11 @@ band1 = ds.read(1)  # read the first band
 
 ## Writing Python UDF to work with raster data
 
-You can write Python UDFs to work with raster data in Python. The UDFs can take `SedonaRaster` objects as input and
-return any Spark data type as output. This is an example of a Python UDF that calculates the mean of the raster data.
+You can write Python UDFs that receive raster data, process it with NumPy, SciPy, or scikit-learn, and return either a scalar value or a new raster.
+
+### Raster to scalar
+
+UDFs can take `SedonaRaster` objects as input and return any Spark data type as output. This is an example of a Python UDF that calculates the mean of the raster data.
 
 ```python
 from pyspark.sql.types import DoubleType
@@ -730,35 +733,38 @@ df_raster.withColumn("mean", expr("mean_udf(rast)")).show()
 +--------------------+------------------+
 ```
 
-It is much trickier to write an UDF that returns a raster object, since Sedona does not support serializing Python raster
-objects yet. However, you can write a UDF that returns the band data as an array and then construct the raster object using
-`RS_MakeRaster`. This is an example of a Python UDF that creates a mask raster based on the first band of the input raster.
+### Raster to raster
+
+You can also write UDFs that return raster objects. Use `raster.with_bands()` to replace pixel data while preserving all spatial metadata (CRS, affine transform, nodata values, etc.). Band count and dtype can change freely.
+
+!!!warning
+    Raster-to-raster UDFs using `with_bands()` only work with **in-db rasters**. Out-db rasters do not carry the metadata needed to round-trip through a UDF. Use `RS_AsInDB` to convert out-db rasters to in-db before passing them to a raster-returning UDF.
 
 ```python
-from pyspark.sql.types import ArrayType, DoubleType
 import numpy as np
+from sedona.spark.sql.types import RasterType
 
 
 def mask_udf(raster):
     band1 = raster.as_numpy()[0, :, :]
-    mask = (band1 < 1400).astype(np.float64)
-    return mask.flatten().tolist()
+    mask = (band1 < 1400).astype(np.float32)
+    return raster.with_bands(mask)  # 1 band, preserves CRS/affine/nodata
 
 
-sedona.udf.register("mask_udf", band_udf, ArrayType(DoubleType()))
-df_raster.withColumn("mask", expr("mask_udf(rast)")).withColumn(
-    "mask_rast", expr("RS_MakeRaster(rast, 'I', mask)")
-).show()
+sedona.udf.register("mask_udf", mask_udf, RasterType())
+df_raster.withColumn("mask_rast", expr("mask_udf(rast)")).show()
 ```
 
 ```
-+--------------------+--------------------+--------------------+
-|                rast|                mask|           mask_rast|
-+--------------------+--------------------+--------------------+
-|GridCoverage2D["g...|[0.0, 0.0, 0.0, 0...|GridCoverage2D["g...|
-+--------------------+--------------------+--------------------+
++--------------------+--------------------+
+|                rast|           mask_rast|
++--------------------+--------------------+
+|GridCoverage2D["g...|GridCoverage2D["g...|
++--------------------+--------------------+
 ```
 
+The `with_bands()` method accepts a NumPy array in CHW order (bands × height × width) or HW order (height × width) for single-band output. The returned `SedonaRaster` carries all the original metadata and is automatically serialized back to the JVM when the UDF returns.
+
 ## Performance optimization
 
 When working with large raster datasets, refer to the [documentation on storing raster geometries in Parquet format](storing-blobs-in-parquet.md) for recommendations to optimize performance.

python/sedona/spark/raster/__init__.py

@@ -14,3 +14,5 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+
+from .sedona_raster import SedonaRaster  # noqa: F401