[SEDONA-350] Modify RS_AddBandFromArray

LICENSE

@@ -246,4 +246,9 @@ core/src/test/resources/primaryroads-linestring.csv
 core/src/test/resources/primaryroads-polygon.csv
 core/src/test/resources/primaryroads.csv
 core/src/test/resources/zcta510-small.csv
-core/src/test/resources/zcta510.csv
+core/src/test/resources/zcta510.csv
+
+
+MOD11C2 v006 from NASA
+---------------
+core/src/test/resources/raster/test5.TIFF

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

@@ -19,14 +19,19 @@
 package org.apache.sedona.common.raster;
 
 import org.apache.sedona.common.utils.RasterUtils;
+import org.geotools.coverage.Category;
 import org.geotools.coverage.GridSampleDimension;
 import org.geotools.coverage.grid.GridCoverage2D;
+import org.geotools.util.NumberRange;
+import org.opengis.util.InternationalString;
 
 import javax.media.jai.RasterFactory;
-import java.awt.Point;
+import java.awt.*;
 import java.awt.image.Raster;
 import java.awt.image.RenderedImage;
 import java.awt.image.WritableRaster;
+import java.util.ArrayList;
+import java.util.List;
 
 public class MapAlgebra
 {
@@ -68,6 +73,21 @@ public static double[] bandAsArray(GridCoverage2D rasterGeom, int bandIndex) {
      * @param bandIndex starts at 1, and no larger than numBands + 1
      * @return
      */
+    public static GridCoverage2D addBandFromArray(GridCoverage2D rasterGeom, double[] bandValues, int bandIndex, Double noDataValue) {
+        int numBands = rasterGeom.getNumSampleDimensions();
+        // Allow the band index to be one larger than the number of bands, which will append the band to the end
+        if (bandIndex < 1 || bandIndex > numBands + 1) {
+            throw new IllegalArgumentException("Band index is out of bounds. Must be between 1 and " + (numBands + 1) + ")");
+        }
+
+        if (bandIndex == numBands + 1) {
+            return copyRasterAndAppendBand(rasterGeom, bandValues, noDataValue);
+        }
+        else {
+            return copyRasterAndReplaceBand(rasterGeom, bandIndex, bandValues, noDataValue, true);
+        }
+    }
+
     public static GridCoverage2D addBandFromArray(GridCoverage2D rasterGeom, double[] bandValues, int bandIndex) {
         int numBands = rasterGeom.getNumSampleDimensions();
         // Allow the band index to be one larger than the number of bands, which will append the band to the end
@@ -101,7 +121,7 @@ public static GridCoverage2D addBandFromArray(GridCoverage2D rasterGeom, double[
      * @param bandValues
      * @return
      */
-    private static GridCoverage2D copyRasterAndAppendBand(GridCoverage2D gridCoverage2D, double[] bandValues) {
+    private static GridCoverage2D copyRasterAndAppendBand(GridCoverage2D gridCoverage2D, double[] bandValues, Double noDataValue) {
         // Get the original image and its properties
         RenderedImage originalImage = gridCoverage2D.getRenderedImage();
         Raster raster = originalImage.getData();
@@ -122,12 +142,25 @@ private static GridCoverage2D copyRasterAndAppendBand(GridCoverage2D gridCoverag
         GridSampleDimension[] originalSampleDimensions = gridCoverage2D.getSampleDimensions();
         GridSampleDimension[] sampleDimensions = new GridSampleDimension[numBand];
         System.arraycopy(originalSampleDimensions, 0, sampleDimensions, 0, originalSampleDimensions.length);
-        sampleDimensions[numBand - 1] = new GridSampleDimension("band" + numBand);
+        if (noDataValue != null) {
+            Category noDataCategory = new Category(
+                    Category.NODATA.getName(),
+                    new Color[] {new Color(0, 0, 0, 0)},
+                    NumberRange.create(noDataValue, noDataValue));
+            Category[] categories = new Category[] {noDataCategory};
+            sampleDimensions[numBand - 1] = new GridSampleDimension("band" + numBand, categories, null);
+        }else {
+            sampleDimensions[numBand - 1] = new GridSampleDimension("band" + numBand);
+        }
         // Construct a GridCoverage2D with the copied image.
         return RasterUtils.create(wr, gridCoverage2D.getGridGeometry(), sampleDimensions);
     }
 
-    private static GridCoverage2D copyRasterAndReplaceBand(GridCoverage2D gridCoverage2D, int bandIndex, double[] bandValues) {
+    private static GridCoverage2D copyRasterAndAppendBand(GridCoverage2D gridCoverage2D, double[] bandValues) {
+        return copyRasterAndAppendBand(gridCoverage2D, bandValues, null);
+    }
+
+    private static GridCoverage2D copyRasterAndReplaceBand(GridCoverage2D gridCoverage2D, int bandIndex, double[] bandValues, Double noDataValue, boolean removeNoDataIfNull) {
         // Do not allow the band index to be out of bounds
         if (bandIndex < 1 || bandIndex > gridCoverage2D.getNumSampleDimensions()) {
             throw new IllegalArgumentException("Band index is out of bounds. Must be between 1 and " + gridCoverage2D.getNumSampleDimensions() + ")");
@@ -144,7 +177,85 @@ private static GridCoverage2D copyRasterAndReplaceBand(GridCoverage2D gridCovera
                 wr.setPixel(i, j, bands);
             }
         }
+        GridSampleDimension[] originalSampleDimensions = gridCoverage2D.getSampleDimensions();
+        GridSampleDimension originalBandDimension = originalSampleDimensions[bandIndex - 1];
+        List<Category> categories =  (originalBandDimension.getCategories());
+        InternationalString defaultBandName = gridCoverage2D.getName();
+        Category defaultCategory = containsDefaultBand(categories, defaultBandName);
+        InternationalString bandToReplace = defaultCategory != null ? defaultBandName : null;
+        List<Category> newCategories = new ArrayList<>();
+        if (noDataValue == null && removeNoDataIfNull) {
+            if (originalBandDimension.getNoDataValues() != null) {
+                // remove noDataCategory
+                newCategories = createCategoryDeepCopy(categories, null, false, bandToReplace);
+            }
+        }else if (noDataValue != null) {
+            if (originalBandDimension.getNoDataValues() == null) {
+               // set a new noDataValue
+                newCategories = createCategoryDeepCopy(categories, noDataValue, true, bandToReplace);
+            }else {
+                double originalNoDataValue = originalBandDimension.getNoDataValues()[0];
+                if (originalNoDataValue != noDataValue) {
+                    //change noDataValue here
+                    newCategories = createCategoryDeepCopy(categories, noDataValue, false, bandToReplace);
+                }
+            }
+        }
+
         // Create a new GridCoverage2D with the copied image
-        return RasterUtils.create(wr, gridCoverage2D.getGridGeometry(), gridCoverage2D.getSampleDimensions());
+        if (!newCategories.isEmpty()) {
+            //create band with new categories altogether, no need to add default category
+            Category[] categoryArray = new Category[newCategories.size()];
+            categoryArray = newCategories.toArray(categoryArray);
+            originalSampleDimensions[bandIndex - 1] = new GridSampleDimension(originalBandDimension.getDescription(), categoryArray, null);
+            return RasterUtils.create(wr, gridCoverage2D.getGridGeometry(), originalSampleDimensions);
+        }else {
+            //Have to add default category
+            if (defaultCategory != null) {
+                //previous default category available, reuse it
+                originalSampleDimensions[bandIndex - 1] = new GridSampleDimension(originalBandDimension.getDescription(), new Category[]{defaultCategory}, null);
+                return RasterUtils.create(wr, gridCoverage2D.getGridGeometry(), originalSampleDimensions);
+            }else {
+                //default category not available, enforce its creation
+                return RasterUtils.create(wr, gridCoverage2D.getGridGeometry(), null);
+            }
+        }
+    }
+
+    private static GridCoverage2D copyRasterAndReplaceBand(GridCoverage2D gridCoverage2D, int bandIndex, double[] bandValues) {
+        return copyRasterAndReplaceBand(gridCoverage2D, bandIndex, bandValues, null, false);
+    }
+
+    private static Category containsDefaultBand(List<Category> categoryNames, InternationalString defaultBandName) {
+        for (Category category: categoryNames) {
+            if (category.getName().equals(defaultBandName)) return category;
+        }
+        return null;
+    }
+
+    private static List<Category> createCategoryDeepCopy(List<Category> categories, Double noDataValue, boolean addNoData, InternationalString defaultBandName) {
+        List<Category> newCategories = new ArrayList<>();
+        for (Category originalCategory : categories) {
+            Category newCategory;
+            if (originalCategory.getName().equals(Category.NODATA.getName())) {
+                //noDataValue removal is intended, skip adding this category to new categories
+                if (noDataValue == null) continue;
+                //noDataValue modification is intended, modify only the range with new noDataValue
+                newCategory = new Category(originalCategory.getName(), originalCategory.getColors(), NumberRange.create(noDataValue, noDataValue));
+            } else {
+                if (defaultBandName != null && originalCategory.getName().equals(defaultBandName))
+                    continue; //we do not need the default band if we're adding a custom band. Default band range can mess with newly added band range
+                newCategory = new Category(originalCategory.getName(), originalCategory.getColors(), originalCategory.getRange()); // add all other categories as is
+            }
+            newCategories.add(newCategory);
+        }
+        if (addNoData) {
+            // Addition of an earlier absent noDataValue is intended, manually add a new category to newCategories
+            newCategories.add(new Category(
+                    Category.NODATA.getName(),
+                    new Color[] {new Color(0, 0, 0, 0)},
+                    NumberRange.create(noDataValue, noDataValue)));
+        }
+        return newCategories;
     }
 }

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

@@ -39,8 +39,10 @@ public void testAddBandAsArrayAppend()
             band1[i] = i;
         }
         double[] band2 = new double[raster.getRenderedImage().getWidth() * raster.getRenderedImage().getHeight()];
+        double[] band3 = new double[raster.getRenderedImage().getWidth() * raster.getRenderedImage().getHeight()];
         for (int i = 0; i < band2.length; i++) {
             band2[i] = i * 2;
+            band3[i] = i * 3;
         }
         // Replace the first band
         GridCoverage2D rasterWithBand1 = MapAlgebra.addBandFromArray(raster, band1, 1);
@@ -49,20 +51,55 @@ public void testAddBandAsArrayAppend()
         assertEquals(raster.getCoordinateReferenceSystem2D(), rasterWithBand1.getCoordinateReferenceSystem2D());
         assertEquals(RasterAccessors.srid(raster), RasterAccessors.srid(rasterWithBand1));
 
-        // Append a new band
+        //replace the first band with a customNoDataValue
+        rasterWithBand1 = MapAlgebra.addBandFromArray(rasterWithBand1, band1, 1, -999d);
+        assertEquals(1, RasterAccessors.numBands(rasterWithBand1));
+        assertEquals(raster.getEnvelope(), rasterWithBand1.getEnvelope());
+        assertEquals(raster.getCoordinateReferenceSystem2D(), rasterWithBand1.getCoordinateReferenceSystem2D());
+        assertEquals(RasterAccessors.srid(raster), RasterAccessors.srid(rasterWithBand1));
+        assertEquals(-999, rasterWithBand1.getSampleDimension(0).getNoDataValues()[0], 1e-9);
+
+        //replace first band with a different customNoDataValue
+        rasterWithBand1 = MapAlgebra.addBandFromArray(rasterWithBand1, band1, 1, -9999d);
+        assertEquals(1, RasterAccessors.numBands(rasterWithBand1));
+        assertEquals(raster.getEnvelope(), rasterWithBand1.getEnvelope());
+        assertEquals(raster.getCoordinateReferenceSystem2D(), rasterWithBand1.getCoordinateReferenceSystem2D());
+        assertEquals(RasterAccessors.srid(raster), RasterAccessors.srid(rasterWithBand1));
+        assertEquals(-9999, rasterWithBand1.getSampleDimension(0).getNoDataValues()[0], 1e-9);
+
+        //remove noDataValue from the first band
+        rasterWithBand1 = MapAlgebra.addBandFromArray(rasterWithBand1, band1, 1, null);
+        assertEquals(1, RasterAccessors.numBands(rasterWithBand1));
+        assertEquals(raster.getEnvelope(), rasterWithBand1.getEnvelope());
+        assertEquals(raster.getCoordinateReferenceSystem2D(), rasterWithBand1.getCoordinateReferenceSystem2D());
+        assertEquals(RasterAccessors.srid(raster), RasterAccessors.srid(rasterWithBand1));
+        assertNull(rasterWithBand1.getSampleDimension(0).getNoDataValues());
+
+        // Append a new band with default noDataValue
         GridCoverage2D rasterWithBand2 = MapAlgebra.addBandFromArray(rasterWithBand1, band2);
         assertEquals(2, RasterAccessors.numBands(rasterWithBand2));
         assertEquals(raster.getEnvelope(), rasterWithBand2.getEnvelope());
         assertEquals(raster.getCoordinateReferenceSystem2D(), rasterWithBand2.getCoordinateReferenceSystem2D());
         assertEquals(RasterAccessors.srid(raster), RasterAccessors.srid(rasterWithBand2));
+        assertNull(rasterWithBand2.getSampleDimension(1).getNoDataValues());
+
+        // Append a new band with custom noDataValue
+        GridCoverage2D rasterWithBand3 = MapAlgebra.addBandFromArray(rasterWithBand2, band3, 3, 2d);
+        assertEquals(3, RasterAccessors.numBands(rasterWithBand3));
+        assertEquals(raster.getEnvelope(), rasterWithBand3.getEnvelope());
+        assertEquals(raster.getCoordinateReferenceSystem2D(), rasterWithBand3.getCoordinateReferenceSystem2D());
+        assertEquals(RasterAccessors.srid(raster), RasterAccessors.srid(rasterWithBand3));
+        assertEquals(2, rasterWithBand3.getSampleDimension(2).getNoDataValues()[0], 1e-9);
 
         // Check the value of the first band when use the raster with only one band
         double[] firstBand = MapAlgebra.bandAsArray(rasterWithBand1, 1);
         for (int i = 0; i < firstBand.length; i++) {
             assertEquals(i, firstBand[i], 0.1);
         }
         // Check the value of the first band when use the raster with two bands
-        firstBand = MapAlgebra.bandAsArray(rasterWithBand2, 1);
+
+        //Check the value of the first band when use the raster with three bands
+        firstBand = MapAlgebra.bandAsArray(rasterWithBand3, 1);
         for (int i = 0; i < firstBand.length; i++) {
             assertEquals(i, firstBand[i], 0.1);
         }
@@ -71,6 +108,12 @@ public void testAddBandAsArrayAppend()
         for (int i = 0; i < secondBand.length; i++) {
             assertEquals(i * 2, secondBand[i], 0.1);
         }
+
+        // Check the value of the third band
+        double[] thirdBand = MapAlgebra.bandAsArray(rasterWithBand3, 3);
+        for (int i = 0; i < secondBand.length; i++) {
+            assertEquals(i * 3, thirdBand[i], 0.1);
+        }
     }
 
     @Test

docs/api/sql/Raster-operators.md

@@ -563,20 +563,30 @@ Output:
 
 Introduction: Add a band to a raster from an array of doubles.
 
-Format: `RS_AddBandFromArray (raster: Raster, band: Array[Double], bandIndex:Int)`.
+Format: `RS_AddBandFromArray (raster: Raster, band: Array[Double])` | `RS_AddBandFromArray (raster: Raster, band: Array[Double], bandIndex:Int)` | `RS_AddBandFromArray (raster: Raster, band: Array[Double], bandIndex:Int, noDataValue:Double)`
 
 Since: `v1.4.1`
 
-The bandIndex is 1-based and must be between 1 and RS_NumBands(raster) + 1. It throws an exception if the bandIndex is out of range or the raster is null.
+The bandIndex is 1-based and must be between 1 and RS_NumBands(raster) + 1. It throws an exception if the bandIndex is out of range or the raster is null. If not specified, the noDataValue of the band is assumed to be null.
 
 When the bandIndex is RS_NumBands(raster) + 1, it appends the band to the end of the raster. Otherwise, it replaces the existing band at the bandIndex.
 
+If the bandIndex and noDataValue is not given, a convenience implementation adds a new band with a null noDataValue.
+
+Adding a new band with a custom noDataValue requires bandIndex = RS_NumBands(raster) + 1 and non-null noDataValue to be explicitly specified.
+
+Modifying or Adding a customNoDataValue is also possible by giving an existing band in RS_AddBandFromArray
+
+In order to remove an existing noDataValue from an existing band, pass null as the noDataValue in the RS_AddBandFromArray.
+
 Note that: `bandIndex == RS_NumBands(raster) + 1` is an experimental feature and might not lead to the loss of raster metadata and properties such as color models.
 
 SQL example:
 
 ```sql
+SELECT RS_AddBandFromArray(raster, RS_MultiplyFactor(RS_BandAsArray(RS_FromGeoTiff(content), 1), 2)) AS raster FROM raster_table
 SELECT RS_AddBandFromArray(raster, RS_MultiplyFactor(RS_BandAsArray(RS_FromGeoTiff(content), 1), 2), 1) AS raster FROM raster_table
+SELECT RS_AddBandFromArray(raster, RS_MultiplyFactor(RS_BandAsArray(RS_FromGeoTiff(content), 1), 2), 1, -999) AS raster FROM raster_table
 ```
 
 Output:

sql/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/raster/MapAlgebra.scala

@@ -805,7 +805,7 @@ case class RS_Append(inputExpressions: Seq[Expression])
 }
 
 case class RS_AddBandFromArray(inputExpressions: Seq[Expression])
-  extends InferredExpression(inferrableFunction3(MapAlgebra.addBandFromArray)) {
+  extends InferredExpression(inferrableFunction3(MapAlgebra.addBandFromArray), inferrableFunction4(MapAlgebra.addBandFromArray), inferrableFunction2(MapAlgebra.addBandFromArray)) {
   protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]) = {
     copy(inputExpressions = newChildren)
   }

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

@@ -375,7 +375,7 @@ class rasteralgebraTest extends TestBaseScala with BeforeAndAfter with GivenWhen
     }
 
     it("Passed RS_AsGeoTiff") {
-      val df = sparkSession.read.format("binaryFile").load(resourceFolder + "raster/*")
+      val df = sparkSession.read.format("binaryFile").load(resourceFolder + "raster/test1.tiff")
       val resultRaw = df.selectExpr("RS_FromGeoTiff(content) as raster").first().get(0)
       val resultLoaded = df.selectExpr("RS_FromGeoTiff(content) as raster")
         .selectExpr("RS_AsGeoTiff(raster) as geotiff")
@@ -474,7 +474,7 @@ class rasteralgebraTest extends TestBaseScala with BeforeAndAfter with GivenWhen
       assertEquals(true, resultOutOfBound)
     }
 
-    it("Passed RS_AddBandFromArray") {
+    it("Passed RS_AddBandFromArray - 3 parameters [default no data value]") {
       var df = sparkSession.read.format("binaryFile").load(resourceFolder + "raster/test1.tiff")
       df = df.selectExpr("RS_FromGeoTiff(content) as raster", "RS_MultiplyFactor(RS_BandAsArray(RS_FromGeoTiff(content), 1), 2) as band")
       val bandNewExpected:Seq[Double] = df.selectExpr("band").first().getSeq(0)
@@ -489,6 +489,24 @@ class rasteralgebraTest extends TestBaseScala with BeforeAndAfter with GivenWhen
       }
     }
 
+    it("Passed RS_AddBandFromArray - adding a new band with 2 parameter convenience function") {
+      var df = sparkSession.read.format("binaryFile").load(resourceFolder + "raster/test1.tiff")
+      df = df.selectExpr("RS_FromGeoTiff(content) as raster", "RS_MultiplyFactor(RS_BandAsArray(RS_FromGeoTiff(content), 1), 2) as band")
+      val bandNewExpected: Seq[Double] = df.selectExpr("band").first().getSeq(0)
+      val bandNewActual: Seq[Double] = df.selectExpr("RS_BandAsArray(RS_AddBandFromArray(raster, band), 2)").first().getSeq(0)
+      for (i <- bandNewExpected.indices) {
+        // The band value needs to be mod 256 because the ColorModel will mod 256.
+        assertEquals(bandNewExpected(i) % 256, bandNewActual(i), 0.001)
+      }
+    }
+
+    it("Passed RS_AddBandFromArray - adding a new band with a custom no data value]") {
+      var df = sparkSession.read.format("binaryFile").load(resourceFolder + "raster/test1.tiff")
+      df = df.selectExpr("RS_FromGeoTiff(content) as raster", "RS_MultiplyFactor(RS_BandAsArray(RS_FromGeoTiff(content), 1), 2) as band")
+      val raster = df.selectExpr("RS_AddBandFromArray(raster, band, 2, 2)").first().getAs[GridCoverage2D](0);
+      assertEquals(2, raster.getSampleDimension(1).getNoDataValues()(0), 1e-9)
+    }
+
     it("Passed RS_Intersects") {
       val df = sparkSession.read.format("binaryFile").load(resourceFolder + "raster/test1.tiff")
         .selectExpr("path", "RS_FromGeoTiff(content) as raster")