Merge remote-tracking branch 'lt/develop' into feature/mask_by_bits

Author: vpipkt

core/src/main/scala/org/locationtech/rasterframes/RasterFunctions.scala

@@ -312,35 +312,19 @@ trait RasterFunctions {
   def rf_mask_by_value(sourceTile: Column, maskTile: Column, maskValue: Int): TypedColumn[Any, Tile] =
     rf_mask_by_value(sourceTile, maskTile, maskValue, false)
 
-  /** Generate a tile with the values from `data_tile`, but where cells in the `mask_tile` are in the `mask_values`
-       list, replace the value with NODATA.
-       If `inverse` is True, the cells in `mask_tile` that are not in `mask_values` list become NODATA */
-  def rf_mask_by_values(sourceTile: Column, maskTile: Column, maskValues: Column, inverse: Boolean): TypedColumn[Any, Tile] =
-    if (!inverse)
-      Mask.MaskByValue(sourceTile, rf_local_is_in(maskTile, maskValues), lit(1))
-    else
-      Mask.InverseMaskByValue(sourceTile, rf_local_is_in(maskTile, maskValues), lit(0))
-
   /** Generate a tile with the values from `data_tile`, but where cells in the `mask_tile` are in the `mask_values`
        list, replace the value with NODATA. */
   def rf_mask_by_values(sourceTile: Column, maskTile: Column, maskValues: Column): TypedColumn[Any, Tile] =
-    rf_mask_by_values(sourceTile, maskTile, maskValues, false)
+      Mask.MaskByValues(sourceTile, maskTile, maskValues)
 
   /** Generate a tile with the values from `data_tile`, but where cells in the `mask_tile` are in the `mask_values`
-       list, replace the value with NODATA.
-       If `inverse` is True, the cells in `mask_tile` that are not in `mask_values` list become NODATA */
-  def rf_mask_by_values(sourceTile: Column, maskTile: Column, maskValues: Iterable[Int], inverse: Boolean): TypedColumn[Any, Tile] = {
+       list, replace the value with NODATA. */
+  def rf_mask_by_values(sourceTile: Column, maskTile: Column, maskValues: Seq[Int]): TypedColumn[Any, Tile] = {
     import org.apache.spark.sql.functions.array
     val valuesCol: Column = array(maskValues.map(lit).toSeq: _*)
-    rf_mask_by_values(sourceTile, maskTile, valuesCol, inverse)
+    rf_mask_by_values(sourceTile, maskTile, valuesCol)
   }
 
-  /** Generate a tile with the values from `data_tile`, but where cells in the `mask_tile` are in the `mask_values`
-       list, replace the value with NODATA.
-       If `inverse` is True, the cells in `mask_tile` that are not in `mask_values` list become NODATA */
-  def rf_mask_by_values(sourceTile: Column, maskTile: Column, maskValues: Iterable[Int]): TypedColumn[Any, Tile] =
-    rf_mask_by_values(sourceTile, maskTile, maskValues, false)
-
   /** Where the `maskTile` does **not** contain `NoData`, replace values in the source tile with `NoData` */
   def rf_inverse_mask(sourceTile: Column, maskTile: Column): TypedColumn[Any, Tile] =
     Mask.InverseMaskByDefined(sourceTile, maskTile)

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

@@ -138,7 +138,7 @@ object TileRasterizerAggregate {
       }
     }
 
-    // Scan table and constuct what the TileLayerMetadata would be in the specified destination CRS.
+    // Scan table and construct what the TileLayerMetadata would be in the specified destination CRS.
     val tlm: TileLayerMetadata[SpatialKey] = df
       .select(
         ProjectedLayerMetadataAggregate(

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

@@ -125,9 +125,10 @@ package object expressions {
     registry.registerExpression[LocalMeanAggregate]("rf_agg_local_mean")
 
     registry.registerExpression[Mask.MaskByDefined]("rf_mask")
+    registry.registerExpression[Mask.InverseMaskByDefined]("rf_inverse_mask")
     registry.registerExpression[Mask.MaskByValue]("rf_mask_by_value")
     registry.registerExpression[Mask.InverseMaskByValue]("rf_inverse_mask_by_value")
-    registry.registerExpression[Mask.InverseMaskByDefined]("rf_inverse_mask")
+    registry.registerExpression[Mask.MaskByValues]("rf_mask_by_values")
 
     registry.registerExpression[DebugRender.RenderAscii]("rf_render_ascii")
     registry.registerExpression[DebugRender.RenderMatrix]("rf_render_matrix")

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

@@ -34,52 +34,58 @@ import org.apache.spark.sql.types.DataType
 import org.apache.spark.sql.{Column, TypedColumn}
 import org.locationtech.rasterframes.encoders.CatalystSerializer._
 import org.locationtech.rasterframes.expressions.DynamicExtractors._
+import org.locationtech.rasterframes.expressions.localops.IsIn
 import org.locationtech.rasterframes.expressions.row
 import org.slf4j.LoggerFactory
 
 abstract class Mask(val left: Expression, val middle: Expression, val right: Expression, inverse: Boolean)
   extends TernaryExpression with CodegenFallback with Serializable {
+  // aliases.
+  def targetExp = left
+  def maskExp = middle
+  def maskValueExp = right
 
   @transient protected lazy val logger = Logger(LoggerFactory.getLogger(getClass.getName))
 
-
   override def children: Seq[Expression] = Seq(left, middle, right)
 
   override def checkInputDataTypes(): TypeCheckResult = {
-    if (!tileExtractor.isDefinedAt(left.dataType)) {
-      TypeCheckFailure(s"Input type '${left.dataType}' does not conform to a raster type.")
-    } else if (!tileExtractor.isDefinedAt(middle.dataType)) {
-      TypeCheckFailure(s"Input type '${middle.dataType}' does not conform to a raster type.")
-    } else if (!intArgExtractor.isDefinedAt(right.dataType)) {
-      TypeCheckFailure(s"Input type '${right.dataType}' isn't an integral type.")
+    if (!tileExtractor.isDefinedAt(targetExp.dataType)) {
+      TypeCheckFailure(s"Input type '${targetExp.dataType}' does not conform to a raster type.")
+    } else if (!tileExtractor.isDefinedAt(maskExp.dataType)) {
+      TypeCheckFailure(s"Input type '${maskExp.dataType}' does not conform to a raster type.")
+    } else if (!intArgExtractor.isDefinedAt(maskValueExp.dataType)) {
+      TypeCheckFailure(s"Input type '${maskValueExp.dataType}' isn't an integral type.")
     } else TypeCheckSuccess
   }
   override def dataType: DataType = left.dataType
 
-  override protected def nullSafeEval(leftInput: Any, middleInput: Any, rightInput: Any): Any = {
+  override def makeCopy(newArgs: Array[AnyRef]): Expression = super.makeCopy(newArgs)
+
+  override protected def nullSafeEval(targetInput: Any, maskInput: Any, maskValueInput: Any): Any = {
     implicit val tileSer = TileUDT.tileSerializer
-    val (leftTile, leftCtx) = tileExtractor(left.dataType)(row(leftInput))
-    val (rightTile, rightCtx) = tileExtractor(middle.dataType)(row(middleInput))
+    val (targetTile, targetCtx) = tileExtractor(targetExp.dataType)(row(targetInput))
+    val (maskTile, maskCtx) = tileExtractor(maskExp.dataType)(row(maskInput))
 
-    if (leftCtx.isEmpty && rightCtx.isDefined)
+    if (targetCtx.isEmpty && maskCtx.isDefined)
       logger.warn(
           s"Right-hand parameter '${middle}' provided an extent and CRS, but the left-hand parameter " +
             s"'${left}' didn't have any. Because the left-hand side defines output type, the right-hand context will be lost.")
 
-    if (leftCtx.isDefined && rightCtx.isDefined && leftCtx != rightCtx)
+    if (targetCtx.isDefined && maskCtx.isDefined && targetCtx != maskCtx)
       logger.warn(s"Both '${left}' and '${middle}' provided an extent and CRS, but they are different. Left-hand side will be used.")
 
-    val maskValue = intArgExtractor(right.dataType)(rightInput)
+    val maskValue = intArgExtractor(maskValueExp.dataType)(maskValueInput)
 
-    val masking = if (maskValue.value == 0) Defined(rightTile)
-    else rightTile
+    val masking = if (maskValue.value == 0) Defined(maskTile)
+    else maskTile
 
     val result = if (inverse)
-      gtInverseMask(leftTile, masking, maskValue.value, raster.NODATA)
+      gtInverseMask(targetTile, masking, maskValue.value, raster.NODATA)
     else
-      gtMask(leftTile, masking, maskValue.value, raster.NODATA)
+      gtMask(targetTile, masking, maskValue.value, raster.NODATA)
 
-    leftCtx match {
+    targetCtx match {
       case Some(ctx) => ctx.toProjectRasterTile(result).toInternalRow
       case None      => result.toInternalRow
     }
@@ -169,4 +175,28 @@ object Mask {
     def apply(srcTile: Column, maskingTile: Column, maskValue: Column): TypedColumn[Any, Tile] =
       new Column(InverseMaskByValue(srcTile.expr, maskingTile.expr, maskValue.expr)).as[Tile]
   }
+
+  @ExpressionDescription(
+    usage = "_FUNC_(data, mask, maskValues) - Generate a tile with the values from `data` tile but where cells in the `mask` tile are in the `maskValues` list, replace the value with NODATA.",
+    arguments = """
+  Arguments:
+    * target - tile to mask
+    * mask - masking definition
+    * maskValues - sequence of values to consider as masks candidates
+        """,
+    examples = """
+  Examples:
+    > SELECT _FUNC_(data, mask, array(1, 2, 3))
+      ..."""
+  )
+  case class MaskByValues(dataTile: Expression, maskTile: Expression)
+    extends Mask(dataTile, maskTile, Literal(1), inverse = false) {
+    def this(dataTile: Expression, maskTile: Expression, maskValues: Expression) =
+      this(dataTile, IsIn(maskTile, maskValues))
+    override def nodeName: String = "rf_mask_by_values"
+  }
+  object MaskByValues {
+    def apply(dataTile: Column, maskTile: Column, maskValues: Column): TypedColumn[Any, Tile] =
+      new Column(MaskByValues(dataTile.expr, IsIn(maskTile, maskValues).expr)).as[Tile]
+  }
 }

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

@@ -24,12 +24,14 @@ package org.locationtech.rasterframes.util
 import geotrellis.raster.render.ColorRamps
 import org.apache.spark.sql.Dataset
 import org.apache.spark.sql.functions.{base64, concat, concat_ws, length, lit, substring, when}
+import org.apache.spark.sql.jts.JTSTypes
 import org.apache.spark.sql.types.{StringType, StructField}
 import org.locationtech.rasterframes.expressions.DynamicExtractors
 import org.locationtech.rasterframes.{rfConfig, rf_render_png, rf_resample}
+import org.apache.spark.sql.rf.WithTypeConformity
 
 /**
-  * DataFrame extensiosn for rendering sample content in a number of ways
+  * DataFrame extension for rendering sample content in a number of ways
   */
 trait DataFrameRenderers {
   private val truncateWidth = rfConfig.getInt("max-truncate-row-element-length")
@@ -47,8 +49,9 @@ trait DataFrameRenderers {
               lit("\"></img>")
             )
           else {
+            val isGeom = WithTypeConformity(c.dataType).conformsTo(JTSTypes.GeometryTypeInstance)
             val str = resolved.cast(StringType)
-            if (truncate)
+            if (truncate || isGeom)
               when(length(str) > lit(truncateWidth),
                 concat(substring(str, 1, truncateWidth), lit("..."))
               )

core/src/test/resources/MCD43A4.A2019111.h30v06.006.2019120033434_01.mrf.aux.xml

@@ -39,7 +39,7 @@ import scala.xml.parsing.XhtmlParser
 class ExtensionMethodSpec extends TestEnvironment with TestData with SubdivideSupport {
   lazy val rf = sampleTileLayerRDD.toLayer
 
-  describe("DataFrame exention methods") {
+  describe("DataFrame extension methods") {
     it("should maintain original type") {
       val df = rf.withPrefixedColumnNames("_foo_")
       "val rf2: RasterFrameLayer = df" should compile
@@ -49,7 +49,7 @@ class ExtensionMethodSpec extends TestEnvironment with TestData with SubdivideSu
       "val Some(col) = df.spatialKeyColumn" should compile
     }
   }
-  describe("RasterFrameLayer exention methods") {
+  describe("RasterFrameLayer extension methods") {
     it("should provide spatial key column") {
       noException should be thrownBy {
         rf.spatialKeyColumn
@@ -124,6 +124,10 @@ class ExtensionMethodSpec extends TestEnvironment with TestData with SubdivideSu
 
       val md3 = rf.toMarkdown(truncate=true, renderTiles = false)
       md3 shouldNot include("<img")
+
+      // Should truncate JTS types even when we don't ask for it.
+      val md4 = rf.withGeometry().select("geometry").toMarkdown(truncate = false)
+      md4 should include ("...")
     }
 
     it("should render HTML") {

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

@@ -28,7 +28,7 @@ import geotrellis.raster._
 import geotrellis.raster.render.ColorRamps
 import geotrellis.raster.testkit.RasterMatchers
 import javax.imageio.ImageIO
-import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.{Column, Encoders, TypedColumn}
 import org.apache.spark.sql.functions._
 import org.locationtech.rasterframes.expressions.accessors.ExtractTile
 import org.locationtech.rasterframes.model.TileDimensions
@@ -701,7 +701,7 @@ class RasterFunctionsSpec extends TestEnvironment with RasterMatchers {
       val withMasked = withMask.withColumn("masked",
         rf_inverse_mask_by_value($"tile", $"mask", mask_value))
         .withColumn("masked2", rf_mask_by_value($"tile", $"mask", lit(mask_value), true))
-
+      withMasked.explain(true)
       val result = withMasked.agg(rf_agg_no_data_cells($"tile") < rf_agg_no_data_cells($"masked")).as[Boolean]
 
       result.first() should be(true)
@@ -712,18 +712,27 @@ class RasterFunctionsSpec extends TestEnvironment with RasterMatchers {
       checkDocs("rf_inverse_mask_by_value")
     }
 
-    it("should mask tile by another identified by specified values") {
-      val squareIncrementingPRT = ProjectedRasterTile(squareIncrementingTile(randPRT.rows), randPRT.extent, randPRT.crs)
-      val df = Seq((randPRT, squareIncrementingPRT))
+    it("should mask tile by another identified by sequence of specified values") {
+      val squareIncrementingPRT = ProjectedRasterTile(squareIncrementingTile(six.rows), six.extent, six.crs)
+      val df = Seq((six, squareIncrementingPRT))
                 .toDF("tile", "mask")
+
       val mask_values = Seq(4, 5, 6, 12)
 
       val withMasked = df.withColumn("masked",
         rf_mask_by_values($"tile", $"mask", mask_values))
 
-      val result = withMasked.agg(rf_agg_no_data_cells($"masked") as "nd").as[Long]
+      val expected = squareIncrementingPRT.toArray().count(v ⇒ mask_values.contains(v))
+
+      val result = withMasked.agg(rf_agg_no_data_cells($"masked") as "masked_nd")
+          .first()
+
+      result.getAs[BigInt](0) should be (expected)
+
+      val withMaskedSql = df.selectExpr("rf_mask_by_values(tile, mask, array(4, 5, 6, 12)) AS masked")
+      val resultSql = withMaskedSql.agg(rf_agg_no_data_cells($"masked")).as[Long]
+      resultSql.first() should be (expected)
 
-      result.first() should be(mask_values.length)
       checkDocs("rf_mask_by_values")
     }
 
@@ -1018,7 +1027,6 @@ class RasterFunctionsSpec extends TestEnvironment with RasterMatchers {
     val e0Result = df.select($"in_expect_0").as[Tile].first()
     e0Result.toArray() should contain only (0)
 
-//    lazy val invalid = df.select(rf_local_is_in($"t", lit("foobar"))).as[Tile].first()
   }
 
   it("should unpack QA bits"){