[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.create(convertedRaster, gridCoverage2D, noDataValue, false);
             } else {
                 // build a new GridCoverage2D from the resultImage
-                return RasterUtils.create(resultImage, gridCoverage2D.getGridGeometry(), null, noDataValue);
+                return RasterUtils.create(resultImage, gridCoverage2D, noDataValue, false);
             }
         } catch (Exception e) {
             throw new RuntimeException("Failed to run map algebra", e);

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

@@ -21,6 +21,7 @@
 import com.sun.media.imageioimpl.common.BogusColorSpace;
 import org.apache.sedona.common.FunctionsGeoTools;
 import org.apache.sedona.common.raster.RasterAccessors;
+import org.apache.sedona.common.raster.RasterBandAccessors;
 import org.geotools.coverage.Category;
 import org.geotools.coverage.CoverageFactoryFinder;
 import org.geotools.coverage.GridSampleDimension;
@@ -57,9 +58,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.
@@ -80,6 +79,89 @@ public static GridCoverage2D create(WritableRaster raster, GridGeometry2D gridGe
         return create(raster, gridGeometry, bands, null);
     }
 
+    /**
+     * Create a new raster, cloning all metadata from the passed reference raster.
+     * @param raster The raster object to be wrapped as an image.
+     * @param referenceRaster The reference raster to clone, all metadata is cloned
+     * @param noDataValue Optionally, a new noDataValue, if passed null, noDataVal of the reference raster is used.
+     * @return cloned raster
+     */
+    public static GridCoverage2D create(WritableRaster raster, GridCoverage2D referenceRaster, Double noDataValue, boolean keepMetadata) {
+        GridSampleDimension[] bands = keepMetadata ? referenceRaster.getSampleDimensions() : null;
+        if (Objects.isNull(noDataValue)) {
+            noDataValue = RasterBandAccessors.getBandNoDataValue(referenceRaster, 1); //resort to using noDataValue of the original raster.
+        }
+        Map propertyMap = null;
+        if (keepMetadata) {
+            propertyMap = referenceRaster.getProperties();
+        }
+        ColorModel originalColorModel = referenceRaster.getRenderedImage().getColorModel();
+        GridGeometry2D gridGeometry = referenceRaster.getGridGeometry();
+        int numBand = raster.getNumBands();
+        int rasterDataType = raster.getDataBuffer().getDataType();
+
+        final ColorSpace cs = new BogusColorSpace(numBand);
+        final int[] nBits = new int[numBand];
+        Arrays.fill(nBits, DataBuffer.getDataTypeSize(rasterDataType));
+        ColorModel colorModel;
+        if (originalColorModel.isCompatibleRaster(raster)) {
+            colorModel = originalColorModel;
+        }else {
+            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, gridGeometry, bands, referenceRasterSources, propertyMap);
+
+    }
+
+    /**
+     * 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 Optionally, a new noDataValue, if passed null, noDataVal of the reference raster is used.
+     * @return cloned raster
+     */
+    public static GridCoverage2D create(RenderedImage image, GridCoverage2D referenceRaster, Double noDataValue, boolean keepMetadata) {
+        int numBand = image.getSampleModel().getNumBands();
+        if (Objects.isNull(noDataValue)) {
+            noDataValue = RasterBandAccessors.getBandNoDataValue(referenceRaster, 1); //resort to using noDataValue of the original raster.
+        }
+        GridSampleDimension[] bands = keepMetadata ? referenceRaster.getSampleDimensions() : null;
+        GridGeometry2D gridGeometry = 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, gridGeometry, bands, referenceRasterSources, propertyMap);
+    }
+
     /**
      * Create a new empty raster from the given WritableRaster object.
      * @param raster The raster object to be wrapped as an image.

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

@@ -1040,7 +1040,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.equals(originalNoDataValue))
           val band = MapAlgebra.bandAsArray(rast, 1)
           val band2 = MapAlgebra.bandAsArray(rast2, 1)
           assert(band.size == band2.size)