Fixes to rf_mask_by_values to compose over IfIn and have SQL docs.

Removed inverse flag from rf_mask_by_values before committing to an approach with it.

Author: metasim

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

@@ -312,32 +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] =
-      Mask.MaskByValues(sourceTile, maskTile, maskValues, lit(inverse))
-
   /** 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/package.scala

@@ -125,10 +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.MaskByValues]("rf_mask_by_values")
     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

@@ -24,11 +24,11 @@ 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.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
@@ -41,47 +41,51 @@ 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
     }
@@ -173,24 +177,26 @@ object Mask {
   }
 
   @ExpressionDescription(
-    usage = "_FUNC_(data, mask, maskValues, inverse) - 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. If `inverse` is true, the cells in `mask` that are not in `maskValues` list become NODATA",
-    arguments =
-      """
-
+    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 =
-      """
-         > SELECT _FUNC_(data, mask, array(1, 2, 3), false)
-
-        """
+    examples = """
+  Examples:
+    > SELECT _FUNC_(data, mask, array(1, 2, 3))
+      ..."""
   )
-  case class MaskByValues(dataTile: Expression, maskTile: Expression, maskValues: Expression, inverse: Boolean)
-  extends Mask(dataTile, IsIn(maskTile, maskValues), inverse, false) {
+  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, inverse: Column): TypedColumn[Any, Tile] =
-      new Column(MaskByValues(dataTile.expr, maskTile.expr, maskValues.expr, inverse.expr)).as[Tile]
+    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
@@ -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,30 +712,28 @@ class RasterFunctionsSpec extends TestEnvironment with RasterMatchers {
       checkDocs("rf_inverse_mask_by_value")
     }
 
-    it("should mask tile by another identified by specified values") {
-      checkDocs("rf_mask_by_values")
-
+    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()
-        .filter(v ⇒ mask_values.contains(v))
-        .length
+      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), false) AS masked")
+      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") {

docs/src/main/paradox/reference.md

@@ -235,12 +235,10 @@ Generate a `tile` with the values from `data_tile`, with NoData in cells where t
 
 ### rf_mask_by_values
 
-    Tile rf_mask_by_values(Tile data_tile, Tile mask_tile, Array mask_values, bool inverse)
-    Tile rf_mask_by_values(Tile data_tile, Tile mask_tile, seq mask_values, bool inverse)
+    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`.
-
-`inverse` is a literal not a Column. If it is True, the `data_tile` cells are set to NoData where the `mask_tile` cells are __not__ in `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
 

pyrasterframes/src/main/python/pyrasterframes/rasterfunctions.py

@@ -474,18 +474,17 @@ def rf_mask_by_value(data_tile, mask_tile, mask_value, inverse=False):
     return Column(jfcn(_to_java_column(data_tile), _to_java_column(mask_tile), _to_java_column(mask_value), inverse))
 
 
-def rf_mask_by_values(data_tile, mask_tile, mask_values, inverse=False):
+def rf_mask_by_values(data_tile, mask_tile, mask_values):
     """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
     """
     from pyspark.sql.functions import array as sql_array
     if isinstance(mask_values, list):
         mask_values = sql_array([lit(v) for v in mask_values])
 
     jfcn = RFContext.active().lookup('rf_mask_by_values')
     col_args = [_to_java_column(c) for c in [data_tile, mask_tile, mask_values]]
-    return Column(jfcn(*col_args, inverse))
+    return Column(jfcn(*col_args))
 
 
 def rf_inverse_mask_by_value(data_tile, mask_tile, mask_value):

pyrasterframes/src/main/python/tests/RasterFunctionsTests.py

@@ -249,21 +249,13 @@ def test_mask_by_values(self):
         tile = Tile(np.random.randint(1, 100, (5, 5)), CellType.uint8())
         mask_tile = Tile(np.array(range(1, 26), 'uint8').reshape(5, 5))
         expected_diag_nd = Tile(np.ma.masked_array(tile.cells, mask=np.eye(5)))
-        expected_off_diag_nd = Tile(np.ma.masked_array(tile.cells, mask=1 - np.eye(5)))
 
         df = self.spark.createDataFrame([Row(t=tile, m=mask_tile)]) \
             .select(rf_mask_by_values('t', 'm', [0, 6, 12, 18, 24]))  # values on the diagonal
         result0 = df.first()
         # assert_equal(result0[0].cells, expected_diag_nd)
         self.assertTrue(result0[0] == expected_diag_nd)
 
-        # mask values off the diagonal! (inverse=True)
-        result1 = self.spark.createDataFrame([Row(t=tile, m=mask_tile)]) \
-            .select(rf_mask_by_values('t', 'm', [0, 6, 12, 18, 24], True)) \
-            .first()
-        # assert_equal(result1[0].cells, expected_off_diag_nd)
-        self.assertTrue(result1[0] == expected_off_diag_nd)
-
     def test_mask(self):
         from pyspark.sql import Row
         from pyrasterframes.rf_types import Tile, CellType