Expand rf_mask_by_value[s] signatures to match python api, expand unit tests

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

Author: vpipkt

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

@@ -304,15 +304,43 @@ trait RasterFunctions {
     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: 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.
        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=false): TypedColumn[Any, Tile] =
+  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)
+
+  /** 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] = {
+    import org.apache.spark.sql.functions.array
+    val valuesCol: Column = array(maskValues.map(lit).toSeq: _*)
+    rf_mask_by_values(sourceTile, maskTile, valuesCol, 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.
+       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)
@@ -321,6 +349,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)(

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

@@ -694,19 +694,39 @@ 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))
 
       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 specified values") {
+      val squareIncrementingPRT = ProjectedRasterTile(squareIncrementingTile(randPRT.rows), randPRT.extent, randPRT.crs)
+      val df = Seq((randPRT, 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]
+
+      result.first() should be(mask_values.length)
+      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"))