Add expression and raster functions around Tile.interpretAs called rf_interpret_cell_type_as

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

Author: vpipkt

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

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

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")

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

@@ -0,0 +1,106 @@
+/*
+ * 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
+
+/**
+ * Change the CellType of a Tile
+ *
+ * @since 9/11/18
+ */
+
+@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

@@ -52,7 +52,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)
@@ -893,4 +893,23 @@ class RasterFunctionsSpec extends TestEnvironment with RasterMatchers {
 
     checkDocs("rf_resample")
   }
+  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)
+
+  }
 }

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

@@ -106,9 +106,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
 

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

@@ -72,6 +72,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"""

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

@@ -279,3 +279,18 @@ def test_cell_type_in_functions(self):
         self.assertEqual(counts["data2"], 0)
         self.assertEqual(counts["nodata2"], 3 * 4)
 
+
+    def test_rf_interpret_cell_type_as(self):
+        from pyspark.sql import Row
+        from pyrasterframes.rf_types import Tile
+        import numpy as np
+
+        df = self.spark.createDataFrame([
+            Row(t=Tile(np.array([[1, 3, 4], [5, 0, 3]]), CellType.uint8().with_no_data_value(5)))
+        ])
+        df = df.withColumn('tile', rf_interpret_cell_type_as('t', 'uint8ud3')) # threes become ND
+        result = df.select(rf_tile_sum(rf_local_equal('t', lit(3))).alias('threes')).first()['threes']
+        self.assertEqual(result, 2)
+
+        result_5 = df.select(rf_tile_sum(rf_local_equal('t', lit(5))).alias('fives')).first()['fives']
+        self.assertEqual(result_5, 0)