Fix for both masking by def and value; expand code comments; update tests

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

Author: vpipkt

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

@@ -24,13 +24,13 @@ 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.{Undefined, 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, NullType}
+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._
@@ -43,10 +43,10 @@ import org.slf4j.LoggerFactory
   * @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 undefined 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)
+abstract class Mask(val left: Expression, val middle: Expression, val right: Expression, undefined: Boolean, inverse: Boolean)
   extends TernaryExpression with CodegenFallback with Serializable {
   // aliases.
   def targetExp = left
@@ -85,9 +85,12 @@ abstract class Mask(val left: Expression, val middle: Expression, val right: Exp
 
     val maskValue = intArgExtractor(maskValueExp.dataType)(maskValueInput)
 
-    val masking = if (defined) Defined(maskTile)
-    else maskTile.localEqual(maskValue.value)
+    // Get a tile where values of 1 indicate locations to set to ND in the target tile
+    // When `undefined` is true, setting targetTile locations to ND for ND locations of the `maskTile`
+    val masking = if (undefined) Undefined(maskTile)
+    else maskTile.localEqual(maskValue.value) // Otherwise if `maskTile` locations equal `maskValue`, set location to ND
 
+    // apply the `masking` where values are 1 set to ND (possibly inverted!)
     val result = if (inverse)
       gtInverseMask(targetTile, masking, 1, raster.NODATA)
     else

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

@@ -642,6 +642,17 @@ class RasterFunctionsSpec extends TestEnvironment with RasterMatchers {
       checkDocs("rf_mask")
     }
 
+    it("should mask with expected results") {
+      val df = Seq((byteArrayTile, maskingTile)).toDF("tile", "mask")
+
+      val withMasked = df.withColumn("masked",
+        rf_mask($"tile", $"mask"))
+
+      val result: Tile = withMasked.select($"masked").as[Tile].first()
+
+      result.localUndefined().toArray() should be (maskingTile.localUndefined().toArray())
+    }
+
     it("should mask without mutating cell type") {
       val result = Seq((byteArrayTile, maskingTile))
         .toDF("tile", "mask")
@@ -1227,6 +1238,7 @@ class RasterFunctionsSpec extends TestEnvironment with RasterMatchers {
       df.createOrReplaceTempView("df_maskbits")
 
       val maskedCol = "cloud_conf_med"
+      // this is the example in the docs
       val result = spark.sql(
         s"""
           |SELECT rf_mask_by_values(

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

@@ -261,19 +261,23 @@ def test_mask(self):
         from pyrasterframes.rf_types import Tile, CellType
 
         np.random.seed(999)
-        ma = np.ma.array(np.random.randint(0, 10, (5, 5), dtype='int8'), mask=np.random.rand(5, 5) > 0.7)
+        # importantly exclude 0 from teh range because that's the nodata value for the `data_tile`'s cell type
+        ma = np.ma.array(np.random.randint(1, 10, (5, 5), dtype='int8'), mask=np.random.rand(5, 5) > 0.7)
         expected_data_values = ma.compressed().size
         expected_no_data_values = ma.size - expected_data_values
         self.assertTrue(expected_data_values > 0, "Make sure random seed is cooperative ")
         self.assertTrue(expected_no_data_values > 0, "Make sure random seed is cooperative ")
 
-        df = self.spark.createDataFrame([
-            Row(t=Tile(np.ones(ma.shape, ma.dtype)), m=Tile(ma))
-        ])
+        data_tile = Tile(np.ones(ma.shape, ma.dtype), CellType.uint8())
+
+        df = self.spark.createDataFrame([Row(t=data_tile, m=Tile(ma))]) \
+            .withColumn('masked_t', rf_mask('t', 'm'))
 
-        df = df.withColumn('masked_t', rf_mask('t', 'm'))
         result = df.select(rf_data_cells('masked_t')).first()[0]
-        self.assertEqual(result, expected_data_values)
+        self.assertEqual(result, expected_data_values,
+                         f"Masked tile should have {expected_data_values} data values but found: {df.select('masked_t').first()[0].cells}."
+                         f"Original data: {data_tile.cells}"
+                         f"Masked by {ma}")
 
         nd_result = df.select(rf_no_data_cells('masked_t')).first()[0]
         self.assertEqual(nd_result, expected_no_data_values)