explicitly allow non-thread safe access to rasters

src/main/scala/com/databricks/labs/mosaic/core/raster/api/GDAL.scala

@@ -117,7 +117,8 @@ object GDAL {
     def readRaster(
                       inputRaster: Any,
                       createInfo: Map[String, String],
-                      inputDT: DataType
+                      inputDT: DataType,
+                      unsafe: Option[Boolean] = None
                   ): MosaicRasterGDAL = {
         if (inputRaster == null) {
             MosaicRasterGDAL(null, createInfo)
@@ -128,7 +129,7 @@ object GDAL {
                 case _: BinaryType =>
                     val bytes = inputRaster.asInstanceOf[Array[Byte]]
                     try {
-                        val rasterObj = MosaicRasterGDAL.readRaster(bytes, createInfo)
+                        val rasterObj = MosaicRasterGDAL.readRaster(bytes, createInfo, unsafe)
                         if (rasterObj.raster == null) {
                             val rasterZipObj = readParentZipBinary(bytes, createInfo)
                             if (rasterZipObj.raster == null) {

src/main/scala/com/databricks/labs/mosaic/core/raster/gdal/MosaicRasterGDAL.scala

@@ -733,13 +733,17 @@ object MosaicRasterGDAL extends RasterReader {
       * @return
       *   A GDAL [[Dataset]] object.
       */
-    def pathAsDataset(path: String, driverShortName: Option[String]): Dataset = {
+    def pathAsDataset(path: String, driverShortName: Option[String], unsafe: Option[Boolean] = None): Dataset = {
+        val openFlags = unsafe match {
+            case Some(value) if (value) => OF_READONLY
+            case _ => OF_THREAD_SAFE
+        }
         driverShortName match {
             case Some(driverShortName) =>
                 val drivers = new JVector[String]()
                 drivers.add(driverShortName)
-                gdal.OpenEx(path, OF_RASTER | OF_THREAD_SAFE, drivers)
-            case None                  => gdal.Open(path, OF_RASTER | OF_THREAD_SAFE)
+                gdal.OpenEx(path, openFlags, drivers)
+            case None                  => gdal.Open(path, openFlags)
         }
     }
 
@@ -772,7 +776,7 @@ object MosaicRasterGDAL extends RasterReader {
       * @return
       *   A [[MosaicRasterGDAL]] object.
      */
-    override def readRaster(contentBytes: Array[Byte], createInfo: Map[String, String]): MosaicRasterGDAL = {
+    override def readRaster(contentBytes: Array[Byte], createInfo: Map[String, String], unsafe: Option[Boolean] = None): MosaicRasterGDAL = {
         if (Option(contentBytes).isEmpty || contentBytes.isEmpty) {
             MosaicRasterGDAL(null, createInfo)
         } else {
@@ -783,12 +787,12 @@ object MosaicRasterGDAL extends RasterReader {
             val tmpPath = PathUtils.createTmpFilePath(extension)
             Files.write(Paths.get(tmpPath), contentBytes)
             // Try reading as a tmp file, if that fails, rename as a zipped file
-            val dataset = pathAsDataset(tmpPath, Some(driverShortName))
+            val dataset = pathAsDataset(tmpPath, Some(driverShortName), unsafe)
             if (dataset == null) {
                 val zippedPath = s"$tmpPath.zip"
                 Files.move(Paths.get(tmpPath), Paths.get(zippedPath), StandardCopyOption.REPLACE_EXISTING)
                 val readPath = PathUtils.getZipPath(zippedPath)
-                val ds1 = pathAsDataset(readPath, Some(driverShortName))
+                val ds1 = pathAsDataset(readPath, Some(driverShortName), unsafe)
                 if (ds1 == null) {
                     // the way we zip using uuid is not compatible with GDAL
                     // we need to unzip and read the file if it was zipped by us
@@ -800,7 +804,7 @@ object MosaicRasterGDAL extends RasterReader {
                     val extension = GDAL.getExtension(driverShortName)
                     val lastExtracted = SysUtils.getLastOutputLine(prompt)
                     val unzippedPath = PathUtils.parseUnzippedPathFromExtracted(lastExtracted, extension)
-                    val ds2 = pathAsDataset(unzippedPath, Some(driverShortName))
+                    val ds2 = pathAsDataset(unzippedPath, Some(driverShortName), unsafe)
                     if (ds2 == null) {
                         throw new Exception(s"Error reading raster from bytes: ${prompt._3}")
                     }

src/main/scala/com/databricks/labs/mosaic/core/raster/io/RasterReader.scala

@@ -40,7 +40,7 @@ trait RasterReader extends Logging {
      * @return
      *   A [[MosaicRasterGDAL]] object.
      */
-    def readRaster(contentBytes: Array[Byte], createInfo: Map[String, String]): MosaicRasterGDAL
+    def readRaster(contentBytes: Array[Byte], createInfo: Map[String, String], unsafe: Option[Boolean]): MosaicRasterGDAL
 
     /**
       * Reads a raster from a file system path. Reads a subdataset if the path
@@ -51,6 +51,8 @@ trait RasterReader extends Logging {
       *   "FORMAT:/path/to/file.tif:subdataset"
       * @param createInfo
       *   Map of create info for the raster.
+      * @param unsafe
+      *   Open this raster in an non-thread safe mode.
       * @return
       *   A [[MosaicRasterGDAL]] object.
       */

src/main/scala/com/databricks/labs/mosaic/core/types/model/MosaicRasterTile.scala

@@ -169,12 +169,12 @@ object MosaicRasterTile {
       * @return
       *   An instance of [[MosaicRasterTile]].
       */
-    def deserialize(row: InternalRow, idDataType: DataType, rasterDataType: DataType): MosaicRasterTile = {
+    def deserialize(row: InternalRow, idDataType: DataType, rasterDataType: DataType, unsafe: Option[Boolean] = None): MosaicRasterTile = {
         val index = row.get(0, idDataType)
         // handle checkpoint related de-serialization
         val rawRaster = row.get(1, rasterDataType)
         val createInfo = extractMap(row.getMap(2))
-        val raster = GDAL.readRaster(rawRaster, createInfo, rasterDataType)
+        val raster = GDAL.readRaster(rawRaster, createInfo, rasterDataType, unsafe)
 
         // noinspection TypeCheckCanBeMatch
         if (Option(index).isDefined) {

src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_AsFormat.scala

@@ -21,6 +21,7 @@ case class RST_AsFormat (
     tileExpr,
     newFormat,
     returnsRaster = true,
+    unsafe = false,
     expressionConfig
 )
     with NullIntolerant

src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_Convolve.scala

@@ -22,6 +22,7 @@ case class RST_Convolve(
       rastersExpr,
       kernelExpr,
       returnsRaster = true,
+      unsafe = false,
       expressionConfig = expressionConfig
     )
       with NullIntolerant

src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_GetSubdataset.scala

@@ -21,6 +21,7 @@ case class RST_GetSubdataset(
       tileExpr,
       subsetName,
       returnsRaster = true,
+      unsafe = false,
       expressionConfig
     )
       with NullIntolerant

src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_SetNoData.scala

@@ -23,6 +23,7 @@ case class RST_SetNoData(
       tileExpr,
       noDataExpr,
       returnsRaster = true,
+      unsafe = false,
       expressionConfig = expressionConfig
     )
       with NullIntolerant

src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_SetSRID.scala

@@ -21,6 +21,7 @@ case class RST_SetSRID(
       rastersExpr,
       sridExpr,
       returnsRaster = true,
+      unsafe = true,
       expressionConfig = expressionConfig
     )
       with NullIntolerant

src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_Transform.scala

@@ -22,6 +22,7 @@ case class RST_Transform(
       tileExpr,
       srid,
       returnsRaster = true,
+      unsafe = false,
       expressionConfig
     )
       with NullIntolerant

src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_UpdateType.scala

@@ -21,6 +21,7 @@ case class RST_UpdateType (
     tileExpr,
     newType,
     returnsRaster = true,
+    unsafe = false,
     expressionConfig
 )
     with NullIntolerant

src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_Write.scala

@@ -36,6 +36,7 @@ case class RST_Write(
       inputExpr,
       dirExpr,
       returnsRaster = true,
+      unsafe = false,
       expressionConfig = expressionConfig
     )
       with NullIntolerant

src/main/scala/com/databricks/labs/mosaic/expressions/raster/base/Raster1ArgExpression.scala

@@ -21,6 +21,10 @@ import scala.reflect.ClassTag
   *   containing the raster file content.
   * @param arg1Expr
   *   The expression for the first argument.
+  * @param returnsRaster
+  *   Does the expression return a raster tile object?
+  * @param unsafe
+  *   Does the operation require a raster opened in non-thread safe mode?
   * @param expressionConfig
   *   Additional arguments for the expression (expressionConfigs).
   * @tparam T
@@ -30,6 +34,7 @@ abstract class Raster1ArgExpression[T <: Expression: ClassTag](
     rasterExpr: Expression,
     arg1Expr: Expression,
     returnsRaster: Boolean,
+    unsafe: Boolean,
     expressionConfig: MosaicExpressionConfig
 ) extends BinaryExpression
       with NullIntolerant
@@ -71,9 +76,10 @@ abstract class Raster1ArgExpression[T <: Expression: ClassTag](
         GDAL.enable(expressionConfig)
         val rasterType = RasterTileType(rasterExpr, expressionConfig.isRasterUseCheckpoint).rasterType
         val tile = MosaicRasterTile.deserialize(
-          input.asInstanceOf[InternalRow],
-          expressionConfig.getCellIdType,
-          rasterType
+            input.asInstanceOf[InternalRow],
+            expressionConfig.getCellIdType,
+            rasterType,
+            Some(unsafe)
         )
         val raster = tile.getRaster
         val result = rasterTransform(tile, arg1)

src/main/scala/com/databricks/labs/mosaic/expressions/raster/base/RasterToGridExpression.scala

@@ -39,7 +39,7 @@ abstract class RasterToGridExpression[T <: Expression: ClassTag, P](
     resolutionExpr: Expression,
     measureType: DataType,
     expressionConfig: MosaicExpressionConfig
-) extends Raster1ArgExpression[T](rasterExpr, resolutionExpr, returnsRaster = false, expressionConfig)
+) extends Raster1ArgExpression[T](rasterExpr, resolutionExpr, returnsRaster = false, unsafe = false, expressionConfig)
       with RasterGridExpression
       with NullIntolerant
       with Serializable {