[SEDONA-394] fix RS_Band data type bug

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

@@ -24,6 +24,7 @@
 import org.geotools.coverage.grid.GridCoverage2D;
 import org.opengis.referencing.FactoryException;
 
+import javax.media.jai.RasterFactory;
 import java.awt.image.Raster;
 import java.awt.image.WritableRaster;
 import java.util.Arrays;
@@ -138,29 +139,48 @@ public static double[] getSummaryStats(GridCoverage2D raster) {
 //        return getSummaryStats(raster, 1, excludeNoDataValue);
 //    }
 
+    /**
+     * @param rasterGeom The raster where the bands will be extracted from.
+     * @param bandIndexes The bands to be added to new raster.
+     *
+     * @return Raster with the specified bands.
+     **/
     public static GridCoverage2D getBand(GridCoverage2D rasterGeom, int[] bandIndexes) throws FactoryException {
         Double noDataValue;
         double[] metadata = RasterAccessors.metadata(rasterGeom);
         int width = (int) metadata[2], height = (int) metadata[3];
         GridCoverage2D resultRaster = RasterConstructors.makeEmptyRaster(bandIndexes.length, width, height,
                 metadata[0], metadata[1], metadata[4], metadata[5], metadata[6], metadata[7], (int) metadata[8]);
 
+        // Get raster data type
+        Raster raster = RasterUtils.getRaster(rasterGeom.getRenderedImage());
+        int dataTypeCode = raster.getDataBuffer().getDataType();
+        boolean isDataTypeIntegral = RasterUtils.isDataTypeIntegral(dataTypeCode);
+
         // Get band data that's required
         int[] bandsDistinct = Arrays.stream(bandIndexes).distinct().toArray();
-        HashMap<Integer, double[]> bandData = new HashMap<>();
+        HashMap<Integer, Object> bandData = new HashMap<>();
         for (int curBand: bandsDistinct) {
             RasterUtils.ensureBand(rasterGeom, curBand);
-            bandData.put(curBand - 1, MapAlgebra.bandAsArray(rasterGeom, curBand));
+            if (isDataTypeIntegral) {
+                bandData.put(curBand - 1, raster.getSamples(0, 0, width, height, curBand - 1, (int[]) null));
+            } else {
+                bandData.put(curBand - 1, raster.getSamples(0, 0, width, height, curBand - 1, (double[]) null));
+            }
         }
 
-        // Get Writable Raster from the resultRaster
-        WritableRaster wr = RasterUtils.getRaster(resultRaster.getRenderedImage()).createCompatibleWritableRaster();
+        // Create Writable Raster with the datatype of given raster
+        WritableRaster wr = RasterFactory.createBandedRaster(dataTypeCode, width, height, bandIndexes.length, null);
 
         GridSampleDimension[] sampleDimensionsOg = rasterGeom.getSampleDimensions();
         GridSampleDimension[] sampleDimensionsResult = resultRaster.getSampleDimensions();
         for (int i = 0; i < bandIndexes.length; i ++) {
             sampleDimensionsResult[i] = sampleDimensionsOg[bandIndexes[i] - 1];
-            wr.setSamples(0, 0, width, height, i, bandData.get(bandIndexes[i] - 1));
+            if (isDataTypeIntegral) {
+                wr.setSamples(0, 0, width, height, i, (int[]) bandData.get(bandIndexes[i] - 1));
+            } else {
+                wr.setSamples(0, 0, width, height, i, (double[]) bandData.get(bandIndexes[i] - 1));
+            }
             noDataValue = RasterBandAccessors.getBandNoDataValue(rasterGeom, bandIndexes[i]);
             GridSampleDimension sampleDimension = sampleDimensionsResult[i];
             if (noDataValue != null) {

common/src/test/java/org/apache/sedona/common/raster/RasterBandAccessorsTest.java

@@ -266,6 +266,24 @@ public void testGetBand() throws FactoryException {
         assertArrayEquals(expectedBandValues, actualBandValues, 0.1d);
     }
 
+    @Test
+    public void testGetBandWithDataTypes() throws FactoryException, IOException {
+        GridCoverage2D emptyRaster = RasterConstructors.makeEmptyRaster( 4, "d", 5, 5, 3, -215, 2, -2, 2, 2, 0);
+        double[] values1 = new double[] {16, 0, 24, 33, 43, 49, 64, 0, 76, 77, 79, 89, 0, 116, 118, 125, 135, 0, 157, 190, 215, 229, 241, 248, 249};
+        emptyRaster = MapAlgebra.addBandFromArray(emptyRaster, values1, 1, 0d);
+        String actual = RasterBandAccessors.getBandType(emptyRaster, 1);
+        String expected = "REAL_64BITS";
+        assertEquals(expected, actual);
+
+        GridCoverage2D raster = rasterFromGeoTiff(resourceFolder + "raster_geotiff_color/FAA_UTM18N_NAD83.tif");
+        raster = RasterBandAccessors.getBand(raster, new int[] {2,1,3});
+        for (int i = 1; i <= RasterAccessors.numBands(raster); i++) {
+            actual = RasterBandAccessors.getBandType(raster, i);
+            expected = "UNSIGNED_8BITS";
+            assertEquals(expected, actual);
+        }
+    }
+
     @Test
     public void testGetBandWithRaster() throws IOException, FactoryException {
         GridCoverage2D raster = rasterFromGeoTiff(resourceFolder + "raster_geotiff_color/FAA_UTM18N_NAD83.tif");

common/src/test/java/org/apache/sedona/common/raster/RasterOutputTest.java

@@ -52,6 +52,18 @@ public void testAsPNG() throws IOException, FactoryException {
         assertEquals("image/png", mimeType);
     }
 
+    @Test
+    public void testAsPNGWithBand() throws IOException, FactoryException {
+        String dirPath = System.getProperty("user.dir") + "/target/testAsPNGFunction/";
+        new File(dirPath).mkdirs();
+        GridCoverage2D raster = rasterFromGeoTiff(resourceFolder + "raster_geotiff_color/FAA_UTM18N_NAD83.tif");
+        byte[] pngData = RasterOutputs.asPNG(RasterBandAccessors.getBand(raster, new int[] {3, 1, 2}));
+        RasterOutputs.writeToDiskFile(pngData, dirPath + "test2.png");
+        File f = new File(dirPath + "test2.png");
+        String mimeType = URLConnection.guessContentTypeFromName(f.getName());
+        assertEquals("image/png", mimeType);
+    }
+
     @Test
     public void testAsGeoTiff() throws IOException {
         GridCoverage2D rasterOg = rasterFromGeoTiff(resourceFolder + "raster/test1.tiff");

docs/api/sql/Raster-writer.md

@@ -95,7 +95,7 @@ root
 
 #### RS_AsPNG
 
-Introduction: Returns a PNG byte array, that can be written to raster files as PNGs using the [sedona function](#write-a-binary-dataframe-to-raster-files). This function can only accept pixel data type of unsigned integer.
+Introduction: Returns a PNG byte array, that can be written to raster files as PNGs using the [sedona function](#write-a-binary-dataframe-to-raster-files). This function can only accept pixel data type of unsigned integer. PNG can accept 1 or 3 bands of data from the raster, refer to [RS_Band](../Raster-operators/#rs_band) for more details.
 
 !!!Note
 	Raster having `UNSIGNED_8BITS` pixel data type will have range of `0 - 255`, whereas rasters having `UNSIGNED_16BITS` pixel data type will have range of `0 - 65535`. If provided pixel value is greater than either `255` for `UNSIGNED_8BITS` or `65535` for `UNSIGNED_16BITS`, then the extra bit will be truncated.
@@ -119,6 +119,18 @@ Output:
 [-119, 80, 78, 71, 13, 10, 26, 10, 0, 0, 0, 13, 73...]
 ```
 
+Spark SQL Example:
+
+```sql
+SELECT RS_AsPNG(RS_Band(raster, Array(3, 1, 2)))
+```
+
+Output:
+
+```
+[-103, 78, 94, -26, 61, -16, -91, -103, -65, -116...]
+```
+
 ### Write a binary DataFrame to raster files
 
 Introduction: You can write a Sedona binary DataFrame to external storage using Sedona's built-in `raster` data source. Note that: `raster` data source does not support reading rasters. Please use Spark built-in `binaryFile` and Sedona RS constructors together to read rasters.