Additional raster functions: rf_interpret_cell_type_as, rf_local_data, rf_local_no_data (#309)

* Add expression and raster functions around Tile.interpretAs called rf_interpret_cell_type_as

* Add rf_local_data and rf_local_no_data #244

Signed-off-by: Jason T. Brown <jason@astraea.earth>

Author: vpipkt

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

@@ -100,6 +100,18 @@ trait RasterFunctions {
   /** Change the Tile's cell type */
   def rf_convert_cell_type(col: Column, cellTypeName: String): Column = SetCellType(col, cellTypeName)
 
+  /** Change the Tile's cell type */
+  def rf_convert_cell_type(col: Column, cellType: Column): Column = SetCellType(col, cellType)
+
+  /** Change the interpretation of the Tile's cell values according to specified CellType */
+  def rf_interpret_cell_type_as(col: Column, cellType: CellType): Column = InterpretAs(col, cellType)
+
+  /** Change the interpretation of the Tile's cell values according to specified CellType */
+  def rf_interpret_cell_type_as(col: Column, cellTypeName: String): Column = InterpretAs(col, cellTypeName)
+
+  /** Change the interpretation of the Tile's cell values according to specified CellType */
+  def rf_interpret_cell_type_as(col: Column, cellType: Column): Column = InterpretAs(col, cellType)
+
   /** Resample tile to different size based on scalar factor or tile whose dimension to match. Scalar less
     * than one will downsample tile; greater than one will upsample. Uses nearest-neighbor. */
   def rf_resample[T: Numeric](tileCol: Column, factorValue: T) = Resample(tileCol, factorValue)
@@ -377,6 +389,12 @@ trait RasterFunctions {
   /** Cellwise inequality comparison between a tile and a scalar. */
   def rf_local_unequal[T: Numeric](tileCol: Column, value: T): Column = Unequal(tileCol, value)
 
+  /** Return a tile with ones where the input is NoData, otherwise zero */
+  def rf_local_no_data(tileCol: Column): Column = Undefined(tileCol)
+
+  /** Return a tile with zeros where the input is NoData, otherwise one*/
+  def rf_local_data(tileCol: Column): Column = Defined(tileCol)
+
   /** Round cell values to nearest integer without chaning cell type. */
   def rf_round(tileCol: Column): Column = Round(tileCol)
 

core/src/main/scala/org/locationtech/rasterframes/expressions/localops/Defined.scala

@@ -0,0 +1,48 @@
+/*
+ * 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.localops
+
+import geotrellis.raster.Tile
+import org.apache.spark.sql.Column
+import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
+import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionDescription}
+import org.locationtech.rasterframes.expressions.{NullToValue, UnaryLocalRasterOp}
+
+@ExpressionDescription(
+  usage = "_FUNC_(tile) - Return a tile with zeros where the input is NoData, otherwise one.",
+  arguments = """
+  Arguments:
+    * tile - tile column to inspect""",
+  examples = """
+  Examples:
+    > SELECT  _FUNC_(tile);
+       ..."""
+)
+case class Defined(child: Expression) extends UnaryLocalRasterOp
+  with NullToValue with CodegenFallback {
+  override def nodeName: String = "rf_local_data"
+  override def na: Any = null
+  override protected def op(child: Tile): Tile = child.localDefined()
+}
+object Defined{
+  def apply(tile: Column): Column = new Column(Defined(tile.expr))
+}

core/src/main/scala/org/locationtech/rasterframes/expressions/localops/Undefined.scala

@@ -0,0 +1,48 @@
+/*
+ * 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.localops
+
+import geotrellis.raster.Tile
+import org.apache.spark.sql.Column
+import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
+import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionDescription}
+import org.locationtech.rasterframes.expressions.{NullToValue, UnaryLocalRasterOp}
+
+@ExpressionDescription(
+  usage = "_FUNC_(tile) - Return a tile with ones where the input is NoData, otherwise zero.",
+  arguments = """
+  Arguments:
+    * tile - tile column to inspect""",
+  examples = """
+  Examples:
+    > SELECT  _FUNC_(tile);
+       ..."""
+)
+case class Undefined(child: Expression) extends UnaryLocalRasterOp
+  with NullToValue with CodegenFallback {
+  override def nodeName: String = "rf_local_no_data"
+  override def na: Any = null
+  override protected def op(child: Tile): Tile = child.localUndefined()
+}
+object Undefined{
+  def apply(tile: Column): Column = new Column(Undefined(tile.expr))
+}

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

@@ -68,6 +68,7 @@ package object expressions {
     registry.registerExpression[ExplodeTiles]("rf_explode_tiles")
     registry.registerExpression[GetCellType]("rf_cell_type")
     registry.registerExpression[SetCellType]("rf_convert_cell_type")
+    registry.registerExpression[InterpretAs]("rf_interpret_cell_type_as")
     registry.registerExpression[SetNoDataValue]("rf_with_no_data")
     registry.registerExpression[GetDimensions]("rf_dimensions")
     registry.registerExpression[ExtentToGeometry]("st_geometry")
@@ -86,6 +87,8 @@ package object expressions {
     registry.registerExpression[GreaterEqual]("rf_local_greater_equal")
     registry.registerExpression[Equal]("rf_local_equal")
     registry.registerExpression[Unequal]("rf_local_unequal")
+    registry.registerExpression[Undefined]("rf_local_no_data")
+    registry.registerExpression[Defined]("rf_local_data")
     registry.registerExpression[Sum]("rf_tile_sum")
     registry.registerExpression[Round]("rf_round")
     registry.registerExpression[Abs]("rf_abs")

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

@@ -0,0 +1,100 @@
+/*
+ * 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.CellType
+import org.apache.spark.sql.Column
+import org.apache.spark.sql.catalyst.InternalRow
+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.{BinaryExpression, Expression, ExpressionDescription}
+import org.apache.spark.sql.functions.lit
+import org.apache.spark.sql.rf.{TileUDT}
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.types.UTF8String
+import org.locationtech.rasterframes.encoders.CatalystSerializer._
+import org.locationtech.rasterframes.expressions.DynamicExtractors.tileExtractor
+import org.locationtech.rasterframes.expressions.row
+
+@ExpressionDescription(
+  usage = "_FUNC_(tile, value) - Change the interpretation of the Tile's cell values according to specified CellType",
+  arguments = """
+  Arguments:
+    * tile - left-hand-side tile
+    * rhs  - a cell type definition""",
+  examples = """
+  Examples:
+    > SELECT _FUNC_(tile, 'int16ud0');
+       ..."""
+)
+case class InterpretAs(tile: Expression, cellType: Expression)
+  extends BinaryExpression with CodegenFallback {
+  def left = tile
+  def right = cellType
+  override def nodeName: String = "rf_interpret_cell_type_as"
+  override def dataType: DataType = left.dataType
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    if (!tileExtractor.isDefinedAt(left.dataType))
+      TypeCheckFailure(s"Input type '${left.dataType}' does not conform to a raster type.")
+    else
+      right.dataType match {
+        case StringType => TypeCheckSuccess
+        case t if t.conformsTo[CellType] => TypeCheckSuccess
+        case _ =>
+          TypeCheckFailure(s"Expected CellType but received '${right.dataType.simpleString}'")
+      }
+  }
+
+  private def toCellType(datum: Any): CellType = {
+    right.dataType match {
+      case StringType =>
+        val text = datum.asInstanceOf[UTF8String].toString
+        CellType.fromName(text)
+      case st if st.conformsTo[CellType] =>
+        row(datum).to[CellType]
+    }
+  }
+
+  override protected def nullSafeEval(tileInput: Any, ctInput: Any): InternalRow = {
+    implicit val tileSer = TileUDT.tileSerializer
+
+    val (tile, ctx) = tileExtractor(left.dataType)(row(tileInput))
+    val ct = toCellType(ctInput)
+    val result = tile.interpretAs(ct)
+
+    ctx match {
+      case Some(c) => c.toProjectRasterTile(result).toInternalRow
+      case None => result.toInternalRow
+    }
+  }
+}
+
+object InterpretAs{
+  def apply(tile: Column, cellType: CellType): Column =
+    new Column(new InterpretAs(tile.expr, lit(cellType.name).expr))
+  def apply(tile: Column, cellType: String): Column =
+    new Column(new InterpretAs(tile.expr, lit(cellType).expr))
+  def apply(tile: Column, cellType: Column): Column =
+    new Column(new InterpretAs(tile.expr, cellType.expr))
+}

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

@@ -101,4 +101,6 @@ object SetCellType {
     new Column(new SetCellType(tile.expr, lit(cellType.name).expr))
   def apply(tile: Column, cellType: String): Column =
     new Column(new SetCellType(tile.expr, lit(cellType).expr))
+  def apply(tile: Column, cellType: Column): Column =
+    new Column(new SetCellType(tile.expr, cellType.expr))
 }

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

@@ -56,7 +56,7 @@ class RasterFunctionsSpec extends TestEnvironment with RasterMatchers {
   lazy val six = ProjectedRasterTile(three * two, three.extent, three.crs)
   lazy val nd = TestData.projectedRasterTile(cols, rows, -2, extent, crs, ct)
   lazy val randPRT = TestData.projectedRasterTile(cols, rows, scala.util.Random.nextInt(), extent, crs, ct)
-  lazy val randNDPRT  = TestData.injectND(numND)(randPRT)
+  lazy val randNDPRT: Tile  = TestData.injectND(numND)(randPRT)
 
   lazy val randDoubleTile = TestData.projectedRasterTile(cols, rows, scala.util.Random.nextGaussian(), extent, crs, DoubleConstantNoDataCellType)
   lazy val randDoubleNDTile  = TestData.injectND(numND)(randDoubleTile)
@@ -948,4 +948,38 @@ class RasterFunctionsSpec extends TestEnvironment with RasterMatchers {
     image.getWidth should be(red.cols)
     image.getHeight should be(red.rows)
   }
+  it("should interpret cell values with a specified cell type") {
+    checkDocs("rf_interpret_cell_type_as")
+    val df = Seq(randNDPRT).toDF("t")
+      .withColumn("tile", rf_interpret_cell_type_as($"t", "int8raw"))
+    val resultTile = df.select("tile").as[Tile].first()
+
+    resultTile.cellType should be (CellType.fromName("int8raw"))
+    // should have same number of values that are -2 the old ND
+    val countOldNd = df.select(
+      rf_tile_sum(rf_local_equal($"tile", ct.noDataValue)),
+      rf_no_data_cells($"t")
+    ).first()
+    countOldNd._1 should be (countOldNd._2)
+
+    // should not have no data any more (raw type)
+    val countNewNd = df.select(rf_no_data_cells($"tile")).first()
+    countNewNd should be (0L)
+
+  }
+
+  it("should return local data and nodata"){
+    checkDocs("rf_local_data")
+    checkDocs("rf_local_no_data")
+
+    val df = Seq(randNDPRT).toDF("t")
+      .withColumn("ld", rf_local_data($"t"))
+      .withColumn("lnd", rf_local_no_data($"t"))
+
+    val ndResult = df.select($"lnd").as[Tile].first()
+    ndResult should be (randNDPRT.localUndefined())
+
+    val dResult = df.select($"ld").as[Tile].first()
+    dResult should be (randNDPRT.localDefined())
+  }
 }

docs/src/main/paradox/release-notes.md

@@ -4,10 +4,11 @@
 
 ### 0.8.1
 
+* Added `rf_local_no_data`, `rf_local_data` and `rf_interpret_cell_type_as` raster functions.
 * Added: `rf_rgb_composite` and `rf_render_png`.
 * Added `toMarkdown` and `toHTML` extension methods for `DataFrame`, and registered them with the IPython formatter system when `rf_ipython` is imported.
 * New documentation theme (thanks [@jonas](https://github.com/jonas)!).
-* Fixed: Removed false return type garauntee in cases where an `Expression` accepts either `Tile` or `ProjectedRasterTile` [(#295)](https://github.com/locationtech/rasterframes/issues/295)
+* Fixed: Removed false return type guarantee in cases where an `Expression` accepts either `Tile` or `ProjectedRasterTile` [(#295)](https://github.com/locationtech/rasterframes/issues/295)
 
 ### 0.8.0
 

pyrasterframes/src/main/python/docs/reference.pymd

@@ -115,9 +115,17 @@ Example: `SELECT rf_mk_crs('EPSG:4326')`
 
 ### rf_convert_cell_type
 
-    Tile rf_convert_cell_type(Tile tileCol, String cellType)
+    Tile rf_convert_cell_type(Tile tile_col, CellType cell_type)
+    Tile rf_convert_cell_type(Tile tile_col, String cell_type)
 
-Convert `tileCol` to a different cell type. Available cell types can be retrieved with the @ref:[rf_cell_types](reference.md#rf-cell-types) function.
+Convert `tile_col` to a different cell type. In Python you can pass a CellType object to `cell_type`.
+
+### rf_interpret_cell_type_as
+
+    Tile rf_interpret_cell_type_as(Tile tile_col, CellType cell_type)
+    Tile rf_interpret_cell_type_as(Tile tile_col, String cell_type)
+
+Change the interpretation of the `tile_col`'s cell values according to specified `cell_type`. In Python you can pass a CellType object to `cell_type`.
 
 ### rf_resample
 
@@ -224,6 +232,20 @@ Generate a `tile` with the values from `data_tile`, with NoData in cells where t
 
 Returns true if `tile` contains only NoData. By definition returns false if cell type does not support NoData. To count NoData cells or data cells, see @ref:[`rf_no_data_cells`](reference.md#rf-no-data-cells), @ref:[`rf_data_cells`](reference.md#rf-data-cells), @ref:[`rf_agg_no_data_cells`](reference.md#rf-agg-no-data-cells), @ref:[`rf_agg_data_cells`](reference.md#rf-agg-data-cells), @ref:[`rf_agg_local_no_data_cells`](reference.md#rf-agg-local-no-data-cells), and @ref:[`rf_agg_local_data_cells`](reference.md#rf-agg-local-data-cells). This function is distinguished from @ref:[`rf_for_all`](reference.md#rf-for-all), which tests that values are not NoData and nonzero.
 
+### rf_local_no_data
+
+    Tile rf_local_no_data(Tile tile)
+
+Returns a tile with values of 1 in each cell where the input tile contains NoData. Otherwise values are 0.
+
+### rf_local_data
+
+    Tile rf_local_no_data(Tile tile)
+
+Returns a tile with values of 0 in each cell where the input tile contains NoData. Otherwise values are 1.
+
+### rf_local_data
+
 ### rf_with_no_data
 
     Tile rf_with_no_data(Tile tile, Double no_data_value)

pyrasterframes/src/main/python/pyrasterframes/rasterfunctions.py

@@ -86,6 +86,11 @@ def rf_convert_cell_type(tile_col, cell_type):
     jfcn = RFContext.active().lookup('rf_convert_cell_type')
     return Column(jfcn(_to_java_column(tile_col), _parse_cell_type(cell_type)))
 
+def rf_interpret_cell_type_as(tile_col, cell_type):
+    """Change the interpretation of the tile_col's cell values according to specified cell_type"""
+    jfcn = RFContext.active().lookup('rf_interpret_cell_type_as')
+    return Column(jfcn(_to_java_column(tile_col), _parse_cell_type(cell_type)))
+
 
 def rf_make_constant_tile(scalar_value, num_cols, num_rows, cell_type=CellType.float64()):
     """Constructor for constant tile column"""
@@ -255,6 +260,13 @@ def rf_local_unequal_int(tile_col, scalar):
     """Return a Tile with values equal 1 if the cell is not equal to a scalar, otherwise 0"""
     return _apply_scalar_to_tile('rf_local_unequal_int', tile_col, scalar)
 
+def rf_local_no_data(tile_col):
+    """Return a tile with ones where the input is NoData, otherwise zero."""
+    return _apply_column_function('rf_local_no_data', tile_col)
+
+def rf_local_data(tile_col):
+    """Return a tile with zeros where the input is NoData, otherwise one."""
+    return _apply_column_function('rf_local_data', tile_col)
 
 def _apply_column_function(name, *args):
     jfcn = RFContext.active().lookup(name)