Merge pull request #330 from s22s/feature/agg-raster-refactor

Incremental step toward easier aggregate raster generation.

Author: metasim

core/src/it/scala/org/locationtech/rasterframes/ref/RasterRefIT.scala

@@ -0,0 +1,65 @@
+/*
+ * This software is licensed under the Apache 2 license, quoted below.
+ *
+ * Copyright 2019 Astraea, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ *     [http://www.apache.org/licenses/LICENSE-2.0]
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package org.locationtech.rasterframes.ref
+
+import java.net.URI
+
+import geotrellis.proj4.LatLng
+import geotrellis.vector.Extent
+import org.locationtech.rasterframes._
+import org.locationtech.rasterframes.expressions.aggregates.TileRasterizerAggregate
+
+class RasterRefIT extends TestEnvironment {
+  describe("practical subregion reads") {
+    ignore("should construct a natural color composite") {
+      import spark.implicits._
+      def scene(idx: Int) = URI.create(s"https://landsat-pds.s3.us-west-2.amazonaws.com" +
+        s"/c1/L8/176/039/LC08_L1TP_176039_20190703_20190718_01_T1/LC08_L1TP_176039_20190703_20190718_01_T1_B$idx.TIF")
+
+      val redScene = RasterSource(scene(4))
+      // [west, south, east, north]
+      val area = Extent(31.115, 29.963, 31.148, 29.99).reproject(LatLng, redScene.crs)
+
+      val red = RasterRef(redScene, 0, Some(area), None)
+      val green = RasterRef(RasterSource(scene(3)), 0, Some(area), None)
+      val blue = RasterRef(RasterSource(scene(2)), 0, Some(area), None)
+
+      val rf = Seq((red, green, blue)).toDF("red", "green", "blue")
+      val df = rf.select(
+        rf_crs($"red"), rf_extent($"red"), rf_tile($"red"), rf_tile($"green"), rf_tile($"blue"))
+        .toDF
+
+      val raster = TileRasterizerAggregate(df, redScene.crs, None, None)
+
+      forEvery(raster.tile.statisticsDouble) { stats =>
+        stats should be ('defined)
+        stats.get.dataCells shouldBe > (1000L)
+      }
+      
+      //import geotrellis.raster.io.geotiff.{GeoTiffOptions, MultibandGeoTiff, Tiled}
+      //import geotrellis.raster.io.geotiff.compression.{DeflateCompression, NoCompression}
+      //import geotrellis.raster.io.geotiff.tags.codes.ColorSpace
+      //val tiffOptions = GeoTiffOptions(Tiled,  DeflateCompression, ColorSpace.RGB)
+      //MultibandGeoTiff(raster, raster.crs, tiffOptions).write("target/composite.tif")
+    }
+  }
+}

core/src/main/scala/org/locationtech/rasterframes/encoders/StandardSerializers.scala

@@ -73,7 +73,7 @@ trait StandardSerializers {
   implicit val gridBoundsSerializer: CatalystSerializer[GridBounds] = new CatalystSerializer[GridBounds] {
     override def schema: StructType = StructType(Seq(
       StructField("colMin", IntegerType, false),
-      StructField("rowlMin", IntegerType, false),
+      StructField("rowMin", IntegerType, false),
       StructField("colMax", IntegerType, false),
       StructField("rowMax", IntegerType, false)
     ))

core/src/main/scala/org/locationtech/rasterframes/expressions/DynamicExtractors.scala

@@ -35,7 +35,7 @@ import org.locationtech.rasterframes.tiles.ProjectedRasterTile
 
 private[rasterframes]
 object DynamicExtractors {
-  /** Partial function for pulling a tile and its contesxt from an input row. */
+  /** Partial function for pulling a tile and its context from an input row. */
   lazy val tileExtractor: PartialFunction[DataType, InternalRow => (Tile, Option[TileContext])] = {
     case _: TileUDT =>
       (row: InternalRow) =>
@@ -47,6 +47,14 @@ object DynamicExtractors {
       }
   }
 
+  lazy val rasterRefExtractor: PartialFunction[DataType, InternalRow => RasterRef] = {
+    case t if t.conformsTo[RasterRef] =>
+      (row: InternalRow) => row.to[RasterRef]
+  }
+
+  lazy val tileableExtractor: PartialFunction[DataType, InternalRow => Tile] =
+    tileExtractor.andThen(_.andThen(_._1)).orElse(rasterRefExtractor.andThen(_.andThen(_.tile)))
+
   lazy val rowTileExtractor: PartialFunction[DataType, Row => (Tile, Option[TileContext])] = {
     case _: TileUDT =>
       (row: Row) =>  (row.to[Tile](TileUDT.tileSerializer), None)

core/src/main/scala/org/locationtech/rasterframes/expressions/accessors/RealizeTile.scala

@@ -22,15 +22,17 @@
 package org.locationtech.rasterframes.expressions.accessors
 
 import geotrellis.raster.Tile
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{TypeCheckFailure, TypeCheckSuccess}
 import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
-import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionDescription}
+import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionDescription, UnaryExpression}
 import org.apache.spark.sql.rf.TileUDT
 import org.apache.spark.sql.types.DataType
 import org.apache.spark.sql.{Column, TypedColumn}
 import org.locationtech.rasterframes._
 import org.locationtech.rasterframes.encoders.CatalystSerializer._
-import org.locationtech.rasterframes.expressions.UnaryRasterOp
-import org.locationtech.rasterframes.model.TileContext
+import org.locationtech.rasterframes.expressions.DynamicExtractors._
+import org.locationtech.rasterframes.expressions._
 
 @ExpressionDescription(
   usage = "_FUNC_(raster) - Extracts the Tile component of a RasterSource, ProjectedRasterTile (or Tile) and ensures the cells are fully fetched.",
@@ -39,14 +41,22 @@ import org.locationtech.rasterframes.model.TileContext
       > SELECT _FUNC_(raster);
          ....
   """)
-case class RealizeTile(child: Expression) extends UnaryRasterOp with CodegenFallback {
+case class RealizeTile(child: Expression) extends UnaryExpression with CodegenFallback {
   override def dataType: DataType = TileType
 
   override def nodeName: String = "rf_tile"
-  implicit val tileSer = TileUDT.tileSerializer
 
-  override protected def eval(tile: Tile, ctx: Option[TileContext]): Any =
+  override def checkInputDataTypes(): TypeCheckResult = {
+    if (!tileableExtractor.isDefinedAt(child.dataType)) {
+      TypeCheckFailure(s"Input type '${child.dataType}' does not conform to a tiled raster type.")
+    } else TypeCheckSuccess
+  }
+  implicit val tileSer = TileUDT.tileSerializer
+  override protected def nullSafeEval(input: Any): Any = {
+    val in = row(input)
+    val tile = tileableExtractor(child.dataType)(in)
     (tile.toArrayTile(): Tile).toInternalRow
+  }
 }
 
 object RealizeTile {

core/src/main/scala/org/locationtech/rasterframes/expressions/aggregates/TileRasterizerAggregate.scala

@@ -24,15 +24,18 @@ package org.locationtech.rasterframes.expressions.aggregates
 import geotrellis.proj4.CRS
 import geotrellis.raster.reproject.Reproject
 import geotrellis.raster.resample.ResampleMethod
-import geotrellis.raster.{ArrayTile, CellType, Raster, Tile}
-import geotrellis.spark.TileLayerMetadata
+import geotrellis.raster.{ArrayTile, CellType, MultibandTile, ProjectedRaster, Raster, Tile}
+import geotrellis.spark.{SpatialKey, TileLayerMetadata}
 import geotrellis.vector.Extent
 import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
 import org.apache.spark.sql.types.{DataType, StructField, StructType}
-import org.apache.spark.sql.{Column, Row, TypedColumn}
+import org.apache.spark.sql.{Column, DataFrame, Row, TypedColumn}
 import org.locationtech.rasterframes._
+import org.locationtech.rasterframes.util._
 import org.locationtech.rasterframes.encoders.CatalystSerializer._
 import org.locationtech.rasterframes.expressions.aggregates.TileRasterizerAggregate.ProjectedRasterDefinition
+import org.locationtech.rasterframes.model.TileDimensions
+import org.slf4j.LoggerFactory
 
 /**
   * Aggregation function for creating a single `geotrellis.raster.Raster[Tile]` from
@@ -88,7 +91,7 @@ class TileRasterizerAggregate(prd: ProjectedRasterDefinition) extends UserDefine
 }
 
 object TileRasterizerAggregate {
-  val nodeName = "rf_tile_rasterizer_aggregate"
+  val nodeName = "rf_agg_raster"
   /**  Convenience grouping of  parameters needed for running aggregate. */
   case class ProjectedRasterDefinition(totalCols: Int, totalRows: Int, cellType: CellType, crs: CRS, extent: Extent, sampler: ResampleMethod = ResampleMethod.DEFAULT)
 
@@ -102,8 +105,73 @@ object TileRasterizerAggregate {
       val rows = actualSize.height
       new ProjectedRasterDefinition(cols, rows, tlm.cellType, tlm.crs, tlm.extent, sampler)
     }
-}
+  }
+
+  @transient
+  private lazy val logger = LoggerFactory.getLogger(getClass)
+
+  def apply(prd: ProjectedRasterDefinition, crsCol: Column, extentCol: Column, tileCol: Column): TypedColumn[Any, Raster[Tile]] = {
+
+    if (prd.totalCols.toDouble * prd.totalRows * 64.0 > Runtime.getRuntime.totalMemory() * 0.5)
+      logger.warn(
+        s"You've asked for the construction of a very large image (${prd.totalCols} x ${prd.totalRows}). Out of memory error likely.")
 
-  def apply(prd: ProjectedRasterDefinition, crsCol: Column, extentCol: Column, tileCol: Column): TypedColumn[Any, Raster[Tile]] =
     new TileRasterizerAggregate(prd)(crsCol, extentCol, tileCol).as(nodeName).as[Raster[Tile]]
+  }
+
+  def apply(df: DataFrame, destCRS: CRS, destExtent: Option[Extent], rasterDims: Option[TileDimensions]): ProjectedRaster[MultibandTile] = {
+    val tileCols = WithDataFrameMethods(df).tileColumns
+    require(tileCols.nonEmpty, "need at least one tile column")
+    // Select the anchoring Tile, Extent and CRS columns
+    val (extCol, crsCol, tileCol) = {
+      // Favor "ProjectedRaster" columns
+      val prCols = df.projRasterColumns
+      if (prCols.nonEmpty) {
+        (rf_extent(prCols.head), rf_crs(prCols.head), rf_tile(prCols.head))
+      } else {
+        // If no "ProjectedRaster" column, look for single Extent and CRS columns.
+        val crsCols = df.crsColumns
+        require(crsCols.size == 1, "Exactly one CRS column must be in DataFrame")
+        val extentCols = df.extentColumns
+        require(extentCols.size == 1, "Exactly one Extent column must be in DataFrame")
+        (extentCols.head, crsCols.head, tileCols.head)
+      }
+    }
+
+    // Scan table and constuct what the TileLayerMetadata would be in the specified destination CRS.
+    val tlm: TileLayerMetadata[SpatialKey] = df
+      .select(
+        ProjectedLayerMetadataAggregate(
+          destCRS,
+          extCol,
+          crsCol,
+          rf_cell_type(tileCol),
+          rf_dimensions(tileCol)
+        ))
+      .first()
+    logger.debug(s"Collected TileLayerMetadata: ${tlm.toString}")
+
+    val c = ProjectedRasterDefinition(tlm)
+
+    val config = rasterDims
+      .map { dims =>
+        c.copy(totalCols = dims.cols, totalRows = dims.rows)
+      }
+      .getOrElse(c)
+
+    destExtent.map { ext =>
+      c.copy(extent = ext)
+    }
+
+    val aggs = tileCols
+      .map(t => TileRasterizerAggregate(config, crsCol, extCol, rf_tile(t))("tile").as(t.columnName))
+
+    val agg = df.select(aggs: _*)
+
+    val row = agg.first()
+
+    val bands = for (i <- 0 until row.size) yield row.getAs[Tile](i)
+
+    ProjectedRaster(MultibandTile(bands), tlm.extent, tlm.crs)
+  }
 }

core/src/main/scala/org/locationtech/rasterframes/expressions/transformers/RGBComposite.scala

@@ -56,9 +56,9 @@ case class RGBComposite(red: Expression, green: Expression, blue: Expression) ex
   override def nodeName: String = "rf_rgb_composite"
 
   override def dataType: DataType = if(
-    red.dataType.conformsTo[ProjectedRasterTile] ||
-    blue.dataType.conformsTo[ProjectedRasterTile] ||
-    green.dataType.conformsTo[ProjectedRasterTile]
+    tileExtractor.isDefinedAt(red.dataType) ||
+      tileExtractor.isDefinedAt(green.dataType) ||
+      tileExtractor.isDefinedAt(blue.dataType)
   ) red.dataType
   else TileType
 

core/src/main/scala/org/locationtech/rasterframes/extensions/DataFrameMethods.scala

@@ -22,6 +22,7 @@
 package org.locationtech.rasterframes.extensions
 
 import geotrellis.proj4.CRS
+import geotrellis.raster.{MultibandTile, ProjectedRaster}
 import geotrellis.spark.io._
 import geotrellis.spark.{SpaceTimeKey, SpatialComponent, SpatialKey, TemporalKey, TileLayerMetadata}
 import geotrellis.util.MethodExtensions
@@ -32,7 +33,9 @@ import org.apache.spark.sql.{Column, DataFrame, TypedColumn}
 import org.locationtech.rasterframes.StandardColumns._
 import org.locationtech.rasterframes.encoders.CatalystSerializer._
 import org.locationtech.rasterframes.encoders.StandardEncoders._
-import org.locationtech.rasterframes.expressions.DynamicExtractors
+import org.locationtech.rasterframes.expressions.{DynamicExtractors, aggregates}
+import org.locationtech.rasterframes.expressions.aggregates.TileRasterizerAggregate
+import org.locationtech.rasterframes.model.TileDimensions
 import org.locationtech.rasterframes.tiles.ProjectedRasterTile
 import org.locationtech.rasterframes.util._
 import org.locationtech.rasterframes.{MetadataKeys, RasterFrameLayer}
@@ -225,7 +228,7 @@ trait DataFrameMethods[DF <: DataFrame] extends MethodExtensions[DF] with Metada
     */
   @throws[IllegalArgumentException]
   def asLayer: RasterFrameLayer = {
-    val potentialRF = certifyRasterframe(self)
+    val potentialRF = certifyLayer(self)
 
     require(
       potentialRF.findSpatialKeyField.nonEmpty,
@@ -301,5 +304,5 @@ trait DataFrameMethods[DF <: DataFrame] extends MethodExtensions[DF] with Metada
 
   /** Internal method for slapping the RasterFreameLayer seal of approval on a DataFrame.
    * Only call if if you are sure it has a spatial key and tile columns and TileLayerMetadata. */
-  private[rasterframes] def certify = certifyRasterframe(self)
+  private[rasterframes] def certify = certifyLayer(self)
 }

core/src/main/scala/org/locationtech/rasterframes/util/package.scala

@@ -77,7 +77,7 @@ package object util extends DataFrameRenderers {
   type KeyMethodsProvider[K1, K2] = K1 ⇒ TilerKeyMethods[K1, K2]
 
   /** Internal method for slapping the RasterFrameLayer seal of approval on a DataFrame. */
-  private[rasterframes] def certifyRasterframe(df: DataFrame): RasterFrameLayer =
+  private[rasterframes] def certifyLayer(df: DataFrame): RasterFrameLayer =
     shapeless.tag[RasterFrameTag][DataFrame](df)
 
 

core/src/test/scala/org/locationtech/rasterframes/RasterFunctionsSpec.scala

@@ -23,53 +23,22 @@ package org.locationtech.rasterframes
 
 import java.io.ByteArrayInputStream
 
-import geotrellis.proj4.LatLng
 import geotrellis.raster
 import geotrellis.raster._
 import geotrellis.raster.render.ColorRamps
 import geotrellis.raster.testkit.RasterMatchers
-import geotrellis.vector.Extent
 import javax.imageio.ImageIO
 import org.apache.spark.sql.Encoders
 import org.apache.spark.sql.functions._
 import org.locationtech.rasterframes.expressions.accessors.ExtractTile
 import org.locationtech.rasterframes.model.TileDimensions
-import org.locationtech.rasterframes.ref.{RasterRef, RasterSource}
 import org.locationtech.rasterframes.stats._
 import org.locationtech.rasterframes.tiles.ProjectedRasterTile
 
 class RasterFunctionsSpec extends TestEnvironment with RasterMatchers {
+  import TestData._
   import spark.implicits._
 
-  val extent = Extent(10, 20, 30, 40)
-  val crs = LatLng
-  val ct = ByteUserDefinedNoDataCellType(-2)
-  val cols = 10
-  val rows = cols
-  val tileSize = cols * rows
-  val tileCount = 10
-  val numND = 4
-  lazy val zero = TestData.projectedRasterTile(cols, rows, 0, extent, crs, ct)
-  lazy val one = TestData.projectedRasterTile(cols, rows, 1, extent, crs, ct)
-  lazy val two = TestData.projectedRasterTile(cols, rows, 2, extent, crs, ct)
-  lazy val three = TestData.projectedRasterTile(cols, rows, 3, extent, crs, ct)
-  lazy val six = ProjectedRasterTile(three * two, three.extent, three.crs)
-  lazy val nd = TestData.projectedRasterTile(cols, rows, -2, extent, crs, ct)
-  lazy val randPRT = TestData.projectedRasterTile(cols, rows, scala.util.Random.nextInt(), extent, crs, ct)
-  lazy val randNDPRT: Tile  = TestData.injectND(numND)(randPRT)
-
-  lazy val randDoubleTile = TestData.projectedRasterTile(cols, rows, scala.util.Random.nextGaussian(), extent, crs, DoubleConstantNoDataCellType)
-  lazy val randDoubleNDTile  = TestData.injectND(numND)(randDoubleTile)
-  lazy val randPositiveDoubleTile = TestData.projectedRasterTile(cols, rows, scala.util.Random.nextDouble() + 1e-6, extent, crs, DoubleConstantNoDataCellType)
-
-  val expectedRandNoData: Long = numND * tileCount.toLong
-  val expectedRandData: Long = cols * rows * tileCount - expectedRandNoData
-  lazy val randNDTilesWithNull = Seq.fill[Tile](tileCount)(TestData.injectND(numND)(
-    TestData.randomTile(cols, rows, UByteConstantNoDataCellType)
-  )).map(ProjectedRasterTile(_, extent, crs)) :+ null
-
-  def lazyPRT = RasterRef(RasterSource(TestData.l8samplePath), 0, None, None).tile
-
   implicit val pairEnc = Encoders.tuple(ProjectedRasterTile.prtEncoder, ProjectedRasterTile.prtEncoder)
   implicit val tripEnc = Encoders.tuple(ProjectedRasterTile.prtEncoder, ProjectedRasterTile.prtEncoder, ProjectedRasterTile.prtEncoder)