[SEDONA-449] Add two raster column support to RS_MapAlgebra (apache#1150

)

* feat: temp push of RS_MapAlgebra implementation

* feat: add multi band raster test

* feat: add 3 argument variant of MapAlgebra

* feat: port RS_MapAlgebra to spark and add test

* docs: add docs for two raster input RS_MapAlgebra

* chore: remove todos

* refactor: use tolerance variable in assert

* refactor: use self join

* refactor: remove redundant check

* docs: add doc to Raster-operators

* docs: add details on type casting

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

@@ -133,16 +133,10 @@ public static GridCoverage2D mapAlgebra(GridCoverage2D gridCoverage2D, String pi
         int rasterDataType = pixelType != null? RasterUtils.getDataTypeCode(pixelType) : renderedImage.getSampleModel().getDataType();
         int width = renderedImage.getWidth();
         int height = renderedImage.getHeight();
-        ColorModel originalColorModel = renderedImage.getColorModel();
         // ImageUtils.createConstantImage is slow, manually constructing a buffered image proved to be faster.
         // It also eliminates the data-copying overhead when converting raster data types after running jiffle script.
         WritableRaster resultRaster = RasterFactory.createBandedRaster(DataBuffer.TYPE_DOUBLE, width, height, 1, null);
-        ColorModel cm;
-        if (originalColorModel.isCompatibleRaster(resultRaster)) {
-            cm = originalColorModel;
-        }else {
-            cm = PlanarImage.createColorModel(resultRaster.getSampleModel());
-        }
+        ColorModel cm = fetchColorModel(renderedImage.getColorModel(), resultRaster);
         WritableRenderedImage resultImage = new BufferedImage(cm, resultRaster, false, null);
         try {
             String prevScript = previousScript.get();
@@ -180,6 +174,77 @@ public static GridCoverage2D mapAlgebra(GridCoverage2D gridCoverage2D, String pi
         }
     }
 
+    public static GridCoverage2D mapAlgebra(GridCoverage2D gridCoverage2D, String pixelType, String script) {
+        return mapAlgebra(gridCoverage2D, pixelType, script, null);
+    }
+
+    private static ColorModel fetchColorModel(ColorModel originalColorModel, WritableRaster resultRaster) {
+        if (originalColorModel.isCompatibleRaster(resultRaster)) {
+            return originalColorModel;
+        }else {
+            return PlanarImage.createColorModel(resultRaster.getSampleModel());
+        }
+    }
+
+    public static GridCoverage2D mapAlgebra(GridCoverage2D rast0, GridCoverage2D rast1, String pixelType, String script, Double noDataValue) {
+        if (rast0 == null || rast1 == null || script == null) {
+            return null;
+        }
+        RasterUtils.isRasterSameShape(rast0, rast1);
+
+        RenderedImage renderedImageRast0 = rast0.getRenderedImage();
+        int rasterDataType = pixelType != null ? RasterUtils.getDataTypeCode(pixelType) : renderedImageRast0.getSampleModel().getDataType();
+        int width = renderedImageRast0.getWidth();
+        int height = renderedImageRast0.getHeight();
+        // ImageUtils.createConstantImage is slow, manually constructing a buffered image proved to be faster.
+        // It also eliminates the data-copying overhead when converting raster data types after running jiffle script.
+        WritableRaster resultRaster = RasterFactory.createBandedRaster(DataBuffer.TYPE_DOUBLE, width, height, 1, null);
+
+        ColorModel cmRast0 = fetchColorModel(renderedImageRast0.getColorModel(), resultRaster);
+        RenderedImage renderedImageRast1 = rast1.getRenderedImage();
+
+        WritableRenderedImage resultImage = new BufferedImage(cmRast0, resultRaster, false, null);
+        try {
+            String prevScript = previousScript.get();
+            JiffleDirectRuntime prevRuntime = previousRuntime.get();
+            JiffleDirectRuntime runtime;
+            if (prevRuntime != null && script.equals(prevScript)) {
+                // Reuse the runtime to avoid recompiling the script
+                runtime = prevRuntime;
+                runtime.setSourceImage("rast0", renderedImageRast0);
+                runtime.setSourceImage("rast1", renderedImageRast1);
+                runtime.setDestinationImage("out", resultImage);
+                runtime.setDefaultBounds();
+            } else {
+                JiffleBuilder builder = new JiffleBuilder();
+                runtime = builder.script(script)
+                        .source("rast0", renderedImageRast0)
+                        .source("rast1", renderedImageRast1)
+                        .dest("out", resultImage)
+                        .getRuntime();
+                previousScript.set(script);
+                previousRuntime.set(runtime);
+            }
+
+            runtime.evaluateAll(null);
+
+            // If pixelType does not match with the data type of the result image (should be double since Jiffle only supports
+            // double destination image), we need to convert the resultImage to the specified pixel type.
+            if (rasterDataType != resultImage.getSampleModel().getDataType()) {
+                // Copy the resultImage to a new raster with the specified pixel type
+                WritableRaster convertedRaster = RasterFactory.createBandedRaster(rasterDataType, width, height, 1, null);
+                double[] samples = resultRaster.getSamples(0, 0, width, height, 0, (double[]) null);
+                convertedRaster.setSamples(0, 0, width, height, 0, samples);
+                return RasterUtils.clone(convertedRaster, null, rast0, noDataValue, false);
+            } else {
+                // build a new GridCoverage2D from the resultImage
+                return RasterUtils.clone(resultImage, null, rast0, noDataValue, false);
+            }
+        } catch (Exception e) {
+            throw new RuntimeException("Failed to run map algebra", e);
+        }
+    }
+
     /**
      * @param band1 band values
      * @param band2 band values

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

@@ -20,7 +20,6 @@
 
 import org.apache.commons.lang3.ArrayUtils;
 import org.apache.commons.lang3.tuple.Pair;
-import org.apache.sedona.common.Functions;
 import org.apache.sedona.common.utils.RasterUtils;
 import org.geotools.coverage.GridSampleDimension;
 import org.geotools.coverage.grid.GridCoverage2D;
@@ -99,7 +98,7 @@ public static GridCoverage2D setBandNoDataValue(GridCoverage2D raster, Double no
     public static GridCoverage2D addBand(GridCoverage2D toRaster, GridCoverage2D fromRaster, int fromBand, int toRasterIndex) {
         RasterUtils.ensureBand(fromRaster, fromBand);
         ensureBandAppend(toRaster, toRasterIndex);
-        isRasterSameShape(toRaster, fromRaster);
+        RasterUtils.isRasterSameShape(toRaster, fromRaster);
 
         int width = RasterAccessors.getWidth(toRaster), height = RasterAccessors.getHeight(toRaster);
 
@@ -160,22 +159,6 @@ private static void ensureBandAppend(GridCoverage2D raster, int band) {
         }
     }
 
-    /**
-     * Check if the two rasters are of the same shape
-     * @param raster1
-     * @param raster2
-     */
-    private static void isRasterSameShape(GridCoverage2D raster1, GridCoverage2D raster2) {
-        int width1 = RasterAccessors.getWidth(raster1), height1 = RasterAccessors.getHeight(raster1);
-        int width2 = RasterAccessors.getWidth(raster2), height2 = RasterAccessors.getHeight(raster2);
-
-        if (width1 != width2 && height1 != height2) {
-            throw new IllegalArgumentException(String.format("Provided rasters are not of same shape. \n" +
-                    "First raster having width of %d and height of %d. \n" +
-                    "Second raster having width of %d and height of %d", width1, height1, width2, height2));
-        }
-    }
-
     /**
      * Return a clipped raster with the specified ROI by the geometry
      * @param raster Raster to clip

common/src/main/java/org/apache/sedona/common/utils/RasterUtils.java

@@ -626,4 +626,20 @@ public static GridCoverage2D copyRasterAndReplaceBand(GridCoverage2D gridCoverag
     public static GridCoverage2D copyRasterAndReplaceBand(GridCoverage2D gridCoverage2D, int bandIndex, Number[] bandValues) {
         return copyRasterAndReplaceBand(gridCoverage2D, bandIndex, bandValues, null, false);
     }
+
+    /**
+     * Check if the two rasters are of the same shape
+     * @param raster1
+     * @param raster2
+     */
+    public static void isRasterSameShape(GridCoverage2D raster1, GridCoverage2D raster2) {
+        int width1 = RasterAccessors.getWidth(raster1), height1 = RasterAccessors.getHeight(raster1);
+        int width2 = RasterAccessors.getWidth(raster2), height2 = RasterAccessors.getHeight(raster2);
+
+        if (width1 != width2 && height1 != height2) {
+            throw new IllegalArgumentException(String.format("Provided rasters are not of same shape. \n" +
+                    "First raster having width of %d and height of %d. \n" +
+                    "Second raster having width of %d and height of %d", width1, height1, width2, height2));
+        }
+    }
 }

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

@@ -436,6 +436,126 @@ public void testCountValue() {
         assertEquals(expected, actual);
     }
 
+    @Test
+    public void testMapAlgebra2Rasters() throws FactoryException {
+        Random random = new Random();
+        String[] pixelTypes = {null, "b", "i", "s", "us", "f", "d"};
+        for (String pixelType : pixelTypes) {
+            int width = random.nextInt(100) + 10;
+            int height = random.nextInt(100) + 10;
+            testMapAlgebra2Rasters(width, height, pixelType, null);
+            testMapAlgebra2Rasters(width, height, pixelType, 100.0);
+            testMapAlgebra2RastersMultiBand(width, height, pixelType, null);
+            testMapAlgebra2RastersMultiBand(width, height, pixelType, 100.0);
+        }
+    }
+
+    private void testMapAlgebra2RastersMultiBand(int width, int height, String pixelType, Double noDataValue) throws FactoryException {
+        GridCoverage2D rast0 = RasterConstructors.makeEmptyRaster(2, "b", width, height, 10, 20, 1);
+        GridCoverage2D rast1 = RasterConstructors.makeEmptyRaster(2, "b", width, height, 10, 20, 1);
+        double[] band1 = new double[width * height];
+        double[] band2 = new double[width * height];
+        double[] band3 = new double[width * height];
+        double[] band4 = new double[width * height];
+        for (int i = 0; i < band1.length; i++) {
+            band1[i] = Math.random() * 10;
+            band2[i] = Math.random() * 10;
+            band3[i] = Math.random() * 10;
+            band4[i] = Math.random() * 10;
+        }
+        rast0 = MapAlgebra.addBandFromArray(rast0, band1, 1);
+        rast0 = MapAlgebra.addBandFromArray(rast0, band2, 2);
+        rast1 = MapAlgebra.addBandFromArray(rast1, band3, 1);
+        rast1 = MapAlgebra.addBandFromArray(rast1, band4, 2);
+        GridCoverage2D result = MapAlgebra.mapAlgebra(rast0, rast1, pixelType, "out = (rast0[0] + rast0[1] + rast1[0] + rast1[1]) * 0.4;", noDataValue);
+        double actualNoDataValue = RasterUtils.getNoDataValue(result.getSampleDimension(0));
+        if (noDataValue != null) {
+            Assert.assertEquals(noDataValue, actualNoDataValue, 1e-9);
+        } else {
+            Assert.assertTrue(Double.isNaN(actualNoDataValue));
+        }
+
+        int resultDataType = result.getRenderedImage().getSampleModel().getDataType();
+        int expectedDataType;
+        if (pixelType != null) {
+            expectedDataType = RasterUtils.getDataTypeCode(pixelType);
+        } else {
+            expectedDataType = rast0.getRenderedImage().getSampleModel().getDataType();
+        }
+        Assert.assertEquals(expectedDataType, resultDataType);
+
+        Assert.assertEquals(rast0.getGridGeometry().getGridToCRS2D(), result.getGridGeometry().getGridToCRS2D());
+        band1 = MapAlgebra.bandAsArray(rast0, 1);
+        band2 = MapAlgebra.bandAsArray(rast0, 2);
+        band3 = MapAlgebra.bandAsArray(rast1, 1);
+        band4 = MapAlgebra.bandAsArray(rast1, 2);
+        double[] bandResult = MapAlgebra.bandAsArray(result, 1);
+        Assert.assertEquals(band1.length, bandResult.length);
+        for (int i = 0; i < band1.length; i++) {
+            double expected = (band1[i] + band2[i] + band3[i] + band4[i]) * 0.4;
+            double actual = bandResult[i];
+            switch (resultDataType) {
+                case DataBuffer.TYPE_BYTE:
+                case DataBuffer.TYPE_SHORT:
+                case DataBuffer.TYPE_USHORT:
+                case DataBuffer.TYPE_INT:
+                    Assert.assertEquals((int) expected, (int) actual);
+                    break;
+                default:
+                    Assert.assertEquals(expected, actual, FP_TOLERANCE);
+            }
+        }
+    }
+
+    private void testMapAlgebra2Rasters(int width, int height, String pixelType, Double noDataValue) throws FactoryException {
+        GridCoverage2D rast0 = RasterConstructors.makeEmptyRaster(1, "b", width, height, 10, 20, 1);
+        GridCoverage2D rast1 = RasterConstructors.makeEmptyRaster(1, "b", width, height, 10, 20, 1);
+        double[] band1 = new double[width * height];
+        double[] band2 = new double[width * height];
+        for (int i = 0; i < band1.length; i++) {
+            band1[i] = Math.random() * 10;
+            band2[i] = Math.random() * 10;
+        }
+        rast0 = MapAlgebra.addBandFromArray(rast0, band1, 1);
+        rast1 = MapAlgebra.addBandFromArray(rast1, band2, 1);
+        GridCoverage2D result = MapAlgebra.mapAlgebra(rast0, rast1, pixelType, "out = (rast0[0] + rast1[0]) * 0.4;", noDataValue);
+        double actualNoDataValue = RasterUtils.getNoDataValue(result.getSampleDimension(0));
+        if (noDataValue != null) {
+            Assert.assertEquals(noDataValue, actualNoDataValue, 1e-9);
+        } else {
+            Assert.assertTrue(Double.isNaN(actualNoDataValue));
+        }
+
+        int resultDataType = result.getRenderedImage().getSampleModel().getDataType();
+        int expectedDataType;
+        if (pixelType != null) {
+            expectedDataType = RasterUtils.getDataTypeCode(pixelType);
+        } else {
+            expectedDataType = rast0.getRenderedImage().getSampleModel().getDataType();
+        }
+        Assert.assertEquals(expectedDataType, resultDataType);
+
+        Assert.assertEquals(rast0.getGridGeometry().getGridToCRS2D(), result.getGridGeometry().getGridToCRS2D());
+        band1 = MapAlgebra.bandAsArray(rast0, 1);
+        band2 = MapAlgebra.bandAsArray(rast1, 1);
+        double[] bandResult = MapAlgebra.bandAsArray(result, 1);
+        Assert.assertEquals(band1.length, bandResult.length);
+        for (int i = 0; i < band1.length; i++) {
+            double expected = (band1[i] + band2[i]) * 0.4;
+            double actual = bandResult[i];
+            switch (resultDataType) {
+                case DataBuffer.TYPE_BYTE:
+                case DataBuffer.TYPE_SHORT:
+                case DataBuffer.TYPE_USHORT:
+                case DataBuffer.TYPE_INT:
+                    Assert.assertEquals((int) expected, (int) actual);
+                    break;
+                default:
+                    Assert.assertEquals(expected, actual, FP_TOLERANCE);
+            }
+        }
+    }
+
     @Test
     public void testMapAlgebra() throws FactoryException {
         Random random = new Random();

docs/api/sql/Raster-map-algebra.md

@@ -7,16 +7,49 @@ Apache Sedona provides two ways to perform map algebra operations:
 1. Using the `RS_MapAlgebra` function.
 2. Using `RS_BandAsArray` and array based map algebra functions, such as `RS_Add`, `RS_Multiply`, etc.
 
-Generally, the `RS_MapAlgebra` function is more flexible and can be used to perform more complex operations. The `RS_MapAlgebra(rast, pixelType, script, [noDataValue])` function takes three to four arguments:
+Generally, the `RS_MapAlgebra` function is more flexible and can be used to perform more complex operations. The function takes three to four arguments:
+
+```sql
+RS_MapAlgebra(rast: Raster, pixelType: String, script: String, [noDataValue: Double])
+```
 
 * `rast`: The raster to apply the map algebra expression to.
 * `pixelType`: The data type of the output raster. This can be one of `D` (double), `F` (float), `I` (integer), `S` (short), `US` (unsigned short) or `B` (byte). If specified `NULL`, the output raster will have the same data type as the input raster.
-* `script`: The map algebra script.
+* `script`: The map algebra script. [Refer here for more details on the format.](#:~:text=The Jiffle script is,current output pixel value)
 * `noDataValue`: (Optional) The nodata value of the output raster.
 
+As of version `v1.5.1`, the `RS_MapAlgebra` function allows two raster column inputs, with multi-band rasters supported. The function accepts 5 parameters:
+
+```sql
+RS_MapAlgebra(rast0: Raster, rast1: Raster, pixelType: String, script: String, noDataValue: Double)
+```
+
+* `rast0`: The first raster to apply the map algebra expression to.
+* `rast1`: The second raster to apply the map algebra expression to.
+* `pixelType`: The data type of the output raster. This can be one of `D` (double), `F` (float), `I` (integer), `S` (short), `US` (unsigned short) or `B` (byte). If specified `NULL`, the output raster will have the same data type as the input raster.
+* `script`: The map algebra script. [Refer here for more details on the format.](#:~:text=The Jiffle script is,current output pixel value)
+* `noDataValue`: (Not optional) The nodata value of the output raster, `null` is allowed.
+
+Spark SQL Example for two raster input `RS_MapAlgebra`:
+
+```sql
+RS_MapAlgebra(rast0, rast1, 'D', 'out = rast0[0] * 0.5 + rast1[0] * 0.5;', null)
+```
+
 `RS_MapAlgebra` also has good performance, since it is backed by [Jiffle](https://github.com/geosolutions-it/jai-ext/wiki/Jiffle) and can be compiled to Java bytecode for
 execution. We'll demonstrate both approaches to implementing commonly used map algebra operations.
 
+!!!Note
+    The `RS_MapAlgebra` function can cast the output raster to a different data type specified by `pixelType`:
+
+    - If `pixelType` is smaller than the input raster data type, narrowing casts will be performed, which may result in loss of data.
+
+    - If `pixelType` is larger, widening casts will retain data accuracy.
+
+    - If `pixelType` matches the input raster data type, no casting occurs.
+
+    This allows controlling the output pixel data type. Users should consider potential precision impacts when coercing to a smaller type. 
+
 ### NDVI
 
 The Normalized Difference Vegetation Index (NDVI) is a simple graphical indicator that can be used to analyze remote sensing measurements, typically, but not necessarily, from a space platform, and assess whether the target being observed contains live green vegetation or not. NDVI has become a de facto standard index used to determine whether a given area contains live green vegetation or not. The NDVI is calculated from these individual measurements as follows:

docs/api/sql/Raster-operators.md

@@ -2202,12 +2202,18 @@ Introduction: Apply a map algebra script on a raster.
 
 Format: 
 
-`RS_MapAlgebra (raster: Raster, pixelType: String, script: String)`
+```
+RS_MapAlgebra (raster: Raster, pixelType: String, script: String)
+```
 
 ```
 RS_MapAlgebra (raster: Raster, pixelType: String, script: String, noDataValue: Double)
 ```
 
+```
+RS_MapAlgebra(rast0: Raster, rast1: Raster, pixelType: String, script: String, noDataValue: Double)
+```
+
 Since: `v1.5.0`
 
 `RS_MapAlgebra` runs a script on a raster. The script is written in a map algebra language called [Jiffle](https://github.com/geosolutions-it/jai-ext/wiki/Jiffle). The script takes a raster
@@ -2232,6 +2238,12 @@ Output:
 +--------------------+
 ```
 
+Spark SQL Example for two raster input `RS_MapAlgebra`:
+
+```sql
+RS_MapAlgebra(rast0, rast1, 'D', 'out = rast0[0] * 0.5 + rast1[0] * 0.5;', null)
+```
+
 For more details and examples about `RS_MapAlgebra`, please refer to the [Map Algebra documentation](../Raster-map-algebra/).
 To learn how to write map algebra script, please refer to [Jiffle language summary](https://github.com/geosolutions-it/jai-ext/wiki/Jiffle---language-summary).
 

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

@@ -193,7 +193,9 @@ case class RS_BandAsArray(inputExpressions: Seq[Expression]) extends InferredExp
 }
 
 case class RS_MapAlgebra(inputExpressions: Seq[Expression])
-  extends InferredExpression(nullTolerantInferrableFunction4(MapAlgebra.mapAlgebra)) {
+  extends InferredExpression(nullTolerantInferrableFunction3(MapAlgebra.mapAlgebra),
+    nullTolerantInferrableFunction4(MapAlgebra.mapAlgebra),
+    nullTolerantInferrableFunction5(MapAlgebra.mapAlgebra)) {
 
   protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]) = {
     copy(inputExpressions = newChildren)