Masking improvements and unit tests.

Mask by value of 0 fixed
Masking cell type mutations in rf_mask_by_values is resolved
Masking by extraction of bits implemented

Signed-off-by: Jason T. Brown <[email protected]>

Author: vpipkt

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

@@ -319,7 +319,7 @@ trait RasterFunctions {
 
   /** 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] = {
+  def rf_mask_by_values(sourceTile: Column, maskTile: Column, maskValues: 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)
@@ -338,22 +338,24 @@ trait RasterFunctions {
     Mask.InverseMaskByValue(sourceTile, maskTile, lit(maskValue))
 
   /** Applies a mask using bit values in the `mask_tile`. Working from the right, extract the bit at `bitPosition` from the `maskTile`. In all locations where these are equal to the `valueToMask`, the returned tile is set to NoData, else the original `dataTile` cell value. */
-  def rf_mask_by_bit(dataTile: Column, maskTile: Column, bitPosition: Int, valueToMask: Boolean): Column =
-    rf_mask_by_bit(dataTile, maskTile, lit(bitPosition), lit(valueToMask))
+  def rf_mask_by_bit(dataTile: Column, maskTile: Column, bitPosition: Int, valueToMask: Boolean): TypedColumn[Any, Tile] =
+    rf_mask_by_bit(dataTile, maskTile, lit(bitPosition), lit(if (valueToMask) 1 else 0))
 
   /** Applies a mask using bit values in the `mask_tile`. Working from the right, extract the bit at `bitPosition` from the `maskTile`. In all locations where these are equal to the `valueToMask`, the returned tile is set to NoData, else the original `dataTile` cell value. */
-  def rf_mask_by_bit(dataTile: Column, maskTile: Column, bitPosition: Column, valueToMask: Column): Column =
-    rf_mask_by_bits(dataTile, maskTile, bitPosition, lit(1), valueToMask)
+  def rf_mask_by_bit(dataTile: Column, maskTile: Column, bitPosition: Column, valueToMask: Column): TypedColumn[Any, Tile] = {
+    import org.apache.spark.sql.functions.array
+    rf_mask_by_bits(dataTile, maskTile, bitPosition, lit(1), array(valueToMask))
+  }
 
   /** Applies a mask from blacklisted bit values in the `mask_tile`. Working from the right, the bits from `start_bit` to `start_bit + num_bits` are @ref:[extracted](reference.md#rf_local_extract_bits) from cell values of the `mask_tile`. In all locations where these are in the `mask_values`, the returned tile is set to NoData; otherwise the original `tile` cell value is returned. */
-  def rf_mask_by_bits(dataTile: Column, maskTile: Column, startBit: Column, numBits: Column, valuesToMask: Column): Column = {
+  def rf_mask_by_bits(dataTile: Column, maskTile: Column, startBit: Column, numBits: Column, valuesToMask: Column): TypedColumn[Any, Tile] = {
     val bitMask =  rf_local_extract_bits(maskTile, startBit, numBits)
     rf_mask_by_values(dataTile, bitMask, valuesToMask)
   }
 
 
   /** Applies a mask from blacklisted bit values in the `mask_tile`. Working from the right, the bits from `start_bit` to `start_bit + num_bits` are @ref:[extracted](reference.md#rf_local_extract_bits) from cell values of the `mask_tile`. In all locations where these are in the `mask_values`, the returned tile is set to NoData; otherwise the original `tile` cell value is returned. */
-  def rf_mask_by_bits(dataTile: Column, maskTile: Column, startBit: Int, numBits: Int, valuesToMask: Int*): Column = {
+  def rf_mask_by_bits(dataTile: Column, maskTile: Column, startBit: Int, numBits: Int, valuesToMask: Int*): TypedColumn[Any, Tile] = {
     import org.apache.spark.sql.functions.array
     val values = array(valuesToMask.map(lit):_*)
     rf_mask_by_bits(dataTile, maskTile, lit(startBit), lit(numBits), values)

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

@@ -24,21 +24,29 @@ package org.locationtech.rasterframes.expressions.transformers
 import com.typesafe.scalalogging.Logger
 import geotrellis.raster
 import geotrellis.raster.Tile
-import geotrellis.raster.mapalgebra.local.{Defined, InverseMask => gtInverseMask, Mask => gtMask}
+import geotrellis.raster.mapalgebra.local.{Defined, InverseMask ⇒ gtInverseMask, Mask ⇒ gtMask}
 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.rf.TileUDT
-import org.apache.spark.sql.types.DataType
+import org.apache.spark.sql.types.{DataType, NullType}
 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)
+/** Convert cells in the `left` to NoData based on another tile's contents
+  *
+  * @param left a tile of data values, with valid nodata cell type
+  * @param middle a tile indicating locations to set to nodata
+  * @param right optional, cell values in the `middle` tile indicating locations to set NoData
+  * @param defined if true, consider NoData in the `middle` as the locations to mask; else use `right` valued cells
+  * @param inverse if true, and defined is true, set `left` to NoData where `middle` is NOT nodata
+  */
+abstract class Mask(val left: Expression, val middle: Expression, val right: Expression, defined: Boolean, inverse: Boolean)
   extends TernaryExpression with CodegenFallback with Serializable {
   // aliases.
   def targetExp = left
@@ -77,13 +85,13 @@ abstract class Mask(val left: Expression, val middle: Expression, val right: Exp
 
     val maskValue = intArgExtractor(maskValueExp.dataType)(maskValueInput)
 
-    val masking = if (maskValue.value == 0) Defined(maskTile)
-    else maskTile
+    val masking = if (defined) Defined(maskTile)
+    else maskTile.localEqual(maskValue.value)
 
     val result = if (inverse)
-      gtInverseMask(targetTile, masking, maskValue.value, raster.NODATA)
+      gtInverseMask(targetTile, masking, 1, raster.NODATA)
     else
-      gtMask(targetTile, masking, maskValue.value, raster.NODATA)
+      gtMask(targetTile, masking, 1, raster.NODATA)
 
     targetCtx match {
       case Some(ctx) => ctx.toProjectRasterTile(result).toInternalRow
@@ -106,7 +114,7 @@ object Mask {
        ..."""
   )
   case class MaskByDefined(target: Expression, mask: Expression)
-    extends Mask(target, mask, Literal(0), false) {
+    extends Mask(target, mask, Literal(0), true, false) {
     override def nodeName: String = "rf_mask"
   }
   object MaskByDefined {
@@ -126,7 +134,7 @@ object Mask {
        ..."""
   )
   case class InverseMaskByDefined(leftTile: Expression, rightTile: Expression)
-    extends Mask(leftTile, rightTile, Literal(0), true) {
+    extends Mask(leftTile, rightTile, Literal(0), true, true) {
     override def nodeName: String = "rf_inverse_mask"
   }
   object InverseMaskByDefined {
@@ -146,7 +154,7 @@ object Mask {
        ..."""
   )
   case class MaskByValue(leftTile: Expression, rightTile: Expression, maskValue: Expression)
-    extends Mask(leftTile, rightTile, maskValue, false) {
+    extends Mask(leftTile, rightTile, maskValue, false, false) {
     override def nodeName: String = "rf_mask_by_value"
   }
   object MaskByValue {
@@ -168,7 +176,7 @@ object Mask {
        ..."""
   )
   case class InverseMaskByValue(leftTile: Expression, rightTile: Expression, maskValue: Expression)
-    extends Mask(leftTile, rightTile, maskValue, true) {
+    extends Mask(leftTile, rightTile, maskValue, false, true) {
     override def nodeName: String = "rf_inverse_mask_by_value"
   }
   object InverseMaskByValue {
@@ -190,7 +198,7 @@ object Mask {
       ..."""
   )
   case class MaskByValues(dataTile: Expression, maskTile: Expression)
-    extends Mask(dataTile, maskTile, Literal(1), inverse = false) {
+    extends Mask(dataTile, maskTile, Literal(1), false, false) {
     def this(dataTile: Expression, maskTile: Expression, maskValues: Expression) =
       this(dataTile, IsIn(maskTile, maskValues))
     override def nodeName: String = "rf_mask_by_values"