Merge pull request #495 from s22s/feature/resampling_method

Choose ResampleMethod in rf_resample and RasterJoin

Author: metasim

core/src/main/scala/org/locationtech/rasterframes/expressions/localops/Resample.scala

@@ -22,53 +22,157 @@
 package org.locationtech.rasterframes.expressions.localops
 
 import geotrellis.raster.Tile
-import geotrellis.raster.resample.NearestNeighbor
+import geotrellis.raster.resample._
+import geotrellis.raster.resample.{ResampleMethod ⇒ GTResampleMethod, Max ⇒ RMax, Min ⇒ RMin}
 import org.apache.spark.sql.Column
 import org.apache.spark.sql.catalyst.InternalRow
+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, Literal, TernaryExpression}
 import org.apache.spark.sql.functions.lit
-import org.locationtech.rasterframes.expressions.BinaryLocalRasterOp
-import org.locationtech.rasterframes.expressions.DynamicExtractors.tileExtractor
+import org.apache.spark.sql.rf.TileUDT
+import org.apache.spark.sql.types.{DataType, StringType}
+import org.apache.spark.unsafe.types.UTF8String
+import org.locationtech.rasterframes.util.ResampleMethod
+import org.locationtech.rasterframes.encoders.CatalystSerializer._
+import org.locationtech.rasterframes.expressions.{fpTile, row}
+import org.locationtech.rasterframes.expressions.DynamicExtractors._
+
+
+abstract class ResampleBase(left: Expression, right: Expression, method: Expression)
+  extends TernaryExpression
+  with CodegenFallback with Serializable {
 
-@ExpressionDescription(
-  usage = "_FUNC_(tile, factor) - Resample tile to different size based on scalar factor or tile whose dimension to match. Scalar less than one will downsample tile; greater than one will upsample. Uses nearest-neighbor value.",
-  arguments = """
-  Arguments:
-    * tile - tile
-    * rhs  - scalar or tile to match dimension""",
-  examples = """
-  Examples:
-    > SELECT _FUNC_(tile, 2.0);
-       ...
-    > SELECT _FUNC_(tile1, tile2);
-       ..."""
-)
-case class Resample(left: Expression, right: Expression) extends BinaryLocalRasterOp
-  with CodegenFallback {
   override val nodeName: String = "rf_resample"
-  override protected def op(left: Tile, right: Tile): Tile = left.resample(right.cols, right.rows, NearestNeighbor)
-  override protected def op(left: Tile, right: Double): Tile = left.resample((left.cols * right).toInt,
-                                                                             (left.rows * right).toInt, NearestNeighbor)
-  override protected def op(left: Tile, right: Int): Tile = op(left, right.toDouble)
+  override def dataType: DataType = left.dataType
+  override def children: Seq[Expression] = Seq(left, right, method)
+
+  def targetFloatIfNeeded(t: Tile, method: GTResampleMethod): Tile =
+     method match {
+      case NearestNeighbor | Mode | RMax | RMin | Sum  ⇒ t
+      case _ ⇒ fpTile(t)
+     }
+
+   // These methods define the core algorithms to be used.
+  def op(left: Tile, right: Tile, method: GTResampleMethod): Tile =
+    op(left, right.cols, right.rows, method)
+
+  def op(left: Tile, right: Double, method: GTResampleMethod): Tile =
+    op(left, (left.cols * right).toInt, (left.rows * right).toInt, method)
+
+  def op(tile: Tile, newCols: Int, newRows: Int, method: GTResampleMethod): Tile =
+    targetFloatIfNeeded(tile, method).resample(newCols, newRows, method)
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    // copypasta from BinaryLocalRasterOp
+    if (!tileExtractor.isDefinedAt(left.dataType)) {
+      TypeCheckFailure(s"Input type '${left.dataType}' does not conform to a raster type.")
+    }
+    else if (!tileOrNumberExtractor.isDefinedAt(right.dataType)) {
+      TypeCheckFailure(s"Input type '${right.dataType}' does not conform to a compatible type.")
+    } else method.dataType match {
+      case StringType ⇒ TypeCheckSuccess
+      case _ ⇒ TypeCheckFailure(s"Cannot interpret value of type `${method.dataType.simpleString}` for resampling method; please provide a String method name.")
+    }
+  }
+
+  override def nullSafeEval(input1: Any, input2: Any, input3: Any): Any = {
+    // more copypasta from BinaryLocalRasterOp
+    implicit val tileSer = TileUDT.tileSerializer
+
+    val (leftTile, leftCtx) = tileExtractor(left.dataType)(row(input1))
+    val methodString = input3.asInstanceOf[UTF8String].toString
+
+    val resamplingMethod = methodString match {
+      case ResampleMethod(mm) => mm
+      case _ => throw new IllegalArgumentException("Unrecognized resampling method specified")
+    }
+
+    val result: Tile = tileOrNumberExtractor(right.dataType)(input2) match {
+      // in this case we expect the left and right contexts to vary. no warnings raised.
+      case TileArg(rightTile, _) ⇒ op(leftTile, rightTile, resamplingMethod)
+      case DoubleArg(d) ⇒ op(leftTile, d, resamplingMethod)
+      case IntegerArg(i) ⇒ op(leftTile, i.toDouble, resamplingMethod)
+    }
+
+    // reassemble the leftTile with its context. Note that this operation does not change Extent and CRS
+    leftCtx match {
+      case Some(ctx) ⇒ ctx.toProjectRasterTile(result).toInternalRow
+      case None ⇒ result.toInternalRow
+    }
+  }
 
   override def eval(input: InternalRow): Any = {
     if(input == null) null
     else {
       val l = left.eval(input)
       val r = right.eval(input)
-      if (l == null && r == null) null
-      else if (l == null) r
-      else if (r == null && tileExtractor.isDefinedAt(right.dataType)) l
-      else if (r == null) null
-      else nullSafeEval(l, r)
+      val m = method.eval(input)
+      if (m == null) null // no method, return null
+      else if (l == null) null  // no l tile, return null
+      else if (r == null) l // no target tile or factor, return l without changin it
+      else nullSafeEval(l, r, m)
     }
   }
+
 }
-object Resample{
-  def apply(left: Column, right: Column): Column =
-    new Column(Resample(left.expr, right.expr))
+
+@ExpressionDescription(
+  usage = "_FUNC_(tile, factor, method_name) - Resample tile to different dimension based on scalar `factor` or a tile whose dimension to match. Scalar less than one will downsample tile; greater than one will upsample. Uses resampling method named in the `method_name`." +
+    "Methods average, mode, median, max, min, and sum aggregate over cells when downsampling",
+  arguments = """
+Arguments:
+  * tile - tile
+  * factor  - scalar or tile to match dimension
+  * method_name - one the following options: nearest_neighbor, bilinear, cubic_convolution, cubic_spline, lanczos, average, mode, median, max, min, sum
+                  This option can be CamelCase as well
+""",
+  examples = """
+Examples:
+  > SELECT _FUNC_(tile, 0.2, median);
+     ...
+  > SELECT _FUNC_(tile1, tile2, lit("cubic_spline"));
+     ..."""
+)
+case class Resample(left: Expression, factor: Expression, method: Expression)
+  extends ResampleBase(left, factor, method)
+
+object Resample {
+  def apply(left: Column, right: Column, methodName: String): Column =
+    new Column(Resample(left.expr, right.expr, lit(methodName).expr))
+
+  def apply(left: Column, right: Column, method: Column): Column =
+    new Column(Resample(left.expr, right.expr, method.expr))
+
+  def apply[N: Numeric](left: Column, right: N, method: String): Column = new Column(Resample(left.expr, lit(right).expr, lit(method).expr))
+
+  def apply[N: Numeric](left: Column, right: N, method: Column): Column = new Column(Resample(left.expr, lit(right).expr, method.expr))
+
+}
+
+@ExpressionDescription(
+ usage = "_FUNC_(tile, factor) - Resample tile to different size based on scalar factor or tile whose dimension to match. Scalar less than one will downsample tile; greater than one will upsample. Uses nearest-neighbor value.",
+ arguments = """
+  Arguments:
+    * tile - tile
+    * rhs  - scalar or tile to match dimension""",
+ examples = """
+  Examples:
+    > SELECT _FUNC_(tile, 2.0);
+       ...
+    > SELECT _FUNC_(tile1, tile2);
+       ...""")
+case class ResampleNearest(tile: Expression, target: Expression)
+  extends ResampleBase(tile, target, Literal("nearest")) {
+  override val nodeName: String = "rf_resample_nearest"
+}
+object ResampleNearest {
+  def apply(tile: Column, target: Column): Column =
+    new Column(ResampleNearest(tile.expr, target.expr))
 
   def apply[N: Numeric](tile: Column, value: N): Column =
-    new Column(Resample(tile.expr, lit(value).expr))
+    new Column(ResampleNearest(tile.expr, lit(value).expr))
 }
+
+

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

@@ -109,6 +109,7 @@ package object expressions {
     registry.registerExpression[ExpM1]("rf_expm1")
     registry.registerExpression[Sqrt]("rf_sqrt")
     registry.registerExpression[Resample]("rf_resample")
+    registry.registerExpression[ResampleNearest]("rf_resample_nearest")
     registry.registerExpression[TileToArrayDouble]("rf_tile_to_array_double")
     registry.registerExpression[TileToArrayInt]("rf_tile_to_array_int")
     registry.registerExpression[DataCells]("rf_data_cells")

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

@@ -23,6 +23,7 @@ package org.locationtech.rasterframes.extensions
 
 import geotrellis.proj4.CRS
 import geotrellis.layer._
+import geotrellis.raster.resample.{NearestNeighbor, ResampleMethod => GTResampleMethod}
 import geotrellis.util.MethodExtensions
 import geotrellis.vector.Extent
 import org.apache.spark.sql.catalyst.expressions.Attribute
@@ -37,6 +38,7 @@ import org.locationtech.rasterframes.util._
 import org.locationtech.rasterframes.{MetadataKeys, RasterFrameLayer}
 import spray.json.JsonFormat
 import org.locationtech.rasterframes.util.JsonCodecs._
+
 import scala.util.Try
 
 /**
@@ -165,9 +167,10 @@ trait DataFrameMethods[DF <: DataFrame] extends MethodExtensions[DF] with Metada
     * }}}
     *
     * @param right Right side of the join.
+    * @param resampleMethod string indicating method to use for resampling.
     * @return joined dataframe
     */
-  def rasterJoin(right: DataFrame): DataFrame = RasterJoin(self, right, None)
+  def rasterJoin(right: DataFrame, resampleMethod: GTResampleMethod = NearestNeighbor): DataFrame = RasterJoin(self, right, resampleMethod, None)
 
   /**
     * Performs a jeft join on the dataframe `right` to this one, reprojecting and merging tiles as necessary.
@@ -183,10 +186,11 @@ trait DataFrameMethods[DF <: DataFrame] extends MethodExtensions[DF] with Metada
     * @param leftCRS this (left) datafrasme's CRS column
     * @param rightExtent right dataframe's CRS extent
     * @param rightCRS right dataframe's CRS column
+    * @param resampleMethod string indicating method to use for resampling.
     * @return joined dataframe
     */
-  def rasterJoin(right: DataFrame, leftExtent: Column, leftCRS: Column, rightExtent: Column, rightCRS: Column): DataFrame =
-    RasterJoin(self, right, leftExtent, leftCRS, rightExtent, rightCRS, None)
+  def rasterJoin(right: DataFrame, leftExtent: Column, leftCRS: Column, rightExtent: Column, rightCRS: Column, resampleMethod: GTResampleMethod): DataFrame =
+    RasterJoin(self, right, leftExtent, leftCRS, rightExtent, rightCRS, resampleMethod, None)
 
   /**
     * Performs a jeft join on the dataframe `right` to this one, reprojecting and merging tiles as necessary.
@@ -200,10 +204,11 @@ trait DataFrameMethods[DF <: DataFrame] extends MethodExtensions[DF] with Metada
     * @param leftCRS this (left) datafrasme's CRS column
     * @param rightExtent right dataframe's CRS extent
     * @param rightCRS right dataframe's CRS column
+    * @param resampleMethod string indicating method to use for resampling.
     * @return joined dataframe
     */
-  def rasterJoin(right: DataFrame, joinExpr: Column, leftExtent: Column, leftCRS: Column, rightExtent: Column, rightCRS: Column): DataFrame =
-    RasterJoin(self, right, joinExpr, leftExtent, leftCRS, rightExtent, rightCRS, None)
+  def rasterJoin(right: DataFrame, joinExpr: Column, leftExtent: Column, leftCRS: Column, rightExtent: Column, rightCRS: Column, resampleMethod: GTResampleMethod): DataFrame =
+    RasterJoin(self, right, joinExpr, leftExtent, leftCRS, rightExtent, rightCRS, resampleMethod, None)
 
 
   /** Layout contents of RasterFrame to a layer. Assumes CRS and extent columns exist. */

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

@@ -21,6 +21,7 @@
 
 package org.locationtech.rasterframes.extensions
 import geotrellis.raster.Dimensions
+import geotrellis.raster.resample.{NearestNeighbor, ResampleMethod => GTResampleMethod}
 import org.apache.spark.sql._
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types.DataType
@@ -36,7 +37,7 @@ import scala.util.Random
 object RasterJoin {
 
   /** Perform a raster join on dataframes that each have proj_raster columns, or crs and extent explicitly included. */
-  def apply(left: DataFrame, right: DataFrame, fallbackDimensions: Option[Dimensions[Int]]): DataFrame = {
+  def apply(left: DataFrame, right: DataFrame, resampleMethod: GTResampleMethod, fallbackDimensions: Option[Dimensions[Int]]): DataFrame = {
     def usePRT(d: DataFrame) =
       d.projRasterColumns.headOption
         .map(p => (rf_crs(p),  rf_extent(p)))
@@ -50,21 +51,21 @@ object RasterJoin {
     val (ldf, lcrs, lextent) = usePRT(left)
     val (rdf, rcrs, rextent) = usePRT(right)
 
-    apply(ldf, rdf, lextent, lcrs, rextent, rcrs, fallbackDimensions)
+    apply(ldf, rdf, lextent, lcrs, rextent, rcrs, resampleMethod, fallbackDimensions)
   }
 
-  def apply(left: DataFrame, right: DataFrame, leftExtent: Column, leftCRS: Column, rightExtent: Column, rightCRS: Column, fallbackDimensions: Option[Dimensions[Int]]): DataFrame = {
+  def apply(left: DataFrame, right: DataFrame, leftExtent: Column, leftCRS: Column, rightExtent: Column, rightCRS: Column, resampleMethod: GTResampleMethod, fallbackDimensions: Option[Dimensions[Int]]): DataFrame = {
     val leftGeom = st_geometry(leftExtent)
     val rightGeomReproj = st_reproject(st_geometry(rightExtent), rightCRS, leftCRS)
     val joinExpr = new Column(SpatialRelation.Intersects(leftGeom.expr, rightGeomReproj.expr))
-    apply(left, right, joinExpr, leftExtent, leftCRS, rightExtent, rightCRS, fallbackDimensions)
+    apply(left, right, joinExpr, leftExtent, leftCRS, rightExtent, rightCRS, resampleMethod, fallbackDimensions)
   }
 
   private def checkType[T](col: Column, description: String, extractor: PartialFunction[DataType, Any => T]): Unit = {
     require(extractor.isDefinedAt(col.expr.dataType), s"Expected column ${col} to be of type $description, but was ${col.expr.dataType}.")
   }
 
-  def apply(left: DataFrame, right: DataFrame, joinExprs: Column, leftExtent: Column, leftCRS: Column, rightExtent: Column, rightCRS: Column, fallbackDimensions: Option[Dimensions[Int]]): DataFrame = {
+  def apply(left: DataFrame, right: DataFrame, joinExprs: Column, leftExtent: Column, leftCRS: Column, rightExtent: Column, rightCRS: Column, resampleMethod: GTResampleMethod = NearestNeighbor, fallbackDimensions: Option[Dimensions[Int]] = None): DataFrame = {
     // Convert resolved column into a symbolic one.
     def unresolved(c: Column): Column = col(c.columnName)
 
@@ -85,7 +86,6 @@ object RasterJoin {
     // Post aggregation right crs. We create a new name.
     val rightCRS2 = id + "crs"
 
-
     // Gathering up various expressions we'll use to construct the result.
     // After joining We will be doing a groupBy the LHS. We have to define the aggregations to perform after the groupBy.
     // On the LHS we just want the first thing (subsequent ones should be identical.
@@ -110,7 +110,7 @@ object RasterJoin {
 
     val reprojCols = rightAggTiles.map(t => {
       reproject_and_merge(
-        col(leftExtent2), col(leftCRS2), col(t.columnName), col(rightExtent2), col(rightCRS2), destDims
+        col(leftExtent2), col(leftCRS2), col(t.columnName), col(rightExtent2), col(rightCRS2), destDims, lit(ResampleMethod(resampleMethod))
       ) as t.columnName
     })
 

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

@@ -22,6 +22,7 @@
 package org.locationtech.rasterframes.extensions
 
 import geotrellis.layer._
+import geotrellis.raster.resample.{NearestNeighbor, ResampleMethod => GTResampleMethod}
 import org.apache.spark.sql._
 import org.apache.spark.sql.functions.broadcast
 import org.locationtech.rasterframes._
@@ -30,7 +31,7 @@ import org.locationtech.rasterframes.util._
 
 /** Algorithm for projecting an arbitrary RasterFrame into a layer with consistent CRS and gridding. */
 object ReprojectToLayer {
-  def apply(df: DataFrame, tlm: TileLayerMetadata[SpatialKey]): RasterFrameLayer = {
+  def apply(df: DataFrame, tlm: TileLayerMetadata[SpatialKey], resampleMethod: Option[GTResampleMethod] = None): RasterFrameLayer = {
     // create a destination dataframe with crs and extend columns
     // use RasterJoin to do the rest.
     val gb = tlm.tileBounds
@@ -48,7 +49,7 @@ object ReprojectToLayer {
     // Create effectively a target RasterFrame, but with no tiles.
     val dest = gridItems.toSeq.toDF(SPATIAL_KEY_COLUMN.columnName, EXTENT_COLUMN.columnName, CRS_COLUMN.columnName)
 
-    val joined = RasterJoin(broadcast(dest), df, Some(tlm.tileLayout.tileDimensions))
+    val joined = RasterJoin(broadcast(dest), df, resampleMethod.getOrElse(NearestNeighbor), Some(tlm.tileLayout.tileDimensions))
 
     joined.asLayer(SPATIAL_KEY_COLUMN, tlm)
   }

core/src/main/scala/org/locationtech/rasterframes/functions/TileFunctions.scala

@@ -34,8 +34,8 @@ import org.locationtech.rasterframes.expressions.transformers.RenderPNG.{RenderC
 import org.locationtech.rasterframes.expressions.transformers._
 import org.locationtech.rasterframes.stats._
 import org.locationtech.rasterframes.tiles.ProjectedRasterTile
-import org.locationtech.rasterframes.util.{withTypedAlias, ColorRampNames, _}
-import org.locationtech.rasterframes.{encoders, singlebandTileEncoder, functions => F}
+import org.locationtech.rasterframes.util.{ColorRampNames, withTypedAlias, _}
+import org.locationtech.rasterframes.{encoders, singlebandTileEncoder, functions ⇒ F}
 
 /** Functions associated with creating and transforming tiles, including tile-wise statistics and rendering. */
 trait TileFunctions {
@@ -104,11 +104,21 @@ trait TileFunctions {
 
   /** Resample tile to different size based on scalar factor or tile whose dimension to match. Scalar less
    * than one will downsample tile; greater than one will upsample. Uses nearest-neighbor. */
-  def rf_resample[T: Numeric](tileCol: Column, factorValue: T) = Resample(tileCol, factorValue)
+  def rf_resample[T: Numeric](tileCol: Column, factorValue: T) = ResampleNearest(tileCol, factorValue)
 
   /** Resample tile to different size based on scalar factor or tile whose dimension to match. Scalar less
-   * than one will downsample tile; greater than one will upsample. Uses nearest-neighbor. */
-  def rf_resample(tileCol: Column, factorCol: Column) = Resample(tileCol, factorCol)
+    * than one will downsample tile; greater than one will upsample. Uses nearest-neighbor. */
+  def rf_resample(tileCol: Column, factorCol: Column) = ResampleNearest(tileCol, factorCol)
+
+  /**  */
+  def rf_resample[T: Numeric](tileCol: Column, factorVal: T, methodName: Column) = Resample(tileCol, factorVal, methodName)
+
+  def rf_resample[T: Numeric](tileCol: Column, factorVal: T, methodName: String) = Resample(tileCol, factorVal, methodName)
+
+  def rf_resample(tileCol: Column, factorCol: Column, methodName: Column) = Resample(tileCol, factorCol, methodName)
+
+  def rf_resample(tileCol: Column, factorCol: Column, methodName: String) = Resample(tileCol, factorCol, lit(methodName))
+
 
   /** Assign a `NoData` value to the tile column. */
   def rf_with_no_data(col: Column, nodata: Double): Column = SetNoDataValue(col, nodata)