Initial support for writing GeoTIFF files using DataSource API. Multiband support included #43.

Signed-off-by: Simeon H.K. Fitch <[email protected]>

Author: metasim

datasource/src/main/scala/astraea/spark/rasterframes/datasource/geotiff/DefaultSource.scala

@@ -22,22 +22,69 @@ package astraea.spark.rasterframes.datasource.geotiff
 import java.net.URI
 
 import org.apache.spark.annotation.Experimental
-import org.apache.spark.sql.SQLContext
-import org.apache.spark.sql.sources.{DataSourceRegister, RelationProvider}
+import org.apache.spark.sql.{DataFrame, SQLContext, SaveMode}
+import org.apache.spark.sql.sources.{BaseRelation, CreatableRelationProvider, DataSourceRegister, RelationProvider}
 import astraea.spark.rasterframes._
+import astraea.spark.rasterframes.util._
+import com.typesafe.scalalogging.LazyLogging
+import geotrellis.raster.io.geotiff.GeoTiff
+
+import scala.util.Try
 
 /**
  *
  * @since 1/14/18
  */
 @Experimental
-class DefaultSource extends DataSourceRegister with RelationProvider {
-  def shortName() = "geotiff"
+class DefaultSource extends DataSourceRegister with RelationProvider with CreatableRelationProvider with LazyLogging {
+  def shortName() = DefaultSource.SHORT_NAME
+
+  def path(parameters: Map[String, String]) =
+    parameters.get(DefaultSource.PATH_PARAM).flatMap(p ⇒ Try(URI.create(p)).toOption)
+
+  def num(key: String, parameters: Map[String, String]) =
+    parameters.get(key).map(_.toLong)
+
 
   def createRelation(sqlContext: SQLContext, parameters: Map[String, String]) = {
-    require(parameters.contains("path"), "'path' parameter required.")
-    val uri: URI = URI.create(parameters("path"))
+    val pathO = path(parameters)
+    require(pathO.isDefined, "Valid URI 'path' parameter required.")
     sqlContext.withRasterFrames
-    GeoTiffRelation(sqlContext, uri)
+    GeoTiffRelation(sqlContext, pathO.get)
   }
+
+  override def createRelation(sqlContext: SQLContext, mode: SaveMode, parameters: Map[String, String], data: DataFrame): BaseRelation = {
+    val pathO = path(parameters)
+    require(pathO.isDefined, "Valid URI 'path' parameter required.")
+    require(pathO.get.getScheme == "file" || pathO.get.getScheme == null, "Currently only 'file://' destinations are supported")
+    sqlContext.withRasterFrames
+
+
+    require(data.isRF, "GeoTIFF can only be constructed from a RasterFrame")
+    val rf = data.certify
+
+    val tl = rf.tileLayerMetadata.widen.layout.tileLayout
+
+    val cols = num(DefaultSource.IMAGE_WIDTH_PARAM, parameters).getOrElse(tl.totalCols)
+    val rows = num(DefaultSource.IMAGE_HEIGHT_PARAM, parameters).getOrElse(tl.totalRows)
+
+    require(cols <= Int.MaxValue && rows <= Int.MaxValue, s"Can't construct a GeoTIFF of size $cols x $rows. (Too big!)")
+
+    // Should we really play traffic cop here?
+    if(cols.toDouble * rows * 64.0 > Runtime.getRuntime.totalMemory() * 0.5)
+      logger.warn(s"You've asked for the construction of a very large image ($cols x $rows), destined for ${pathO.get}. Out of memory error likely.")
+
+    println()
+    val raster = rf.toMultibandRaster(rf.tileColumns, cols.toInt, rows.toInt)
+
+    GeoTiff(raster).write(pathO.get.getPath)
+    GeoTiffRelation(sqlContext, pathO.get)
+  }
+}
+
+object DefaultSource {
+  final val SHORT_NAME = "geotiff"
+  final val PATH_PARAM = "path"
+  final val IMAGE_WIDTH_PARAM = "imageWidth"
+  final val IMAGE_HEIGHT_PARAM = "imageWidth"
 }

datasource/src/main/scala/astraea/spark/rasterframes/datasource/geotiff/GeoTiffRelation.scala

@@ -22,23 +22,21 @@ package astraea.spark.rasterframes.datasource.geotiff
 import java.net.URI
 
 import astraea.spark.rasterframes._
+import astraea.spark.rasterframes.util._
+import geotrellis.raster.TileLayout
 import geotrellis.raster.io.geotiff.reader.GeoTiffReader
 import geotrellis.spark._
 import geotrellis.spark.io._
 import geotrellis.spark.io.hadoop._
 import geotrellis.spark.tiling.LayoutDefinition
 import geotrellis.util._
-import geotrellis.vector.Extent
 import org.apache.hadoop.fs.Path
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.gt.types.TileUDT
 import org.apache.spark.sql.sources._
 import org.apache.spark.sql.types._
 import org.apache.spark.sql.{Row, SQLContext}
-import astraea.spark.rasterframes.util._
-import geotrellis.raster.TileLayout
-import geotrellis.raster.io.geotiff.MultibandGeoTiff
 
 /**
  * Spark SQL data source over a single GeoTiff file. Works best with CoG compliant ones.

datasource/src/main/scala/astraea/spark/rasterframes/datasource/geotiff/package.scala

@@ -22,25 +22,34 @@ package astraea.spark.rasterframes.datasource
 import java.net.URI
 
 import astraea.spark.rasterframes._
-import org.apache.spark.sql.DataFrameReader
+import org.apache.spark.sql.{DataFrameReader, DataFrameWriter}
 import shapeless.tag
 import shapeless.tag.@@
 
 /**
+ * Extension methods enabled by this module.
  *
  * @since 1/16/18
  */
 package object geotiff {
-
   /** Tagged type construction for enabling type-safe extension methods for loading
    * a RasterFrame in expected form. */
   type GeoTiffRasterFrameReader = DataFrameReader @@ GeoTiffRasterFrameReaderTag
   trait GeoTiffRasterFrameReaderTag
+  /** Tagged type construction for enabling type-safe extension methods for writing
+   * a RasterFrame to a GeoTIFF. */
+  type GeoTiffRasterFrameWriter[T] = DataFrameWriter[T] @@ GeoTiffRasterFrameWriterTag
+  trait GeoTiffRasterFrameWriterTag
 
   /** Adds `geotiff` format specifier to `DataFrameReader`. */
   implicit class DataFrameReaderHasGeoTiffFormat(val reader: DataFrameReader) {
     def geotiff: GeoTiffRasterFrameReader =
-      tag[GeoTiffRasterFrameReaderTag][DataFrameReader](reader.format("geotiff"))
+      tag[GeoTiffRasterFrameReaderTag][DataFrameReader](reader.format(DefaultSource.SHORT_NAME))
+  }
+
+  implicit class DataFrameWriterHasGeoTiffFormat[T](val writer: DataFrameWriter[T]) {
+    def geotiff: GeoTiffRasterFrameWriter[T] =
+      tag[GeoTiffRasterFrameWriterTag][DataFrameWriter[T]](writer.format(DefaultSource.SHORT_NAME))
   }
 
   /** Adds `loadRF` to appropriately tagged `DataFrameReader` */

datasource/src/main/scala/astraea/spark/rasterframes/datasource/geotrellis/DefaultSource.scala

@@ -38,7 +38,7 @@ import scala.util.Try
  */
 @Experimental
 class DefaultSource extends DataSourceRegister with RelationProvider with CreatableRelationProvider {
-  def shortName(): String = "geotrellis"
+  def shortName(): String = DefaultSource.SHORT_NAME
 
   /**
    * Create a GeoTrellis data source.
@@ -107,3 +107,7 @@ class DefaultSource extends DataSourceRegister with RelationProvider with Creata
     createRelation(sqlContext, parameters)
   }
 }
+
+object DefaultSource {
+  final val SHORT_NAME = "geotrellis"
+}

datasource/src/main/scala/astraea/spark/rasterframes/datasource/geotrellis/package.scala

@@ -28,7 +28,7 @@ import shapeless.tag
 import shapeless.tag.@@
 
 /**
- * Module utilities.
+ * Extension:win
  *
  * @since 1/12/18
  */
@@ -55,12 +55,12 @@ package object geotrellis {
       reader.format("geotrellis-catalog").load(base.toASCIIString)
 
     def geotrellis: GeoTrellisRasterFrameReader =
-      tag[GeoTrellisRasterFrameReaderTag][DataFrameReader](reader.format("geotrellis"))
+      tag[GeoTrellisRasterFrameReaderTag][DataFrameReader](reader.format(DefaultSource.SHORT_NAME))
   }
 
   implicit class DataFrameWriterHasGeotrellisFormat[T](val writer: DataFrameWriter[T]) {
     def geotrellis: GeoTrellisRasterFrameWriter[T] =
-      tag[GeoTrellisRasterFrameWriterTag][DataFrameWriter[T]](writer.format("geotrellis"))
+      tag[GeoTrellisRasterFrameWriterTag][DataFrameWriter[T]](writer.format(DefaultSource.SHORT_NAME))
   }
 
   implicit class GeoTrellisWriterAddLayer[T](val writer: GeoTrellisRasterFrameWriter[T]) {

datasource/src/test/scala/astraea/spark/rasterframes/datasource/geotiff/GeoTiffDataSourceSpec.scala

@@ -18,6 +18,8 @@
  */
 package astraea.spark.rasterframes.datasource.geotiff
 
+import java.nio.file.Paths
+
 import astraea.spark.rasterframes._
 
 /**
@@ -80,11 +82,24 @@ class GeoTiffDataSourceSpec
 
     }
 
-    it("should write RF to parquet") {
+    it("should write GeoTIFF RF to parquet") {
       val rf = spark.read
         .geotiff
         .loadRF(cogPath)
       assert(write(rf))
     }
+
+    it("should write GeoTIFF") {
+
+      val rf = spark.read
+        .geotiff
+        .loadRF(cogPath)
+
+      val out = Paths.get(outputLocalPath, "example-geotiff.tiff")
+      //val out = Paths.get("target", "example-geotiff.tiff")
+      noException shouldBe thrownBy {
+        rf.write.geotiff.save(out.toString)
+      }
+    }
   }
 }

docs/src/main/tut/release-notes.md

@@ -4,12 +4,12 @@
 
 ### 0.6.2
 
+* Added for for writing GeoTIFFs from RasterFrames via with DataFrameWriter.
 * _Breaking_: Renamed/moved `astraea.spark.rasterframes.functions.CellStatsAggregateFunction.Statistics` to
 `astraea.spark.rasterframes.stats.CellStatistics`.
 * _Breaking_: `HistogramAggregateFunction` now generates a `astraea.spark.rasterframes.stats.CellHistogram`.
 
 
-
 ### 0.6.1
 
 * Added support for reading striped GeoTiffs (#64).