diff --git a/src/main/scala/com/databricks/labs/mosaic/core/raster/api/GDAL.scala b/src/main/scala/com/databricks/labs/mosaic/core/raster/api/GDAL.scala index bc358d19d..8e54b8aaa 100644 --- a/src/main/scala/com/databricks/labs/mosaic/core/raster/api/GDAL.scala +++ b/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) { 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 index 9a6a0d4aa..992ea9721 100644 --- 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 @@ -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}") } 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 index 07add8e0e..74fb624a9 100644 --- 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 @@ -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. */ diff --git a/src/main/scala/com/databricks/labs/mosaic/core/types/model/MosaicRasterTile.scala b/src/main/scala/com/databricks/labs/mosaic/core/types/model/MosaicRasterTile.scala index dde04f0a6..0a7f9c510 100644 --- a/src/main/scala/com/databricks/labs/mosaic/core/types/model/MosaicRasterTile.scala +++ b/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) { diff --git a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_AsFormat.scala b/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_AsFormat.scala index b27525be6..dcc489be9 100644 --- a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_AsFormat.scala +++ b/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 diff --git a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_Convolve.scala b/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_Convolve.scala index 90325feff..3eb1f11c2 100644 --- a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_Convolve.scala +++ b/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 diff --git a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_GetSubdataset.scala b/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_GetSubdataset.scala index 88df43f25..6eb7b17b2 100644 --- a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_GetSubdataset.scala +++ b/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 diff --git a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_SetNoData.scala b/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_SetNoData.scala index a501791aa..86cb7f577 100644 --- a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_SetNoData.scala +++ b/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 diff --git a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_SetSRID.scala b/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_SetSRID.scala index 51d234fd4..5613cbe66 100644 --- a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_SetSRID.scala +++ b/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 diff --git a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_Transform.scala b/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_Transform.scala index cb4bd8d8f..ae9dc3925 100644 --- a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_Transform.scala +++ b/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 diff --git a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_UpdateType.scala b/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_UpdateType.scala index 821b051d0..cbe9d972a 100644 --- a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/RST_UpdateType.scala +++ b/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 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 index 8c37dddd8..bb7d75583 100644 --- 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 @@ -36,6 +36,7 @@ case class RST_Write( inputExpr, dirExpr, returnsRaster = true, + unsafe = false, expressionConfig = expressionConfig ) with NullIntolerant diff --git a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/base/Raster1ArgExpression.scala b/src/main/scala/com/databricks/labs/mosaic/expressions/raster/base/Raster1ArgExpression.scala index 01285b652..49e855ca6 100644 --- a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/base/Raster1ArgExpression.scala +++ b/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) diff --git a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/base/RasterToGridExpression.scala b/src/main/scala/com/databricks/labs/mosaic/expressions/raster/base/RasterToGridExpression.scala index a69f73151..7cb5020ae 100644 --- a/src/main/scala/com/databricks/labs/mosaic/expressions/raster/base/RasterToGridExpression.scala +++ b/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 {