Attempt to create Expression for MaskByValues to enable sql api

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

Author: vpipkt

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

@@ -316,10 +316,7 @@ trait RasterFunctions {
        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] =
-    if (!inverse)
-      Mask.MaskByValue(sourceTile, rf_local_is_in(maskTile, maskValues), lit(1))
-    else
-      Mask.MaskByValue(sourceTile, rf_local_is_in(maskTile, maskValues), lit(0))
+      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. */

core/src/main/scala/org/locationtech/rasterframes/expressions/package.scala

@@ -126,6 +126,7 @@ package object expressions {
 
     registry.registerExpression[Mask.MaskByDefined]("rf_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")
 

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

@@ -24,16 +24,18 @@ 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.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
 
@@ -169,4 +171,26 @@ 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, 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 =
+      """
+
+        """,
+    examples =
+      """
+         > SELECT _FUNC_(data, mask, array(1, 2, 3), false)
+
+        """
+  )
+  case class MaskByValues(dataTile: Expression, maskTile: Expression, maskValues: Expression, inverse: Boolean)
+  extends Mask(dataTile, IsIn(maskTile, maskValues), inverse, false) {
+    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]
+  }
+
 }

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

@@ -713,6 +713,8 @@ class RasterFunctionsSpec extends TestEnvironment with RasterMatchers {
     }
 
     it("should mask tile by another identified by specified values") {
+      checkDocs("rf_mask_by_values")
+
       val squareIncrementingPRT = ProjectedRasterTile(squareIncrementingTile(six.rows), six.extent, six.crs)
       val df = Seq((six, squareIncrementingPRT))
                 .toDF("tile", "mask")