[SEDONA-448] RS_SetBandNoDataValue should have replace option (apac…

…he#1160) * feat: add replace option, not implementation * feat: add replace option and testing * feat: port new function to spark * docs: add docs for replace option * chore: clean imports * refactor: remove duplicate variable
Kontinuation · Dec 22, 2023 · f41d908 · f41d908
1 parent 73741be
commit f41d908
Show file tree

Hide file tree

Showing 4 changed files with 113 additions and 19 deletions.
diff --git a/common/src/main/java/org/apache/sedona/common/raster/RasterBandEditors.java b/common/src/main/java/org/apache/sedona/common/raster/RasterBandEditors.java
@@ -23,12 +23,8 @@
 import org.apache.sedona.common.utils.RasterUtils;
 import org.geotools.coverage.GridSampleDimension;
 import org.geotools.coverage.grid.GridCoverage2D;
-import org.geotools.coverage.grid.GridEnvelope2D;
-import org.geotools.coverage.grid.GridGeometry2D;
 import org.geotools.coverage.processing.operation.Crop;
-import org.geotools.referencing.operation.transform.AffineTransform2D;
 import org.locationtech.jts.geom.Geometry;
-import org.opengis.metadata.spatial.PixelOrientation;
 import org.opengis.parameter.ParameterValueGroup;
 import org.opengis.referencing.FactoryException;
 import org.opengis.referencing.operation.TransformException;
@@ -46,9 +42,10 @@ public class RasterBandEditors {
      * @param raster Source raster to add no-data value
      * @param bandIndex Band index to add no-data value
      * @param noDataValue Value to set as no-data value, if null then remove existing no-data value
+     * @param replace if true replaces the previous no-data value with the specified no-data value
      * @return Raster with no-data value
      */
-    public static GridCoverage2D setBandNoDataValue(GridCoverage2D raster, int bandIndex, Double noDataValue) {
+    public static GridCoverage2D setBandNoDataValue(GridCoverage2D raster, int bandIndex, Double noDataValue, boolean replace) {
         RasterUtils.ensureBand(raster, bandIndex);
         Double rasterNoData = RasterBandAccessors.getBandNoDataValue(raster, bandIndex);
 
@@ -62,30 +59,55 @@ public static GridCoverage2D setBandNoDataValue(GridCoverage2D raster, int bandI
             return RasterUtils.clone(raster.getRenderedImage(), null, sampleDimensions, raster, null, true);
         }
 
-        if ( !(rasterNoData == null) && rasterNoData.equals(noDataValue)) {
+        if ( rasterNoData != null && rasterNoData.equals(noDataValue)) {
             return raster;
         }
         GridSampleDimension[] bands = raster.getSampleDimensions();
         bands[bandIndex - 1] = RasterUtils.createSampleDimensionWithNoDataValue(bands[bandIndex - 1], noDataValue);
 
-        int width = RasterAccessors.getWidth(raster), height = RasterAccessors.getHeight(raster);
-        AffineTransform2D affine = RasterUtils.getGDALAffineTransform(raster);
-        GridGeometry2D gridGeometry2D = new GridGeometry2D(
-                new GridEnvelope2D(0, 0, width, height),
-                PixelOrientation.UPPER_LEFT,
-                affine, raster.getCoordinateReferenceSystem2D(), null
-        );
+        if (replace) {
+            if (rasterNoData == null) {
+                throw new IllegalArgumentException("The raster provided doesn't have a no-data value. Please provide a raster that has a no-data value to use `replace` option.");
+            }
+
+            Raster rasterData = RasterUtils.getRaster(raster.getRenderedImage());
+            int dataTypeCode = rasterData.getDataBuffer().getDataType();
+            int numBands = RasterAccessors.numBands(raster);
+            int height = RasterAccessors.getHeight(raster);
+            int width = RasterAccessors.getWidth(raster);
+            WritableRaster wr = RasterFactory.createBandedRaster(dataTypeCode, width, height, numBands, null);
+            double[] bandData = rasterData.getSamples(0, 0, width, height, bandIndex - 1, (double[]) null);
+            for (int i = 0; i < bandData.length; i++) {
+                if (bandData[i] == rasterNoData) {
+                    bandData[i] = noDataValue;
+                }
+            }
+            wr.setSamples(0, 0, width, height, bandIndex - 1, bandData);
+            return RasterUtils.clone(wr, null, bands, raster, null, true);
+        }
+
         return RasterUtils.clone(raster.getRenderedImage(), null, bands, raster, null, true);
     }
 
+    /**
+     * Adds no-data value to the raster.
+     * @param raster Source raster to add no-data value
+     * @param bandIndex Band index to add no-data value
+     * @param noDataValue Value to set as no-data value, if null then remove existing no-data value
+     * @return Raster with no-data value
+     */
+    public static GridCoverage2D setBandNoDataValue(GridCoverage2D raster, int bandIndex, Double noDataValue) {
+        return setBandNoDataValue(raster, bandIndex, noDataValue, false);
+    }
+
     /**
      * Adds no-data value to the raster.
      * @param raster Source raster to add no-data value
      * @param noDataValue Value to set as no-data value, if null then remove existing no-data value
      * @return Raster with no-data value
      */
     public static GridCoverage2D setBandNoDataValue(GridCoverage2D raster, Double noDataValue) {
-        return setBandNoDataValue(raster, 1, noDataValue);
+        return setBandNoDataValue(raster, 1, noDataValue, false);
     }
 
     /**

diff --git a/common/src/test/java/org/apache/sedona/common/raster/RasterBandEditorsTest.java b/common/src/test/java/org/apache/sedona/common/raster/RasterBandEditorsTest.java
@@ -19,21 +19,21 @@
 package org.apache.sedona.common.raster;
 
 import org.apache.sedona.common.Constructors;
-import org.apache.sedona.common.Functions;
-import org.apache.sedona.common.FunctionsGeoTools;
-import org.apache.sedona.common.utils.RasterUtils;
 import org.geotools.coverage.grid.GridCoverage2D;
 import org.junit.Test;
 import org.locationtech.jts.geom.Geometry;
 import org.locationtech.jts.io.ParseException;
 import org.opengis.referencing.FactoryException;
 import org.opengis.referencing.operation.TransformException;
 
-import java.awt.geom.Point2D;
 import java.io.IOException;
+import java.text.DecimalFormat;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.stream.Collectors;
 
 import static org.junit.Assert.*;
 
@@ -64,6 +64,60 @@ public void testSetBandNoDataValueWithNull() throws IOException {
         assertEquals(expected, actual);
     }
 
+    @Test
+    public void testSetBandNoDataValueWithReplaceOptionRaster() throws IOException {
+        GridCoverage2D raster = rasterFromGeoTiff(resourceFolder + "raster/raster_with_no_data/test5.tiff");
+        double[] originalSummary = RasterBandAccessors.getSummaryStats(raster, 1, false);
+        int sumOG = (int) originalSummary[1];
+
+        assertEquals(206233487, sumOG);
+        GridCoverage2D resultRaster = RasterBandEditors.setBandNoDataValue(raster, 1, 10.0, true);
+        double[] resultSummary = RasterBandAccessors.getSummaryStats(resultRaster, 1, false);
+        int sumActual = (int) resultSummary[1];
+
+        // 108608 is the total no-data values in the raster
+        // 10.0 is the new no-data value
+        int sumExpected = sumOG + (10 * 108608);
+        assertEquals(sumExpected, sumActual);
+
+        // Not replacing previous no-data value
+        resultRaster = RasterBandEditors.setBandNoDataValue(raster, 1, 10.0);
+        resultSummary = RasterBandAccessors.getSummaryStats(resultRaster, 1, false);
+        sumActual = (int) resultSummary[1];
+        assertEquals(sumOG, sumActual);
+    }
+
+    @Test
+    public void testSetBandNoDataValueWithReplaceOption() throws FactoryException {
+        GridCoverage2D raster = RasterConstructors.makeEmptyRaster(1, "d", 10, 20, 10, 20, 1);
+        double[] band1 = new double[200];
+        DecimalFormat df = new DecimalFormat("0.00");
+        for (int i = 0; i < band1.length; i++) {
+            if (i % 3 == 0) {
+                band1[i] = 15;
+                continue;
+            }
+            band1[i] = Double.parseDouble(df.format(Math.random() * 10));
+        }
+        raster = MapAlgebra.addBandFromArray(raster, band1, 1);
+        // setting the noData property
+        raster = RasterBandEditors.setBandNoDataValue(raster, 1, 15.0);
+
+        // invoking replace option.
+        GridCoverage2D result = RasterBandEditors.setBandNoDataValue(raster, 1, 20.0, true);
+        double[] resultBand = MapAlgebra.bandAsArray(result, 1);
+
+        Map<Double, Long> resultMap = Arrays.stream(resultBand)
+                .boxed()
+                .collect(Collectors.groupingBy(Function.identity(), Collectors.counting()));
+
+        Map<Double, Long> actualMap = Arrays.stream(band1)
+                .boxed()
+                .collect(Collectors.groupingBy(Function.identity(), Collectors.counting()));
+
+        assertEquals(actualMap.get(15.0), resultMap.get(20.0));
+    }
+
     @Test
     public void testSetBandNoDataValueWithEmptyRaster() throws FactoryException {
         GridCoverage2D emptyRaster = RasterConstructors.makeEmptyRaster(1, 20, 20, 0, 0, 8, 8, 0.1, 0.1, 4326);

diff --git a/docs/api/sql/Raster-operators.md b/docs/api/sql/Raster-operators.md
@@ -1583,7 +1583,24 @@ Output:
 
 Introduction: This sets the no data value for a specified band in the raster. If the band index is not provided, band 1 is assumed by default. Passing a `null` value for `noDataValue` will remove the no data value and that will ensure all pixels are included in functions rather than excluded as no data.
 
-Format: `RS_SetBandNoDataValue(raster: Raster, bandIndex: Integer = 1, noDataValue: Double)`
+Since `v1.5.1`, this function supports the ability to replace the current no-data value with the new `noDataValue`. 
+
+!!!Note
+    When `replace` is true, any pixels matching the provided `noDataValue` will be considered as no-data in the output raster.
+
+    An `IllegalArgumentException` will be thrown if the input raster does not already have a no-data value defined. Replacing existing values with `noDataValue` requires a defined no-data baseline to evaluate against.
+
+    To use this for no-data replacement, the input raster must first set its no-data value, which can then be selectively replaced via this function.
+
+Format: 
+
+```
+RS_SetBandNoDataValue(raster: Raster, bandIndex: Integer, noDataValue: Double, replace: Boolean)
+```
+
+```
+RS_SetBandNoDataValue(raster: Raster, bandIndex: Integer = 1, noDataValue: Double)
+```
 
 Since: `v1.5.0`
 

diff --git a/...src/main/scala/org/apache/spark/sql/sedona_sql/expressions/raster/RasterBandEditors.scala b/...src/main/scala/org/apache/spark/sql/sedona_sql/expressions/raster/RasterBandEditors.scala
@@ -24,6 +24,7 @@ import org.apache.spark.sql.sedona_sql.expressions.InferrableFunctionConverter._
 import org.apache.spark.sql.sedona_sql.expressions.InferredExpression
 
 case class RS_SetBandNoDataValue(inputExpressions: Seq[Expression]) extends InferredExpression(
+  inferrableFunction4(RasterBandEditors.setBandNoDataValue),
   inferrableFunction3(RasterBandEditors.setBandNoDataValue),
   inferrableFunction2(RasterBandEditors.setBandNoDataValue)
 ) {