GeoTiff spec more info logged (+1 squashed commit)

Squashed commits:
[d104949] spark.write.geotiff correct full res output for heterogenous tile size
Signed-off-by: Simeon H.K. Fitch <[email protected]>

Author: vpipkt

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

@@ -25,7 +25,8 @@ import astraea.spark.rasterframes.datasource._
 import astraea.spark.rasterframes.util._
 import com.typesafe.scalalogging.LazyLogging
 import org.apache.spark.sql.sources.{BaseRelation, CreatableRelationProvider, DataSourceRegister, RelationProvider}
-import org.apache.spark.sql.{DataFrame, SQLContext, SaveMode}
+import org.apache.spark.sql.types.LongType
+import org.apache.spark.sql.{DataFrame, SQLContext, SaveMode, functions ⇒ F}
 
 /**
  * Spark SQL data source over GeoTIFF files.
@@ -50,22 +51,44 @@ class DefaultSource extends DataSourceRegister with RelationProvider with Creata
     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.merge.layout.tileLayout
+    // If no desired image size is given, write at full size.
+    lazy val (fullResCols, fullResRows) = {
+      // get the layout size given that the tiles may be heterogenously sized
+      // first get any valid row and column in the spatial key structure
+      val sk = rf.select(SPATIAL_KEY_COLUMN).first()
+
+      val tc = rf.tileColumns.head
+
+      val c = rf
+        .where(SPATIAL_KEY_COLUMN("row") === sk.row)
+        .agg(
+          F.sum(tileDimensions(tc)("cols") cast(LongType))
+        ).first()
+        .getLong(0)
+
+      val r = rf
+        .where(SPATIAL_KEY_COLUMN("col") === sk.col)
+        .agg(
+          F.sum(tileDimensions(tc)("rows") cast(LongType))
+        ).first()
+        .getLong(0)
+
+      (c, r)
+    }
 
-    val cols = numParam(DefaultSource.IMAGE_WIDTH_PARAM, parameters).getOrElse(tl.totalCols)
-    val rows = numParam(DefaultSource.IMAGE_HEIGHT_PARAM, parameters).getOrElse(tl.totalRows)
+    val cols = numParam(DefaultSource.IMAGE_WIDTH_PARAM, parameters).getOrElse(fullResCols)
+    val rows = numParam(DefaultSource.IMAGE_HEIGHT_PARAM, parameters).getOrElse(fullResRows)
 
     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()
+    logger.debug(s"Writing DataFrame to GeoTIFF ($cols by $rows) at ${pathO.get}")
     val raster = rf.toMultibandRaster(rf.tileColumns, cols.toInt, rows.toInt)
 
     GeoTiff(raster).write(pathO.get.getPath)

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

@@ -108,11 +108,28 @@ class GeoTiffDataSourceSpec
         .geotiff
         .loadRF(cogPath)
 
+      logger.info(s"Read extent: ${rf.tileLayerMetadata.merge.extent}")
+
+      val cellsize = rf.groupBy(
+        tileDimensions(col("tile_1"))("rows") as "tileRows",
+        tileDimensions(col("tile_1"))("cols") as "tileCols"
+      ).count()
+
+      cellsize.show(false)
+
+      rf.withBounds().select(
+      tileDimensions(col("tile_1"))("rows") as "tileRows",
+      tileDimensions(col("tile_1"))("cols") as "tileCols",
+        BOUNDS_COLUMN)
+        .show(false)
+
       val out = Paths.get(outputLocalPath, "example-geotiff.tiff")
+      logger.info(s"Writing to $out")
       //val out = Paths.get("target", "example-geotiff.tiff")
       noException shouldBe thrownBy {
         rf.write.geotiff.save(out.toString)
       }
+      logger.info("Done!")
     }
   }
 }