Merge branch 'develop' into feature/raster-spatial-index

Author: metasim

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

@@ -339,7 +339,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)
@@ -357,6 +357,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

@@ -139,5 +139,8 @@ package object expressions {
     registry.registerExpression[Z2Indexer]("rf_z2_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"

core/src/main/scala/org/locationtech/rasterframes/extensions/RasterJoin.scala

@@ -23,22 +23,37 @@ package org.locationtech.rasterframes.extensions
 import org.apache.spark.sql._
 import org.apache.spark.sql.functions._
 import org.locationtech.rasterframes._
+import org.locationtech.rasterframes.expressions.SpatialRelation
+import org.locationtech.rasterframes.expressions.accessors.ExtractTile
 import org.locationtech.rasterframes.functions.reproject_and_merge
 import org.locationtech.rasterframes.util._
 
 import scala.util.Random
 
 object RasterJoin {
 
+  /** Perform a raster join on dataframes that each have proj_raster columns, or crs and extent explicitly included. */
   def apply(left: DataFrame, right: DataFrame): DataFrame = {
-    val df = apply(left, right, left("extent"), left("crs"), right("extent"), right("crs"))
-    df.drop(right("extent")).drop(right("crs"))
+    def usePRT(d: DataFrame) =
+      d.projRasterColumns.headOption
+        .map(p => (rf_crs(p),  rf_extent(p)))
+        .orElse(Some(col("crs"), col("extent")))
+        .map { case (crs, extent) =>
+          val d2 = d.withColumn("crs", crs).withColumn("extent", extent)
+          (d2, d2("crs"), d2("extent"))
+        }
+        .get
+
+    val (ldf, lcrs, lextent) = usePRT(left)
+    val (rdf, rcrs, rextent) = usePRT(right)
+
+    apply(ldf, rdf, lextent, lcrs, rextent, rcrs)
   }
 
   def apply(left: DataFrame, right: DataFrame, leftExtent: Column, leftCRS: Column, rightExtent: Column, rightCRS: Column): DataFrame = {
     val leftGeom = st_geometry(leftExtent)
     val rightGeomReproj = st_reproject(st_geometry(rightExtent), rightCRS, leftCRS)
-    val joinExpr = st_intersects(leftGeom, rightGeomReproj)
+    val joinExpr = new Column(SpatialRelation.Intersects(leftGeom.expr, rightGeomReproj.expr))
     apply(left, right, joinExpr, leftExtent, leftCRS, rightExtent, rightCRS)
   }
 
@@ -65,7 +80,7 @@ object RasterJoin {
     val leftAggCols = left.columns.map(s => first(left(s), true) as s)
     // On the RHS we collect result as a list.
     val rightAggCtx = Seq(collect_list(rightExtent) as rightExtent2, collect_list(rightCRS) as rightCRS2)
-    val rightAggTiles = right.tileColumns.map(c => collect_list(c) as c.columnName)
+    val rightAggTiles = right.tileColumns.map(c => collect_list(ExtractTile(c)) as c.columnName)
     val rightAggOther = right.notTileColumns
       .filter(n => n.columnName != rightExtent.columnName && n.columnName != rightCRS.columnName)
       .map(c => collect_list(c) as (c.columnName + "_agg"))