[SEDONA-426] Change raster cloning to be able to include metadata.

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

@@ -133,10 +133,16 @@ 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 = PlanarImage.createColorModel(resultRaster.getSampleModel());
+        ColorModel cm;
+        if (originalColorModel.isCompatibleRaster(resultRaster)) {
+            cm = originalColorModel;
+        }else {
+            cm = PlanarImage.createColorModel(resultRaster.getSampleModel());
+        }
         WritableRenderedImage resultImage = new BufferedImage(cm, resultRaster, false, null);
         try {
             String prevScript = previousScript.get();
@@ -164,10 +170,10 @@ public static GridCoverage2D mapAlgebra(GridCoverage2D gridCoverage2D, String pi
                 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.create(convertedRaster, gridCoverage2D.getGridGeometry(), null, noDataValue);
+                return RasterUtils.clone(convertedRaster,null, gridCoverage2D, noDataValue, false);
             } else {
                 // build a new GridCoverage2D from the resultImage
-                return RasterUtils.create(resultImage, gridCoverage2D.getGridGeometry(), null, noDataValue);
+                return RasterUtils.clone(resultImage, null, gridCoverage2D, noDataValue, false);
             }
         } catch (Exception e) {
             throw new RuntimeException("Failed to run map algebra", e);

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

@@ -19,7 +19,6 @@
 package org.apache.sedona.common.raster;
 
 import org.apache.sedona.common.utils.RasterUtils;
-import org.geotools.coverage.grid.GridCoordinates2D;
 import org.geotools.coverage.grid.GridCoverage2D;
 import org.geotools.geometry.DirectPosition2D;
 import org.locationtech.jts.geom.Coordinate;
@@ -67,7 +66,7 @@ public static GridCoverage2D setValues(GridCoverage2D raster, int band, int colX
         for (int j = rowY; j < rowY + height; j++) {
             for (int i = colX; i < colX + width; i++) {
                 double[] pixel = rasterCopied.getPixel(i, j, (double[]) null);
-                if (keepNoData && noDataValue == pixel[band - 1]) {
+                if (keepNoData && noDataValue != null && noDataValue == pixel[band - 1]) {
                     iterator++;
                     continue;
                 } else {
@@ -77,7 +76,7 @@ public static GridCoverage2D setValues(GridCoverage2D raster, int band, int colX
                 iterator++;
             }
         }
-        return RasterUtils.create(rasterCopied, raster.getGridGeometry(), raster.getSampleDimensions());
+        return RasterUtils.clone(rasterCopied, raster.getSampleDimensions(), raster, null, true); //Keep metadata since this is essentially the same raster
     }
 
     /**
@@ -160,7 +159,7 @@ public static GridCoverage2D setValues(GridCoverage2D raster, int band, Geometry
                     double[] pixel = rasterCopied.getPixel(i, j, (double[]) null);
                     // [0] as only one band in the rasterized Geometry
                     double pixelNew = rasterizedGeomData.getPixel(k, l, (double[]) null)[0];
-                    if (pixelNew == 0 || keepNoData && noDataValue == pixel[band - 1]) {
+                    if (keepNoData && noDataValue != null && noDataValue == pixel[band - 1]) {
                         continue;
                     } else {
                         pixel[band - 1] = pixelNew;
@@ -169,8 +168,7 @@ public static GridCoverage2D setValues(GridCoverage2D raster, int band, Geometry
                 }
             }
         }
-
-        return RasterUtils.create(rasterCopied, raster.getGridGeometry(), raster.getSampleDimensions());
+        return RasterUtils.clone(rasterCopied, raster.getSampleDimensions(), raster, null, true); // keep metadata since this is essentially the same raster
     }
 
     /**

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

@@ -20,7 +20,6 @@
 
 import org.apache.sedona.common.utils.RasterUtils;
 import org.geotools.coverage.GridSampleDimension;
-import org.geotools.coverage.grid.GridCoordinates2D;
 import org.geotools.coverage.grid.GridCoverage2D;
 import org.opengis.referencing.FactoryException;
 
@@ -184,8 +183,7 @@ public static GridCoverage2D getBand(GridCoverage2D rasterGeom, int[] bandIndexe
                 sampleDimensionsResult[i] = RasterUtils.createSampleDimensionWithNoDataValue(sampleDimension, noDataValue);
             }
         }
-
-        return RasterUtils.create(wr, resultRaster.getGridGeometry(), sampleDimensionsResult);
+        return RasterUtils.clone(wr, sampleDimensionsResult, rasterGeom, null, false); //do not keep meta-data since this will most probably be a new raster
     }
 
     public static String getBandType(GridCoverage2D raster, int band) {

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

@@ -58,7 +58,9 @@ public static GridCoverage2D setBandNoDataValue(GridCoverage2D raster, int bandI
             }
             GridSampleDimension[] sampleDimensions = raster.getSampleDimensions();
             sampleDimensions [bandIndex - 1] = RasterUtils.removeNoDataValue(sampleDimensions[bandIndex - 1]);
-            return RasterUtils.create(raster.getRenderedImage(), raster.getGridGeometry(), sampleDimensions, null);
+            //return RasterUtils.create(raster.getRenderedImage(), raster)
+            return RasterUtils.clone(raster.getRenderedImage(), null, sampleDimensions, raster, null, true);
+            //return RasterUtils.create(raster.getRenderedImage(), raster.getGridGeometry(), sampleDimensions, null);
         }
 
         if ( !(rasterNoData == null) && rasterNoData.equals(noDataValue)) {
@@ -74,8 +76,8 @@ public static GridCoverage2D setBandNoDataValue(GridCoverage2D raster, int bandI
                 PixelOrientation.UPPER_LEFT,
                 affine, raster.getCoordinateReferenceSystem2D(), null
         );
-
-        return RasterUtils.create(raster.getRenderedImage(), gridGeometry2D, bands, null);
+        return RasterUtils.clone(raster.getRenderedImage(), null, bands, raster, null, true);
+        //return RasterUtils.create(raster.getRenderedImage(), gridGeometry2D, bands, null);
     }
 
     /**
@@ -247,10 +249,12 @@ public static GridCoverage2D clip(GridCoverage2D raster, int band, Geometry geom
                     }
                 }
             }
-            newRaster = RasterUtils.create(resultRaster, raster.getGridGeometry(), newRaster.getSampleDimensions(), noDataValue);
+            newRaster = RasterUtils.clone(resultRaster, raster.getGridGeometry(), newRaster.getSampleDimensions(), newRaster, noDataValue, true);
+           // newRaster = RasterUtils.create(resultRaster, raster.getGridGeometry(), newRaster.getSampleDimensions(), noDataValue);
         } else {
             // to add no-data value
-            newRaster = RasterUtils.create(newRaster.getRenderedImage(), newRaster.getGridGeometry(), newRaster.getSampleDimensions(), noDataValue);
+            newRaster = RasterUtils.clone(newRaster.getRenderedImage(), newRaster.getGridGeometry(), newRaster.getSampleDimensions(), newRaster, noDataValue, true);
+            //newRaster = RasterUtils.create(newRaster.getRenderedImage(), newRaster.getGridGeometry(), newRaster.getSampleDimensions(), noDataValue);
         }
 
         return newRaster;

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

@@ -41,7 +41,6 @@
 import javax.media.jai.RasterFactory;
 import java.awt.image.WritableRaster;
 import java.io.IOException;
-import java.util.Arrays;
 
 public class RasterConstructors
 {
@@ -119,7 +118,7 @@ public static GridCoverage2D asRaster(Geometry geom, GridCoverage2D raster, Stri
         double [] samples = RasterUtils.getRaster(rasterized.getRenderedImage()).getSamples(0, 0, width, height, 0, (double[]) null);
         writableRaster.setSamples(0, 0, width, height, 0, samples);
 
-        return RasterUtils.create(writableRaster, rasterized.getGridGeometry(), rasterized.getSampleDimensions(), noDataValue);
+        return RasterUtils.clone(writableRaster, rasterized.getSampleDimensions(), rasterized, noDataValue, false); //no need to original raster metadata since this is a new raster.
     }
 
     /**

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

@@ -92,9 +92,7 @@ public static GridCoverage2D setGeoReference(GridCoverage2D raster, String geoRe
                 PixelOrientation.UPPER_LEFT,
                 affine, raster.getCoordinateReferenceSystem(), null
         );
-
-        return RasterUtils.create(raster.getRenderedImage(), gridGeometry2D, raster.getSampleDimensions(),
-                null);
+        return RasterUtils.clone(raster.getRenderedImage(), gridGeometry2D, raster.getSampleDimensions(), raster, null, true);
     }
 
     public static GridCoverage2D setGeoReference(GridCoverage2D raster, String geoRefCoords) {

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

@@ -57,9 +57,7 @@
 import java.awt.image.Raster;
 import java.awt.image.RenderedImage;
 import java.awt.image.WritableRaster;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.List;
+import java.util.*;
 
 /**
  * Utility functions for working with GridCoverage2D objects.
@@ -81,12 +79,100 @@ public static GridCoverage2D create(WritableRaster raster, GridGeometry2D gridGe
     }
 
     /**
-     * Create a new empty raster from the given WritableRaster object.
+     * Create a new raster, cloning all metadata from the passed reference raster.
      * @param raster The raster object to be wrapped as an image.
-     * @param gridGeometry The grid geometry of the raster.
-     * @param bands The bands of the raster.
-     * @return A new GridCoverage2D object.
+     * @param referenceRaster The reference raster to clone, all metadata is cloned
+     * @param noDataValue new noDataValue
+     * @return cloned raster
      */
+    public static GridCoverage2D clone(WritableRaster raster, GridGeometry2D gridGeometry2D, GridSampleDimension[] bands, GridCoverage2D referenceRaster, Double noDataValue, boolean keepMetadata) {
+        Map propertyMap = null;
+        if (keepMetadata) {
+            propertyMap = referenceRaster.getProperties();
+        }
+        ColorModel originalColorModel = referenceRaster.getRenderedImage().getColorModel();
+        if (Objects.isNull(gridGeometry2D)) {
+            gridGeometry2D = referenceRaster.getGridGeometry();
+        }
+        int numBand = raster.getNumBands();
+        int rasterDataType = raster.getDataBuffer().getDataType();
+        ColorModel colorModel;
+        if (originalColorModel.isCompatibleRaster(raster)) {
+            colorModel = originalColorModel;
+        }else {
+            final ColorSpace cs = new BogusColorSpace(numBand);
+            final int[] nBits = new int[numBand];
+            Arrays.fill(nBits, DataBuffer.getDataTypeSize(rasterDataType));
+            colorModel = new ComponentColorModel(cs, nBits, false, true, Transparency.OPAQUE, rasterDataType);
+        }
+        if (noDataValue != null) {
+            GridSampleDimension[] newBands = new GridSampleDimension[numBand];
+            for (int k = 0; k < numBand; k++) {
+                if (bands != null) {
+                    newBands[k] = createSampleDimensionWithNoDataValue(bands[k], noDataValue);
+                } else {
+                    newBands[k] = createSampleDimensionWithNoDataValue("band_" + k, noDataValue);
+                }
+            }
+            bands = newBands;
+        }
+
+        GridCoverage2D[] referenceRasterSources = keepMetadata ? referenceRaster.getSources().toArray(new GridCoverage2D[0]) : null;
+        CharSequence rasterName = keepMetadata ? referenceRaster.getName() : "genericCoverage";
+
+        final RenderedImage image = new BufferedImage(colorModel, raster, false, null);
+        return gridCoverageFactory.create(rasterName, image, gridGeometry2D, bands, referenceRasterSources, propertyMap);
+
+    }
+
+    public static GridCoverage2D clone(WritableRaster raster, GridSampleDimension[] bands, GridCoverage2D referenceRaster, Double noDataValue, boolean keepMetadata) {
+        return RasterUtils.clone(raster, null, bands, referenceRaster, noDataValue, keepMetadata);
+    }
+
+    /**
+     * Create a new raster, cloning all metadata from the passed reference raster.
+     * @param image The RenderedImage to be used as data for the created GridCoverage2D
+     * @param referenceRaster The reference raster to clone, all metadata is cloned
+     * @param noDataValue new noDataValue
+     * @return cloned raster
+     */
+
+
+    public static GridCoverage2D clone(RenderedImage image, GridSampleDimension[] bands, GridCoverage2D referenceRaster, Double noDataValue, boolean keepMetadata) {
+        return RasterUtils.clone(image, null, bands, referenceRaster, noDataValue, keepMetadata);
+    }
+
+    public static GridCoverage2D clone(RenderedImage image, GridGeometry2D gridGeometry2D, GridSampleDimension[] bands, GridCoverage2D referenceRaster, Double noDataValue, boolean keepMetadata) {
+        int numBand = image.getSampleModel().getNumBands();
+        if (Objects.isNull(gridGeometry2D)) {
+            gridGeometry2D = referenceRaster.getGridGeometry();
+        }
+        if (noDataValue != null) {
+            GridSampleDimension[] newBands = new GridSampleDimension[numBand];
+            for (int k = 0; k < numBand; k++) {
+                if (bands != null) {
+                    newBands[k] = createSampleDimensionWithNoDataValue(bands[k], noDataValue);
+                } else {
+                    newBands[k] = createSampleDimensionWithNoDataValue("band_" + k, noDataValue);
+                }
+            }
+            bands = newBands;
+        }
+        GridCoverage2D[] referenceRasterSources = keepMetadata ? referenceRaster.getSources().toArray(new GridCoverage2D[0]) : null;
+        Map propertyMap = null;
+        if (keepMetadata) {
+            propertyMap = referenceRaster.getProperties();
+        }        CharSequence rasterName = keepMetadata ? referenceRaster.getName() : "genericCoverage";
+        return gridCoverageFactory.create(rasterName, image, gridGeometry2D, bands, referenceRasterSources, propertyMap);
+    }
+
+        /**
+         * Create a new empty raster from the given WritableRaster object.
+         * @param raster The raster object to be wrapped as an image.
+         * @param gridGeometry The grid geometry of the raster.
+         * @param bands The bands of the raster.
+         * @return A new GridCoverage2D object.
+         */
     public static GridCoverage2D create(WritableRaster raster, GridGeometry2D gridGeometry, GridSampleDimension[] bands, Double noDataValue) {
         int numBand = raster.getNumBands();
         int rasterDataType = raster.getDataBuffer().getDataType();
@@ -421,7 +507,8 @@ public static GridCoverage2D copyRasterAndAppendBand(GridCoverage2D gridCoverage
             sampleDimensions[numBand - 1] = new GridSampleDimension("band" + numBand);
         }
         // Construct a GridCoverage2D with the copied image.
-        return create(wr, gridCoverage2D.getGridGeometry(), sampleDimensions);
+        return clone(wr, gridCoverage2D.getGridGeometry(), sampleDimensions, gridCoverage2D, null, true);
+        //return create(wr, gridCoverage2D.getGridGeometry(), sampleDimensions);
     }
 
     public static GridCoverage2D copyRasterAndAppendBand(GridCoverage2D gridCoverage2D, Number[] bandValues) {
@@ -456,7 +543,8 @@ public static GridCoverage2D copyRasterAndReplaceBand(GridCoverage2D gridCoverag
         } else if (noDataValue != null) {
             sampleDimensions[bandIndex - 1] = createSampleDimensionWithNoDataValue(sampleDimension, noDataValue);
         }
-        return create(wr, gridCoverage2D.getGridGeometry(), sampleDimensions);
+        return clone(wr, gridCoverage2D.getGridGeometry(), sampleDimensions, gridCoverage2D, null, true);
+        //return create(wr, gridCoverage2D.getGridGeometry(), sampleDimensions);
     }
 
     public static GridCoverage2D copyRasterAndReplaceBand(GridCoverage2D gridCoverage2D, int bandIndex, Number[] bandValues) {

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

@@ -1074,7 +1074,8 @@ class rasteralgebraTest extends TestBaseScala with BeforeAndAfter with GivenWhen
           }
           assert(rast2.getRenderedImage.getSampleModel.getDataType == dataType)
           val noDataValue = RasterUtils.getNoDataValue(rast2.getSampleDimension(0))
-          assert(noDataValue.isNaN)
+          val originalNoDataValue = RasterUtils.getNoDataValue(rast.getSampleDimension(0))
+          assert(noDataValue.isNaN())
           val band = MapAlgebra.bandAsArray(rast, 1)
           val band2 = MapAlgebra.bandAsArray(rast2, 1)
           assert(band.size == band2.size)