[sedona] branch master updated: [SEDONA-319] Make result of RS_BandAsArray serializable (#899)

[jiayu]

This is an automated email from the ASF dual-hosted git repository.

jiayu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/sedona.git


The following commit(s) were added to refs/heads/master by this push:
     new 3a1e8c6c [SEDONA-319] Make result of RS_BandAsArray serializable (#899)
3a1e8c6c is described below

commit 3a1e8c6cb725d680cfb5c6e54b612e1e5b245645
Author: Kristin Cowalcijk <b...@wherobots.com>
AuthorDate: Thu Jul 13 01:40:13 2023 +0800

    [SEDONA-319] Make result of RS_BandAsArray serializable (#899)
---
 .../apache/sedona/common/raster/MapAlgebra.java    | 35 +++++++++++++++++++---
 .../org/apache/sedona/sql/rasteralgebraTest.scala  | 11 +++++++
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git 
a/common/src/main/java/org/apache/sedona/common/raster/MapAlgebra.java 
b/common/src/main/java/org/apache/sedona/common/raster/MapAlgebra.java
index 0a24cce2..67e8ad18 100644
--- a/common/src/main/java/org/apache/sedona/common/raster/MapAlgebra.java
+++ b/common/src/main/java/org/apache/sedona/common/raster/MapAlgebra.java
@@ -18,16 +18,25 @@
  */
 package org.apache.sedona.common.raster;
 
+import com.sun.media.imageioimpl.common.BogusColorSpace;
 import org.geotools.coverage.CoverageFactoryFinder;
+import org.geotools.coverage.GridSampleDimension;
 import org.geotools.coverage.grid.GridCoverage2D;
 import org.geotools.coverage.grid.GridCoverageFactory;
 
 import javax.media.jai.RasterFactory;
 
 import java.awt.Point;
+import java.awt.Transparency;
+import java.awt.color.ColorSpace;
+import java.awt.image.BufferedImage;
+import java.awt.image.ColorModel;
+import java.awt.image.ComponentColorModel;
+import java.awt.image.DataBuffer;
 import java.awt.image.Raster;
 import java.awt.image.RenderedImage;
 import java.awt.image.WritableRaster;
+import java.util.Arrays;
 
 public class MapAlgebra
 {
@@ -118,9 +127,14 @@ public class MapAlgebra
                 wr.setPixel(i, j, copiedPixels);
             }
         }
-        // Create a new GridCoverage2D with the copied image
-        GridCoverageFactory gridCoverageFactory = 
CoverageFactoryFinder.getGridCoverageFactory(null);
-        return gridCoverageFactory.create(gridCoverage2D.getName(), wr, 
gridCoverage2D.getEnvelope());
+        // Add a sample dimension for newly added band
+        int numBand = wr.getNumBands();
+        GridSampleDimension[] originalSampleDimensions = 
gridCoverage2D.getSampleDimensions();
+        GridSampleDimension[] sampleDimensions = new 
GridSampleDimension[numBand];
+        System.arraycopy(originalSampleDimensions, 0, sampleDimensions, 0, 
originalSampleDimensions.length);
+        sampleDimensions[numBand - 1] = new GridSampleDimension("band" + 
numBand);
+        // Construct a GridCoverage2D with the copied image.
+        return createCompatibleGridCoverage2D(gridCoverage2D, wr, 
sampleDimensions);
     }
 
     private static GridCoverage2D copyRasterAndReplaceBand(GridCoverage2D 
gridCoverage2D, int bandIndex, double[] bandValues) {
@@ -141,7 +155,20 @@ public class MapAlgebra
             }
         }
         // Create a new GridCoverage2D with the copied image
+        return createCompatibleGridCoverage2D(gridCoverage2D, wr, 
gridCoverage2D.getSampleDimensions());
+    }
+
+    private static GridCoverage2D 
createCompatibleGridCoverage2D(GridCoverage2D gridCoverage2D, WritableRaster 
wr, GridSampleDimension[] bands) {
+        int rasterDataType = wr.getDataBuffer().getDataType();
+        int numBand = wr.getNumBands();
+        final ColorSpace cs = new BogusColorSpace(numBand);
+        final int[] nBits = new int[numBand];
+        Arrays.fill(nBits, DataBuffer.getDataTypeSize(rasterDataType));
+        ColorModel colorModel =
+                new ComponentColorModel(cs, nBits, false, true, 
Transparency.OPAQUE, rasterDataType);
+        final RenderedImage image = new BufferedImage(colorModel, wr, false, 
null);
         GridCoverageFactory gridCoverageFactory = 
CoverageFactoryFinder.getGridCoverageFactory(null);
-        return gridCoverageFactory.create(gridCoverage2D.getName(), wr, 
gridCoverage2D.getEnvelope());
+        return gridCoverageFactory.create(gridCoverage2D.getName(), image, 
gridCoverage2D.getCoordinateReferenceSystem(),
+                gridCoverage2D.getGridGeometry().getGridToCRS(), bands, null, 
null);
     }
 }
diff --git 
a/sql/common/src/test/scala/org/apache/sedona/sql/rasteralgebraTest.scala 
b/sql/common/src/test/scala/org/apache/sedona/sql/rasteralgebraTest.scala
index 8ef24399..7728b772 100644
--- a/sql/common/src/test/scala/org/apache/sedona/sql/rasteralgebraTest.scala
+++ b/sql/common/src/test/scala/org/apache/sedona/sql/rasteralgebraTest.scala
@@ -452,5 +452,16 @@ class rasteralgebraTest extends TestBaseScala with 
BeforeAndAfter with GivenWhen
       assert(df.selectExpr("RS_Intersects(raster, 
ST_SetSRID(ST_Point(33.81798,-117.47993), 4326))").first().getBoolean(0))
       assert(!df.selectExpr("RS_Intersects(raster, 
ST_SetSRID(ST_Point(33.97896,-117.27868), 4326))").first().getBoolean(0))
     }
+
+    it("Passed RS_AddBandFromArray collect generated raster") {
+      var df = sparkSession.read.format("binaryFile").load(resourceFolder + 
"raster/test1.tiff")
+      df = df.selectExpr("RS_FromGeoTiff(content) as raster", 
"RS_BandAsArray(RS_FromGeoTiff(content), 1) as band")
+      df = df.selectExpr("RS_AddBandFromArray(raster, band, 1) as raster", 
"band")
+      var raster = df.collect().head.getAs[GridCoverage2D](0)
+      assert(raster.getNumSampleDimensions == 1)
+      df = df.selectExpr("RS_AddBandFromArray(raster, band, 2) as raster", 
"band")
+      raster = df.collect().head.getAs[GridCoverage2D](0)
+      assert(raster.getNumSampleDimensions == 2)
+    }
   }
 }