Merge pull request #404 from s22s/feature/local_is_in

Mask by list of values

Author: metasim

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

@@ -292,12 +292,38 @@ trait RasterFunctions {
   }
 
   /** Where the rf_mask tile contains NODATA, replace values in the source tile with NODATA */
-  def rf_mask(sourceTile: Column, maskTile: Column): TypedColumn[Any, Tile] =
-    Mask.MaskByDefined(sourceTile, maskTile)
+  def rf_mask(sourceTile: Column, maskTile: Column): TypedColumn[Any, Tile] = rf_mask(sourceTile, maskTile, false)
+
+  /** Where the rf_mask tile contains NODATA, replace values in the source tile with NODATA */
+  def rf_mask(sourceTile: Column, maskTile: Column, inverse: Boolean=false): TypedColumn[Any, Tile] =
+    if(!inverse) Mask.MaskByDefined(sourceTile, maskTile)
+    else Mask.InverseMaskByDefined(sourceTile, maskTile)
+
+  /** Where the `maskTile` equals `maskValue`, replace values in the source tile with `NoData` */
+  def rf_mask_by_value(sourceTile: Column, maskTile: Column, maskValue: Column, inverse: Boolean=false): TypedColumn[Any, Tile] =
+    if (!inverse) Mask.MaskByValue(sourceTile, maskTile, maskValue)
+    else Mask.InverseMaskByValue(sourceTile, maskTile, maskValue)
 
   /** Where the `maskTile` equals `maskValue`, replace values in the source tile with `NoData` */
-  def rf_mask_by_value(sourceTile: Column, maskTile: Column, maskValue: Column): TypedColumn[Any, Tile] =
-    Mask.MaskByValue(sourceTile, maskTile, maskValue)
+  def rf_mask_by_value(sourceTile: Column, maskTile: Column, maskValue: Int, inverse: Boolean): TypedColumn[Any, Tile] =
+    rf_mask_by_value(sourceTile, maskTile, lit(maskValue), inverse)
+
+  /** Where the `maskTile` equals `maskValue`, replace values in the source tile with `NoData` */
+  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. */
+  def rf_mask_by_values(sourceTile: Column, maskTile: Column, maskValues: Column): TypedColumn[Any, Tile] =
+      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. */
+  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)
+  }
 
   /** 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] =
@@ -307,6 +333,10 @@ trait RasterFunctions {
   def rf_inverse_mask_by_value(sourceTile: Column, maskTile: Column, maskValue: Column): TypedColumn[Any, Tile] =
     Mask.InverseMaskByValue(sourceTile, maskTile, maskValue)
 
+  /** Where the `maskTile` does **not** equal `maskValue`, replace values in the source tile with `NoData` */
+  def rf_inverse_mask_by_value(sourceTile: Column, maskTile: Column, maskValue: Int): TypedColumn[Any, Tile] =
+    Mask.InverseMaskByValue(sourceTile, maskTile, lit(maskValue))
+
   /** Create a tile where cells in the grid defined by cols, rows, and bounds are filled with the given value. */
   def rf_rasterize(geometry: Column, bounds: Column, value: Column, cols: Int, rows: Int): TypedColumn[Any, Tile] =
     withTypedAlias("rf_rasterize", geometry)(
@@ -408,6 +438,9 @@ trait RasterFunctions {
   /** Test if each cell value is in provided array */
   def rf_local_is_in(tileCol: Column, arrayCol: Column) = IsIn(tileCol, arrayCol)
 
+  /** Test if each cell value is in provided array */
+  def rf_local_is_in(tileCol: Column, array: Array[Int]) = IsIn(tileCol, array)
+
   /** Return a tile with ones where the input is NoData, otherwise zero */
   def rf_local_no_data(tileCol: Column): Column = Undefined(tileCol)
 

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

@@ -85,4 +85,12 @@ case class IsIn(left: Expression, right: Expression) extends BinaryExpression wi
 object IsIn {
   def apply(left: Column, right: Column): Column =
     new Column(IsIn(left.expr, right.expr))
+
+  def apply(left: Column, right: Array[Int]): Column = {
+    import org.apache.spark.sql.functions.lit
+    import org.apache.spark.sql.functions.array
+    val arrayExpr = array(right.map(lit):_*).expr
+    new Column(IsIn(left.expr, arrayExpr))
+  }
+
 }

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

@@ -28,58 +28,64 @@ import geotrellis.raster.mapalgebra.local.{Defined, InverseMask => gtInverseMask
 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, Literal, TernaryExpression}
+import org.apache.spark.sql.catalyst.expressions.{AttributeSet, Expression, ExpressionDescription, Literal, TernaryExpression}
+import org.apache.spark.sql.functions.lit
 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.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 {
+  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/test/scala/org/locationtech/rasterframes/RasterFunctionsSpec.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
@@ -694,19 +694,48 @@ class RasterFunctionsSpec extends TestEnvironment with RasterMatchers {
         rf_local_multiply(rf_convert_cell_type(
           rf_local_greater($"tile", 50),
           "uint8"),
-          lit(mask_value)
+          mask_value
         )
       )
 
       val withMasked = withMask.withColumn("masked",
-        rf_inverse_mask_by_value($"tile", $"mask", lit(mask_value)))
-
+        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)
+
+      val result2 = withMasked.agg(rf_agg_no_data_cells($"tile") < rf_agg_no_data_cells($"masked2")).as[Boolean]
+      result2.first() should be(true)
+
       checkDocs("rf_inverse_mask_by_value")
     }
 
+    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 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)
+
+      checkDocs("rf_mask_by_values")
+    }
+
     it("should render ascii art") {
       val df = Seq[Tile](ProjectedRasterTile(TestData.l8Labels)).toDF("tile")
       val r1 = df.select(rf_render_ascii($"tile"))
@@ -983,6 +1012,7 @@ class RasterFunctionsSpec extends TestEnvironment with RasterMatchers {
       .withColumn("ten", lit(10))
       .withColumn("in_expect_2", rf_local_is_in($"t", array($"one", $"five")))
       .withColumn("in_expect_1", rf_local_is_in($"t", array($"ten", $"five")))
+      .withColumn("in_expect_1a", rf_local_is_in($"t", Array(10, 5)))
       .withColumn("in_expect_0", rf_local_is_in($"t", array($"ten")))
 
     val e2Result = df.select(rf_tile_sum($"in_expect_2")).as[Double].first()
@@ -991,9 +1021,11 @@ class RasterFunctionsSpec extends TestEnvironment with RasterMatchers {
     val e1Result = df.select(rf_tile_sum($"in_expect_1")).as[Double].first()
     e1Result should be (1.0)
 
+    val e1aResult = df.select(rf_tile_sum($"in_expect_1a")).as[Double].first()
+    e1aResult should be (1.0)
+
     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()
   }
 }

docs/src/main/paradox/reference.md

@@ -207,35 +207,51 @@ SQL implementation does not accept a cell_type argument. It returns a float64 ce
 
 ## Masking and NoData
 
-See @ref:[NoData handling](nodata-handling.md) for conceptual discussion of cell types and NoData.
+See the @ref:[masking](masking.md) page for conceptual discussion of masking operations.
 
 There are statistical functions of the count of data and NoData values per `tile` and aggregate over a `tile` column: @ref:[`rf_data_cells`](reference.md#rf-data-cells), @ref:[`rf_no_data_cells`](reference.md#rf-no-data-cells), @ref:[`rf_agg_data_cells`](reference.md#rf-agg-data-cells), and @ref:[`rf_agg_no_data_cells`](reference.md#rf-agg-no-data-cells).
 
 Masking is a raster operation that sets specific cells to NoData based on the values in another raster.
 
 ### rf_mask
 
-    Tile rf_mask(Tile tile, Tile mask)
+    Tile rf_mask(Tile tile, Tile mask, bool inverse)
 
 Where the `mask` contains NoData, replace values in the `tile` with NoData.
 
 Returned `tile` cell type will be coerced to one supporting NoData if it does not already.
 
+`inverse` is a literal not a Column. If `inverse` is true, return the `tile` with NoData in locations where the `mask` _does not_ contain NoData. Equivalent to @ref:[`rf_inverse_mask`](reference.md#rf-inverse-mask).
+
 See also @ref:[`rf_rasterize`](reference.md#rf-rasterize).
 
+### rf_mask_by_value
+
+    Tile rf_mask_by_value(Tile data_tile, Tile mask_tile, Int mask_value, bool inverse)
+
+Generate a `tile` with the values from `data_tile`, with NoData in cells where the `mask_tile` is equal to `mask_value`.
+
+`inverse` is a literal not a Column. If `inverse` is true, return the `data_tile` with NoData in locations where the `mask_tile` value is _not equal_ to `mask_value`. Equivalent to @ref:[`rf_inverse_mask_by_value`](reference.md#rf-inverse-mask-by-value).
+
+### rf_mask_by_values
+
+    Tile rf_mask_by_values(Tile data_tile, Tile mask_tile, Array mask_values)
+    Tile rf_mask_by_values(Tile data_tile, Tile mask_tile, seq mask_values)
+
+Generate a `tile` with the values from `data_tile`, with NoData in cells where the `mask_tile` is in the `mask_values` Array or list. `mask_values` can be a [`pyspark.sql.ArrayType`][Array] or a `list`.  
 
 ### rf_inverse_mask
 
     Tile rf_inverse_mask(Tile tile, Tile mask)
 
 Where the `mask` _does not_ contain NoData, replace values in `tile` with NoData.
 
-### rf_mask_by_value
 
-    Tile rf_mask_by_value(Tile data_tile, Tile mask_tile, Int mask_value)
+### rf_inverse_mask_by_value
 
-Generate a `tile` with the values from `data_tile`, with NoData in cells where the `mask_tile` is equal to `mask_value`.
+    Tile rf_inverse_mask_by_value(Tile data_tile, Tile mask_tile, Int mask_value)
 
+Generate a `tile` with the values from `data_tile`, with NoData in cells where the `mask_tile` is not equal to `mask_value`. In other words, only keep `data_tile` cells in locations where the `mask_tile` is equal to `mask_value`.
 
 ### rf_is_no_data_tile
 

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

@@ -5,6 +5,7 @@
 ### 0.8.4
 
 * Upgraded to Spark 2.4.4
+* Add `rf_mask_by_values` and `rf_local_is_in` raster functions; added optional `inverse` argument to `rf_mask` functions. ([#403](https://github.com/locationtech/rasterframes/pull/403), [#384](https://github.com/locationtech/rasterframes/issues/384))
 * Added forced truncation of WKT types in Markdown/HTML rendering. ([#408](https://github.com/locationtech/rasterframes/pull/408))
 * Add `rf_local_is_in` raster function. ([#400](https://github.com/locationtech/rasterframes/pull/400))
 * Added partitioning to catalogs before processing in RasterSourceDataSource ([#397](https://github.com/locationtech/rasterframes/pull/397))