easier diff

databrickslabs · Nov 27, 2024 · f0ff785 · f0ff785
1 parent 06e96fc
commit f0ff785
Show file tree

Hide file tree

Showing 11 changed files with 25 additions and 20 deletions.
diff --git a/src/main/scala/com/databricks/labs/mosaic/core/raster/gdal/MosaicRasterGDAL.scala b/src/main/scala/com/databricks/labs/mosaic/core/raster/gdal/MosaicRasterGDAL.scala
@@ -180,7 +180,7 @@ case class MosaicRasterGDAL(
             } else {
                 subdatasets.values
                     .filter(_.toLowerCase(Locale.ROOT).startsWith(getDriversShortName.toLowerCase(Locale.ROOT)))
-                    .flatMap(bp => readRaster(createInfo + ("path" -> bp), None).getBands)
+                    .flatMap(bp => readRaster(createInfo + ("path" -> bp)).getBands)
                     .takeWhile(_.isEmpty)
                     .nonEmpty
             }
@@ -761,7 +761,7 @@ object MosaicRasterGDAL extends RasterReader {
       *   A [[MosaicRasterGDAL]] object.
       */
     override def readBand(bandIndex: Int, createInfo: Map[String, String]): MosaicRasterBandGDAL = {
-        val raster = readRaster(createInfo, None)
+        val raster = readRaster(createInfo)
         // Note: Raster and Band are coupled, this can cause a pointer leak
         raster.getBand(bandIndex)
     }

diff --git a/src/main/scala/com/databricks/labs/mosaic/core/raster/io/RasterReader.scala b/src/main/scala/com/databricks/labs/mosaic/core/raster/io/RasterReader.scala
@@ -42,6 +42,9 @@ trait RasterReader extends Logging {
      */
     def readRaster(contentBytes: Array[Byte], createInfo: Map[String, String], unsafe: Option[Boolean]): MosaicRasterGDAL
 
+    def readRaster(contentBytes: Array[Byte], createInfo: Map[String, String]): MosaicRasterGDAL =
+        readRaster(contentBytes, createInfo, None)
+
     /**
       * Reads a raster from a file system path. Reads a subdataset if the path
       * is to a subdataset. Assumes "path" is a key in createInfo.
@@ -58,4 +61,7 @@ trait RasterReader extends Logging {
       */
     def readRaster(createInfo: Map[String, String], unsafe: Option[Boolean]): MosaicRasterGDAL
 
+    def readRaster(createInfo: Map[String, String]): MosaicRasterGDAL =
+        readRaster(createInfo, None)
+
 }
diff --git a/src/main/scala/com/databricks/labs/mosaic/core/raster/operator/rasterize/GDALRasterize.scala b/src/main/scala/com/databricks/labs/mosaic/core/raster/operator/rasterize/GDALRasterize.scala
@@ -75,7 +75,7 @@ object GDALRasterize {
                 "last_error" -> errorMsg,
                 "all_parents" -> ""
             )
-            return MosaicRasterGDAL.readRaster(createInfo, unsafe = Option(true))
+            return MosaicRasterGDAL.readRaster(createInfo)
         }
 
         val valuesToBurn = values.getOrElse(geoms.map(_.getAnyPoint.getZ)) // can come back and make this the mean
@@ -100,7 +100,7 @@ object GDALRasterize {
           "last_error" -> errorMsg,
           "all_parents" -> ""
         )
-        MosaicRasterGDAL.readRaster(createInfo, unsafe = Option(true))
+        MosaicRasterGDAL.readRaster(createInfo)
     }
 
 

diff --git a/src/main/scala/com/databricks/labs/mosaic/datasource/gdal/ReTileOnRead.scala b/src/main/scala/com/databricks/labs/mosaic/datasource/gdal/ReTileOnRead.scala
@@ -158,7 +158,7 @@ object ReTileOnRead extends ReadStrategy {
     def localSubdivide(inPath: String, parentPath: String, sizeInMB: Int): Seq[MosaicRasterTile] = {
         val cleanPath = PathUtils.getCleanPath(inPath)
         val createInfo = Map("path" -> cleanPath, "parentPath" -> parentPath)
-        val raster = MosaicRasterGDAL.readRaster(createInfo, None)
+        val raster = MosaicRasterGDAL.readRaster(createInfo)
         val inTile = new MosaicRasterTile(null, raster)
         val tiles = BalancedSubdivision.splitRaster(inTile, sizeInMB)
         RasterCleaner.dispose(raster)

diff --git a/src/main/scala/com/databricks/labs/mosaic/datasource/gdal/ReadAsPath.scala b/src/main/scala/com/databricks/labs/mosaic/datasource/gdal/ReadAsPath.scala
@@ -93,7 +93,7 @@ object ReadAsPath extends ReadStrategy {
 
         val tmpPath = PathUtils.copyToTmp(inPath)
         val createInfo = Map("path" -> tmpPath, "parentPath" -> inPath)
-        val raster = MosaicRasterGDAL.readRaster(createInfo, None)
+        val raster = MosaicRasterGDAL.readRaster(createInfo)
         val tile = MosaicRasterTile(null, raster)
 
         val trimmedSchema = StructType(requiredSchema.filter(field => field.name != TILE))

diff --git a/src/main/scala/com/databricks/labs/mosaic/datasource/gdal/ReadInMemory.scala b/src/main/scala/com/databricks/labs/mosaic/datasource/gdal/ReadInMemory.scala
@@ -83,7 +83,7 @@ object ReadInMemory extends ReadStrategy {
             "path" -> readPath,
             "parentPath" -> inPath
         )
-        val raster = MosaicRasterGDAL.readRaster(createInfo, None)
+        val raster = MosaicRasterGDAL.readRaster(createInfo)
         val uuid = getUUID(status)
 
         val fields = requiredSchema.fieldNames.filter(_ != TILE).map {

diff --git a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_FromFile.scala b/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_FromFile.scala
@@ -70,7 +70,7 @@ case class RST_FromFile(
         val currentSize = Files.size(Paths.get(PathUtils.replaceDBFSTokens(readPath)))
         if (targetSize <= 0 && currentSize <= Integer.MAX_VALUE) {
             val createInfo = Map("path" -> readPath, "parentPath" -> path)
-            var raster = MosaicRasterGDAL.readRaster(createInfo, None)
+            var raster = MosaicRasterGDAL.readRaster(createInfo)
             var tile = MosaicRasterTile(null, raster)
             val row = tile.formatCellId(indexSystem).serialize(rasterType)
             RasterCleaner.dispose(raster)

diff --git a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_Write.scala b/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_Write.scala
@@ -75,8 +75,7 @@ case class RST_Write(
             .toString
 
         MosaicRasterGDAL.readRaster(
-          Map("path" -> outPath, "driver" -> inDriver, "parentPath" -> inPath),
-            None
+          Map("path" -> outPath, "driver" -> inDriver, "parentPath" -> inPath)
         )
     }
 

diff --git a/src/test/scala/com/databricks/labs/mosaic/core/raster/TestRasterBandGDAL.scala b/src/test/scala/com/databricks/labs/mosaic/core/raster/TestRasterBandGDAL.scala
@@ -14,7 +14,7 @@ class TestRasterBandGDAL extends SharedSparkSessionGDAL {
           "path" -> filePath("/modis/MCD43A4.A2018185.h10v07.006.2018194033728_B01.TIF"),
           "parentPath" -> filePath("/modis/MCD43A4.A2018185.h10v07.006.2018194033728_B01.TIF")
         )
-        val testRaster = MosaicRasterGDAL.readRaster(createInfo, None)
+        val testRaster = MosaicRasterGDAL.readRaster(createInfo)
         val testBand = testRaster.getBand(1)
         testBand.getBand
         testBand.index shouldBe 1
@@ -41,7 +41,7 @@ class TestRasterBandGDAL extends SharedSparkSessionGDAL {
             "path" -> filePath("/binary/grib-cams/adaptor.mars.internal-1650626995.380916-11651-14-ca8e7236-16ca-4e11-919d-bdbd5a51da35.grb"),
             "parentPath" -> filePath("/binary/grib-cams/adaptor.mars.internal-1650626995.380916-11651-14-ca8e7236-16ca-4e11-919d-bdbd5a51da35.grb")
         )
-        val testRaster = MosaicRasterGDAL.readRaster(createInfo, None)
+        val testRaster = MosaicRasterGDAL.readRaster(createInfo)
         val testBand = testRaster.getBand(1)
         testBand.description shouldBe "1[-] HYBL=\"Hybrid level\""
         testBand.dataType shouldBe 7
@@ -61,13 +61,13 @@ class TestRasterBandGDAL extends SharedSparkSessionGDAL {
             "path" -> filePath("/binary/netcdf-coral/ct5km_baa-max-7d_v3.1_20220101.nc"),
             "parentPath" -> filePath("/binary/netcdf-coral/ct5km_baa-max-7d_v3.1_20220101.nc")
         )
-        val superRaster = MosaicRasterGDAL.readRaster(createInfo, None)
+        val superRaster = MosaicRasterGDAL.readRaster(createInfo)
         val subdatasetPath = superRaster.subdatasets("bleaching_alert_area")
         val sdCreate  = Map(
             "path" -> subdatasetPath,
             "parentPath" -> subdatasetPath
         )
-        val testRaster = MosaicRasterGDAL.readRaster(sdCreate, None)
+        val testRaster = MosaicRasterGDAL.readRaster(sdCreate)
 
         val testBand = testRaster.getBand(1)
         testBand.dataType shouldBe 1

diff --git a/src/test/scala/com/databricks/labs/mosaic/core/raster/TestRasterGDAL.scala b/src/test/scala/com/databricks/labs/mosaic/core/raster/TestRasterGDAL.scala
@@ -62,7 +62,7 @@ class TestRasterGDAL extends SharedSparkSessionGDAL {
           "path" -> filePath("/modis/MCD43A4.A2018185.h10v07.006.2018194033728_B01.TIF"),
           "parentPath" -> filePath("/modis/MCD43A4.A2018185.h10v07.006.2018194033728_B01.TIF")
         )
-        val testRaster = MosaicRasterGDAL.readRaster(createInfo, None)
+        val testRaster = MosaicRasterGDAL.readRaster(createInfo)
         testRaster.xSize shouldBe 2400
         testRaster.ySize shouldBe 2400
         testRaster.numBands shouldBe 1
@@ -84,7 +84,7 @@ class TestRasterGDAL extends SharedSparkSessionGDAL {
           "path" -> filePath("/binary/grib-cams/adaptor.mars.internal-1650626995.380916-11651-14-ca8e7236-16ca-4e11-919d-bdbd5a51da35.grb"),
           "parentPath" -> filePath("/binary/grib-cams/adaptor.mars.internal-1650626995.380916-11651-14-ca8e7236-16ca-4e11-919d-bdbd5a51da35.grb")
         )
-        val testRaster = MosaicRasterGDAL.readRaster(createInfo, None)
+        val testRaster = MosaicRasterGDAL.readRaster(createInfo)
         testRaster.xSize shouldBe 14
         testRaster.ySize shouldBe 14
         testRaster.numBands shouldBe 14
@@ -102,14 +102,14 @@ class TestRasterGDAL extends SharedSparkSessionGDAL {
           "path" -> filePath("/binary/netcdf-coral/ct5km_baa-max-7d_v3.1_20220101.nc"),
           "parentPath" -> filePath("/binary/netcdf-coral/ct5km_baa-max-7d_v3.1_20220101.nc")
         )
-        val superRaster = MosaicRasterGDAL.readRaster(createInfo, None)
+        val superRaster = MosaicRasterGDAL.readRaster(createInfo)
         val subdatasetPath = superRaster.subdatasets("bleaching_alert_area")
 
         val sdCreateInfo = Map(
           "path" -> subdatasetPath,
           "parentPath" -> subdatasetPath
         )
-        val testRaster = MosaicRasterGDAL.readRaster(sdCreateInfo, None)
+        val testRaster = MosaicRasterGDAL.readRaster(sdCreateInfo)
 
         testRaster.xSize shouldBe 7200
         testRaster.ySize shouldBe 3600
@@ -129,7 +129,7 @@ class TestRasterGDAL extends SharedSparkSessionGDAL {
           "path" -> filePath("/modis/MCD43A4.A2018185.h10v07.006.2018194033728_B01.TIF"),
           "parentPath" -> filePath("/modis/MCD43A4.A2018185.h10v07.006.2018194033728_B01.TIF")
         )
-        val testRaster = MosaicRasterGDAL.readRaster(createInfo, None)
+        val testRaster = MosaicRasterGDAL.readRaster(createInfo)
 
         testRaster.pixelXSize - 463.312716527 < 0.0000001 shouldBe true
         testRaster.pixelYSize - -463.312716527 < 0.0000001 shouldBe true

diff --git a/src/test/scala/com/databricks/labs/mosaic/expressions/raster/RST_ClipBehaviors.scala b/src/test/scala/com/databricks/labs/mosaic/expressions/raster/RST_ClipBehaviors.scala
@@ -61,7 +61,7 @@ trait RST_ClipBehaviors extends QueryTest {
         val testPath = filePath("/binary/geotiff-small/chicago_sp27.tif")
 
         val createInfo = Map("path" -> testPath, "parentPath" -> testPath)
-        val testRaster = MosaicRasterGDAL.readRaster(createInfo, None)
+        val testRaster = MosaicRasterGDAL.readRaster(createInfo)
 
         val ftMeters = 0.3 // ~0.3 ft in meter
         val ftUnits = 0.3 // epsg:26771 0.3 ft per unit