[GH-3330] Preserve other raster bands when replacing NoData

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

@@ -85,6 +85,7 @@ public static GridCoverage2D setBandNoDataValue(
       int width = RasterAccessors.getWidth(raster);
       WritableRaster wr =
           RasterFactory.createBandedRaster(dataTypeCode, width, height, numBands, null);
+      wr.setRect(rasterData);
       double[] bandData =
           rasterData.getSamples(0, 0, width, height, bandIndex - 1, (double[]) null);
       for (int i = 0; i < bandData.length; i++) {

common/src/test/java/org/apache/sedona/common/raster/RasterBandEditorsTest.java

@@ -164,6 +164,48 @@ public void testSetBandNoDataValueWithReplaceOption() throws FactoryException {
     assertEquals(actualMap.get(15.0), resultMap.get(20.0));
   }
 
+  @Test
+  public void testSetBandNoDataValueWithReplacePreservesOtherBands() throws FactoryException {
+    double[][] originalBands = {{1, 5, 3, 4}, {11, 5, 13, 14}, {21, 22, 5, 24}};
+    double[][] replacedBands = {{1, 99, 3, 4}, {11, 5, 99, 14}, {21, 22, 5, 99}};
+    double[] noDataValues = {5, 13, 24};
+    for (String dataType : new String[] {"b", "us", "s", "i", "f", "d"}) {
+      GridCoverage2D raster =
+          RasterConstructors.makeNonEmptyRaster(
+              3, dataType, 2, 2, 10, 20, 2, -3, 0.25, 0.5, 4326, originalBands);
+      for (int band = 1; band <= 3; band++) {
+        raster = RasterBandEditors.setBandNoDataValue(raster, band, noDataValues[band - 1]);
+      }
+
+      for (int targetBand = 1; targetBand <= 3; targetBand++) {
+        GridCoverage2D result =
+            RasterBandEditors.setBandNoDataValue(raster, targetBand, 99.0, true);
+        assertEquals(raster.getGridGeometry(), result.getGridGeometry());
+        assertEquals(
+            raster.getRenderedImage().getSampleModel().getDataType(),
+            result.getRenderedImage().getSampleModel().getDataType());
+        for (int band = 1; band <= 3; band++) {
+          String context = dataType + ", replacing band " + targetBand + ", checking band " + band;
+          assertArrayEquals(
+              context,
+              band == targetBand ? replacedBands[band - 1] : originalBands[band - 1],
+              MapAlgebra.bandAsArray(result, band),
+              0);
+          assertEquals(
+              context,
+              band == targetBand ? 99.0 : noDataValues[band - 1],
+              RasterBandAccessors.getBandNoDataValue(result, band),
+              0);
+          // Replacing pixels in the result must not mutate the input raster.
+          assertArrayEquals(
+              context, originalBands[band - 1], MapAlgebra.bandAsArray(raster, band), 0);
+          assertEquals(
+              noDataValues[band - 1], RasterBandAccessors.getBandNoDataValue(raster, band), 0);
+        }
+      }
+    }
+  }
+
   @Test
   public void testSetBandNoDataValueWithEmptyRaster() throws FactoryException {
     GridCoverage2D emptyRaster =

spark/common/src/test/scala/org/apache/sedona/sql/rasteralgebraTest.scala

@@ -18,7 +18,7 @@
  */
 package org.apache.sedona.sql
 
-import org.apache.sedona.common.raster.MapAlgebra
+import org.apache.sedona.common.raster.{MapAlgebra, RasterBandAccessors}
 import org.apache.sedona.common.utils.RasterUtils
 import org.apache.spark.sql.expressions.Window
 import org.apache.spark.sql.functions._
@@ -895,6 +895,27 @@ class rasteralgebraTest extends TestBaseScala with BeforeAndAfter with GivenWhen
       assertNull(sparkSession.sql("SELECT RS_SetBandNoDataValue(null, -999)").first().get(0))
     }
 
+    it("Passed RS_SetBandNoDataValue replacement preserves other bands") {
+      val input =
+        Seq((Seq(1.25, 5.0, 3.0, 4.0), Seq(11.0, 5.0, 13.0, 14.0), Seq(21.0, 22.0, 5.0, 24.5)))
+          .toDF("band1", "band2", "band3")
+      val raster = input.selectExpr(
+        "RS_AddBandFromArray(RS_AddBandFromArray(RS_AddBandFromArray(" +
+          "RS_MakeEmptyRaster(3, 'd', 2, 2, 0, 2, 1, -1, 0, 0, 4326), " +
+          "band1, 1, 5d), band2, 2, 13d), band3, 3, 24.5d) AS raster")
+      // Collect the raster itself to cover the four-argument SQL binding and serialization.
+      val result = raster
+        .selectExpr("RS_SetBandNoDataValue(raster, 2, -999d, true)")
+        .first()
+        .getAs[GridCoverage2D](0)
+      assert(MapAlgebra.bandAsArray(result, 1).toSeq == Seq(1.25, 5.0, 3.0, 4.0))
+      assert(MapAlgebra.bandAsArray(result, 2).toSeq == Seq(11.0, 5.0, -999.0, 14.0))
+      assert(MapAlgebra.bandAsArray(result, 3).toSeq == Seq(21.0, 22.0, 5.0, 24.5))
+      assertEquals(5.0, RasterBandAccessors.getBandNoDataValue(result, 1), 0)
+      assertEquals(-999.0, RasterBandAccessors.getBandNoDataValue(result, 2), 0)
+      assertEquals(24.5, RasterBandAccessors.getBandNoDataValue(result, 3), 0)
+    }
+
     it("Passed RS_SetBandNoDataValue clearing a band other than band 1") {
       // The clear consults the target band's no-data value; band 1 having none
       // must not short-circuit the removal on band 2.