Merge pull request #415 from s22s/feature/mask_by_bits

Masking by bit values

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)
@@ -337,6 +337,52 @@ trait RasterFunctions {
   def rf_inverse_mask_by_value(sourceTile: Column, maskTile: Column, maskValue: Int): TypedColumn[Any, Tile] =
     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): 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): 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): 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*): 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)
+  }
+
+  /** Extract value from specified bits of the cells' underlying binary data.
+    * `startBit` is the first bit to consider, working from the right. It is zero indexed.
+    * `numBits` is the number of bits to take moving further to the left. */
+  def rf_local_extract_bits(tile: Column, startBit: Column, numBits: Column): Column =
+    ExtractBits(tile, startBit, numBits)
+
+  /** Extract value from specified bits of the cells' underlying binary data.
+    * `bitPosition` is bit to consider, working from the right. It is zero indexed. */
+  def rf_local_extract_bits(tile: Column, bitPosition: Column): Column =
+    rf_local_extract_bits(tile, bitPosition, lit(1))
+
+  /** Extract value from specified bits of the cells' underlying binary data.
+    * `startBit` is the first bit to consider, working from the right. It is zero indexed.
+    * `numBits` is the number of bits to take, moving further to the left. */
+  def rf_local_extract_bits(tile: Column, startBit: Int, numBits: Int): Column =
+    rf_local_extract_bits(tile, lit(startBit), lit(numBits))
+
+  /** Extract value from specified bits of the cells' underlying binary data.
+    * `bitPosition` is bit to consider, working from the right. It is zero indexed. */
+  def rf_local_extract_bits(tile: Column, bitPosition: Int): Column =
+    rf_local_extract_bits(tile, lit(bitPosition))
+
   /** 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/main/scala/org/locationtech/rasterframes/expressions/package.scala

@@ -138,5 +138,8 @@ package object expressions {
     registry.registerExpression[XZ2Indexer]("rf_spatial_index")
 
     registry.registerExpression[transformers.ReprojectGeometry]("st_reproject")
+
+    registry.registerExpression[ExtractBits]("rf_local_extract_bits")
+    registry.registerExpression[ExtractBits]("rf_local_extract_bit")
   }
 }

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

@@ -0,0 +1,96 @@
+/*
+ * This software is licensed under the Apache 2 license, quoted below.
+ *
+ * Copyright 2019 Astraea, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ *     [http://www.apache.org/licenses/LICENSE-2.0]
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package org.locationtech.rasterframes.expressions.transformers
+
+import geotrellis.raster.Tile
+import org.apache.spark.sql.{Column, TypedColumn}
+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, TernaryExpression}
+import org.apache.spark.sql.rf.TileUDT
+import org.apache.spark.sql.types.DataType
+import org.locationtech.rasterframes.encoders.CatalystSerializer._
+import org.locationtech.rasterframes.expressions.DynamicExtractors._
+import org.locationtech.rasterframes.expressions._
+
+@ExpressionDescription(
+  usage = "_FUNC_(tile, start_bit, num_bits) - In each cell of `tile`, extract `num_bits` from the cell value, starting at `start_bit` from the left.",
+  arguments = """
+  Arguments:
+    * tile - tile column to extract values
+    * start_bit -
+    * num_bits -
+  """,
+  examples = """
+  Examples:
+    > SELECT  _FUNC_(tile, lit(4), lit(2))
+       ..."""
+)
+case class ExtractBits(child1: Expression, child2: Expression, child3: Expression) extends TernaryExpression with CodegenFallback with Serializable {
+  override val nodeName: String = "rf_local_extract_bits"
+
+  override def children: Seq[Expression] = Seq(child1, child2, child3)
+
+  override def dataType: DataType = child1.dataType
+
+  override def checkInputDataTypes(): TypeCheckResult =
+    if(!tileExtractor.isDefinedAt(child1.dataType)) {
+      TypeCheckFailure(s"Input type '${child1.dataType}' does not conform to a raster type.")
+    } else if (!intArgExtractor.isDefinedAt(child2.dataType)) {
+      TypeCheckFailure(s"Input type '${child2.dataType}' isn't an integral type.")
+    } else if (!intArgExtractor.isDefinedAt(child3.dataType)) {
+      TypeCheckFailure(s"Input type '${child3.dataType}' isn't an integral type.")
+    } else TypeCheckSuccess
+
+
+  override protected def nullSafeEval(input1: Any, input2: Any, input3: Any): Any = {
+    implicit val tileSer = TileUDT.tileSerializer
+    val (childTile, childCtx) = tileExtractor(child1.dataType)(row(input1))
+
+    val startBits = intArgExtractor(child2.dataType)(input2).value
+
+    val numBits = intArgExtractor(child2.dataType)(input3).value
+
+    childCtx match {
+      case Some(ctx) => ctx.toProjectRasterTile(op(childTile, startBits, numBits)).toInternalRow
+      case None => op(childTile, startBits, numBits).toInternalRow
+    }
+  }
+
+  protected def op(tile: Tile, startBit: Int, numBits: Int): Tile = ExtractBits(tile, startBit, numBits)
+
+}
+
+object ExtractBits{
+  def apply(tile: Column, startBit: Column, numBits: Column): Column =
+    new Column(ExtractBits(tile.expr, startBit.expr, numBits.expr))
+
+  def apply(tile: Tile, startBit: Int, numBits: Int): Tile = {
+    assert(!tile.cellType.isFloatingPoint, "ExtractBits operation requires integral CellType")
+    // this is the last `numBits` positions of "111111111111111"
+    val widthMask = Int.MaxValue >> (63 - numBits)
+    // map preserving the nodata structure
+    tile.mapIfSet(x ⇒ x >> startBit & widthMask)
+  }
+
+}

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

@@ -24,7 +24,7 @@ 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
@@ -38,7 +38,15 @@ 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 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, undefined: Boolean, inverse: Boolean)
   extends TernaryExpression with CodegenFallback with Serializable {
   // aliases.
   def targetExp = left
@@ -77,13 +85,16 @@ 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
+    // 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, 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 +117,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 +137,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 +157,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 +179,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 +201,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"