[SEDONA-436] Fix RS_SetValues bug (#1132)

apache · Nov 29, 2023 · 8aee765 · 8aee765
1 parent 1642af9
commit 8aee765
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 17 deletions.
diff --git a/common/src/main/java/org/apache/sedona/common/raster/PixelFunctionEditors.java b/common/src/main/java/org/apache/sedona/common/raster/PixelFunctionEditors.java
@@ -122,7 +122,7 @@ public static GridCoverage2D setValues(GridCoverage2D raster, int band, Geometry
 
         Raster rasterizedGeomData = RasterUtils.getRaster(rasterizedGeom.getRenderedImage());
         double colX = RasterAccessors.getUpperLeftX(rasterizedGeom), rowY = RasterAccessors.getUpperLeftY(rasterizedGeom);
-        int height = RasterAccessors.getHeight(rasterizedGeom), width = RasterAccessors.getWidth(rasterizedGeom);
+        int heightGeometryRaster = RasterAccessors.getHeight(rasterizedGeom), widthGeometryRaster = RasterAccessors.getWidth(rasterizedGeom);
         int heightOriginalRaster = RasterAccessors.getHeight(raster), widthOriginalRaster = RasterAccessors.getWidth(raster);
         WritableRaster rasterCopied = makeCopiedRaster(raster);
 
@@ -144,27 +144,28 @@ public static GridCoverage2D setValues(GridCoverage2D raster, int band, Geometry
             int[] pixelLocation = RasterUtils.getGridCoordinatesFromWorld(raster, colX, rowY);
             int x = pixelLocation[0], y = pixelLocation[1];
 
-            // lower-bound if the rasterized geometry starts at more north or west from the given raster than move rasterized geometry starting pixel accordingly.
-            if (x < 0) {
-                x = Math.abs(x);
-            }
-            if (y < 0) {
-                y = Math.abs(y);
-            }
-            // i & j is for main raster
-            // k & l is for rasterized geom
-            // added an upperbound if the rasterized geometry is bigger than provided raster
-            for (int j = 0, l = y; j < heightOriginalRaster && l < height; j++, l++) {
-                for (int i = 0, k = x; i < widthOriginalRaster && k < width; i++, k++) {
-                    double[] pixel = rasterCopied.getPixel(i, j, (double[]) null);
+            // rasterX & rasterY are the starting pixels for the target raster
+            int rasterX = Math.max(x, 0);
+            int rasterY = Math.max(y, 0);
+            // geometryX & geometryY are the starting pixels for the geometry raster
+            int geometryX = rasterX - x;
+            int geometryY = rasterY - y;
+            // widthRegion & heightRegion are the size of the region to update
+            int widthRegion = Math.min(widthGeometryRaster - geometryX, widthOriginalRaster - rasterX);
+            int heightRegion = Math.min(heightGeometryRaster - geometryY, heightOriginalRaster - rasterY);
+
+            for (int j = 0; j < heightRegion; j++) {
+                for (int i = 0; i < widthRegion; i++) {
+                    double[] pixel = rasterCopied.getPixel(rasterX + i, rasterY + j, (double[]) null);
                     // [0] as only one band in the rasterized Geometry
-                    double pixelNew = rasterizedGeomData.getPixel(k, l, (double[]) null)[0];
-                    if (keepNoData && noDataValue != null && noDataValue == pixel[band - 1]) {
+                    double pixelNew = rasterizedGeomData.getPixel(geometryX + i, geometryY + j, (double[]) null)[0];
+                    // skipping 0 from the rasterized geometry as
+                    if (pixelNew == 0 || keepNoData && noDataValue != null && noDataValue == pixel[band - 1]) {
                         continue;
                     } else {
                         pixel[band - 1] = pixelNew;
                     }
-                    rasterCopied.setPixel(i, j, pixel);
+                    rasterCopied.setPixel(rasterX + i, rasterY + j, pixel);
                 }
             }
         }

diff --git a/common/src/test/java/org/apache/sedona/common/raster/FunctionEditorsTest.java b/common/src/test/java/org/apache/sedona/common/raster/FunctionEditorsTest.java
@@ -26,6 +26,7 @@
 import org.opengis.referencing.FactoryException;
 import org.opengis.referencing.operation.TransformException;
 
+import java.io.IOException;
 import java.util.Arrays;
 
 import static org.junit.Assert.*;
@@ -49,6 +50,28 @@ public void testSetValuesWithEmptyRaster() throws FactoryException {
         assertArrayEquals(actual, expected, 0.0);
     }
 
+    @Test
+    public void testSetValuesWithGeomInRaster() throws IOException, ParseException, FactoryException, TransformException {
+        GridCoverage2D raster = rasterFromGeoTiff(resourceFolder + "raster_geotiff_color/FAA_UTM18N_NAD83.tif");
+        String polygon = "POLYGON ((236722 4204770, 243900 4204770, 243900 4197590, 236722 4197590, 236722 4204770))";
+        Geometry geom = Constructors.geomFromWKT(polygon, 26918);
+
+        GridCoverage2D result = PixelFunctionEditors.setValues(raster, 1, geom, 10, false);
+
+        Geometry point = Constructors.geomFromWKT("POINT (243700 4197797)", 26918);
+        double actual = PixelFunctions.value(result, point, 1);
+        double expected = 10.0;
+        assertEquals(expected, actual, 0d);
+
+        point = Constructors.geomFromWKT("POINT (240311 4202806)", 26918);
+        actual = PixelFunctions.value(result, point, 1);
+        assertEquals(expected, actual, 0d);
+
+        point = Constructors.geomFromWKT("POINT (241800 4199660)", 26918);
+        actual = PixelFunctions.value(result, point, 1);
+        assertEquals(expected, actual, 0d);
+    }
+
     @Test
     public void testSetValuesGeomVariant() throws FactoryException, ParseException, TransformException {
         GridCoverage2D emptyRaster = RasterConstructors.makeEmptyRaster(1, 4, 6, 1, -1, 1, -1, 0, 0, 0);