Reworked rf_tile to accept RasterRef as argument.

Wrote unit tests for rf_tile.

Author: metasim

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

@@ -35,7 +35,7 @@ import org.locationtech.rasterframes.tiles.ProjectedRasterTile
 
 private[rasterframes]
 object DynamicExtractors {
-  /** Partial function for pulling a tile and its contesxt from an input row. */
+  /** Partial function for pulling a tile and its context from an input row. */
   lazy val tileExtractor: PartialFunction[DataType, InternalRow => (Tile, Option[TileContext])] = {
     case _: TileUDT =>
       (row: InternalRow) =>
@@ -47,6 +47,14 @@ object DynamicExtractors {
       }
   }
 
+  lazy val rasterRefExtractor: PartialFunction[DataType, InternalRow => RasterRef] = {
+    case t if t.conformsTo[RasterRef] =>
+      (row: InternalRow) => row.to[RasterRef]
+  }
+
+  lazy val tileableExtractor: PartialFunction[DataType, InternalRow => Tile] =
+    tileExtractor.andThen(_.andThen(_._1)).orElse(rasterRefExtractor.andThen(_.andThen(_.tile)))
+
   lazy val rowTileExtractor: PartialFunction[DataType, Row => (Tile, Option[TileContext])] = {
     case _: TileUDT =>
       (row: Row) =>  (row.to[Tile](TileUDT.tileSerializer), None)

core/src/main/scala/org/locationtech/rasterframes/expressions/accessors/RealizeTile.scala

@@ -22,15 +22,17 @@
 package org.locationtech.rasterframes.expressions.accessors
 
 import geotrellis.raster.Tile
+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}
+import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionDescription, UnaryExpression}
 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._
 import org.locationtech.rasterframes.encoders.CatalystSerializer._
-import org.locationtech.rasterframes.expressions.UnaryRasterOp
-import org.locationtech.rasterframes.model.TileContext
+import org.locationtech.rasterframes.expressions.DynamicExtractors._
+import org.locationtech.rasterframes.expressions._
 
 @ExpressionDescription(
   usage = "_FUNC_(raster) - Extracts the Tile component of a RasterSource, ProjectedRasterTile (or Tile) and ensures the cells are fully fetched.",
@@ -39,14 +41,22 @@ import org.locationtech.rasterframes.model.TileContext
       > SELECT _FUNC_(raster);
          ....
   """)
-case class RealizeTile(child: Expression) extends UnaryRasterOp with CodegenFallback {
+case class RealizeTile(child: Expression) extends UnaryExpression with CodegenFallback {
   override def dataType: DataType = TileType
 
   override def nodeName: String = "rf_tile"
-  implicit val tileSer = TileUDT.tileSerializer
 
-  override protected def eval(tile: Tile, ctx: Option[TileContext]): Any =
+  override def checkInputDataTypes(): TypeCheckResult = {
+    if (!tileableExtractor.isDefinedAt(child.dataType)) {
+      TypeCheckFailure(s"Input type '${child.dataType}' does not conform to a tiled raster type.")
+    } else TypeCheckSuccess
+  }
+  implicit val tileSer = TileUDT.tileSerializer
+  override protected def nullSafeEval(input: Any): Any = {
+    val in = row(input)
+    val tile = tileableExtractor(child.dataType)(in)
     (tile.toArrayTile(): Tile).toInternalRow
+  }
 }
 
 object RealizeTile {

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

@@ -23,53 +23,22 @@ package org.locationtech.rasterframes
 
 import java.io.ByteArrayInputStream
 
-import geotrellis.proj4.LatLng
 import geotrellis.raster
 import geotrellis.raster._
 import geotrellis.raster.render.ColorRamps
 import geotrellis.raster.testkit.RasterMatchers
-import geotrellis.vector.Extent
 import javax.imageio.ImageIO
 import org.apache.spark.sql.Encoders
 import org.apache.spark.sql.functions._
 import org.locationtech.rasterframes.expressions.accessors.ExtractTile
 import org.locationtech.rasterframes.model.TileDimensions
-import org.locationtech.rasterframes.ref.{RasterRef, RasterSource}
 import org.locationtech.rasterframes.stats._
 import org.locationtech.rasterframes.tiles.ProjectedRasterTile
 
 class RasterFunctionsSpec extends TestEnvironment with RasterMatchers {
+  import TestData._
   import spark.implicits._
 
-  val extent = Extent(10, 20, 30, 40)
-  val crs = LatLng
-  val ct = ByteUserDefinedNoDataCellType(-2)
-  val cols = 10
-  val rows = cols
-  val tileSize = cols * rows
-  val tileCount = 10
-  val numND = 4
-  lazy val zero = TestData.projectedRasterTile(cols, rows, 0, extent, crs, ct)
-  lazy val one = TestData.projectedRasterTile(cols, rows, 1, extent, crs, ct)
-  lazy val two = TestData.projectedRasterTile(cols, rows, 2, extent, crs, ct)
-  lazy val three = TestData.projectedRasterTile(cols, rows, 3, extent, crs, ct)
-  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: 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)
-  lazy val randPositiveDoubleTile = TestData.projectedRasterTile(cols, rows, scala.util.Random.nextDouble() + 1e-6, extent, crs, DoubleConstantNoDataCellType)
-
-  val expectedRandNoData: Long = numND * tileCount.toLong
-  val expectedRandData: Long = cols * rows * tileCount - expectedRandNoData
-  lazy val randNDTilesWithNull = Seq.fill[Tile](tileCount)(TestData.injectND(numND)(
-    TestData.randomTile(cols, rows, UByteConstantNoDataCellType)
-  )).map(ProjectedRasterTile(_, extent, crs)) :+ null
-
-  def lazyPRT = RasterRef(RasterSource(TestData.l8samplePath), 0, None, None).tile
-
   implicit val pairEnc = Encoders.tuple(ProjectedRasterTile.prtEncoder, ProjectedRasterTile.prtEncoder)
   implicit val tripEnc = Encoders.tuple(ProjectedRasterTile.prtEncoder, ProjectedRasterTile.prtEncoder, ProjectedRasterTile.prtEncoder)
 

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

@@ -38,8 +38,8 @@ import org.apache.spark.SparkContext
 import org.apache.spark.sql.SparkSession
 import org.locationtech.jts.geom.{Coordinate, GeometryFactory}
 import org.locationtech.rasterframes.expressions.tilestats.NoDataCells
+import org.locationtech.rasterframes.ref.{RasterRef, RasterSource}
 import org.locationtech.rasterframes.tiles.ProjectedRasterTile
-import spray.json.JsObject
 
 import scala.reflect.ClassTag
 
@@ -49,8 +49,15 @@ import scala.reflect.ClassTag
  * @since 4/3/17
  */
 trait TestData {
+  val extent = Extent(10, 20, 30, 40)
+  val crs = LatLng
+  val ct = ByteUserDefinedNoDataCellType(-2)
+  val cols = 10
+  val rows = cols
+  val tileSize = cols * rows
+  val tileCount = 10
+  val numND = 4
   val instant = ZonedDateTime.now()
-  val extent = Extent(1, 2, 3, 4)
   val sk = SpatialKey(37, 41)
   val stk = SpaceTimeKey(sk, instant)
   val pe = ProjectedExtent(extent, LatLng)
@@ -153,6 +160,29 @@ trait TestData {
   lazy val l8samplePath: URI = getClass.getResource("/L8-B1-Elkton-VA.tiff").toURI
   lazy val modisConvertedMrfPath: URI = getClass.getResource("/MCD43A4.A2019111.h30v06.006.2019120033434_01.mrf").toURI
 
+
+
+  lazy val zero = TestData.projectedRasterTile(cols, rows, 0, extent, crs, ct)
+  lazy val one = TestData.projectedRasterTile(cols, rows, 1, extent, crs, ct)
+  lazy val two = TestData.projectedRasterTile(cols, rows, 2, extent, crs, ct)
+  lazy val three = TestData.projectedRasterTile(cols, rows, 3, extent, crs, ct)
+  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: 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)
+  lazy val randPositiveDoubleTile = TestData.projectedRasterTile(cols, rows, scala.util.Random.nextDouble() + 1e-6, extent, crs, DoubleConstantNoDataCellType)
+
+  val expectedRandNoData: Long = numND * tileCount.toLong
+  val expectedRandData: Long = cols * rows * tileCount - expectedRandNoData
+  lazy val randNDTilesWithNull = Seq.fill[Tile](tileCount)(TestData.injectND(numND)(
+    TestData.randomTile(cols, rows, UByteConstantNoDataCellType)
+  )).map(ProjectedRasterTile(_, extent, crs)) :+ null
+
+  def lazyPRT = RasterRef(RasterSource(TestData.l8samplePath), 0, None, None).tile
+
   object GeomData {
     val fact = new GeometryFactory()
     val c1 = new Coordinate(1, 2)

core/src/test/scala/org/locationtech/rasterframes/encoders/CatalystSerializerSpec.scala

@@ -35,7 +35,9 @@ import org.locationtech.rasterframes.model.{CellContext, TileContext, TileDataCo
 import org.locationtech.rasterframes.ref.{RasterRef, RasterSource}
 import org.scalatest.Assertion
 
-class CatalystSerializerSpec extends TestEnvironment with TestData {
+class CatalystSerializerSpec extends TestEnvironment {
+  import TestData._
+
   val dc = TileDataContext(UShortUserDefinedNoDataCellType(3), TileDimensions(12, 23))
   val tc = TileContext(Extent(1, 2, 3, 4), WebMercator)
   val cc = CellContext(tc, dc, 34, 45)

core/src/test/scala/org/locationtech/rasterframes/ref/RasterRefSpec.scala

@@ -21,14 +21,14 @@
 
 package org.locationtech.rasterframes.ref
 
-import org.locationtech.rasterframes._
-import org.locationtech.rasterframes.expressions.accessors._
-import org.locationtech.rasterframes.expressions.generators._
-import RasterRef.RasterRefTile
-import geotrellis.raster.Tile
+import geotrellis.raster.{ByteConstantNoDataCellType, Tile}
 import geotrellis.vector.Extent
 import org.apache.spark.sql.Encoders
-import org.locationtech.rasterframes.TestEnvironment
+import org.locationtech.rasterframes.{TestEnvironment, _}
+import org.locationtech.rasterframes.expressions.accessors._
+import org.locationtech.rasterframes.expressions.generators._
+import org.locationtech.rasterframes.ref.RasterRef.RasterRefTile
+import org.locationtech.rasterframes.tiles.ProjectedRasterTile
 
 /**
  *
@@ -199,12 +199,45 @@ class RasterRefSpec extends TestEnvironment with TestData {
 
       refs.count() shouldBe > (1L)
 
-
       val dims = refs.select(rf_dimensions($"proj_raster")).distinct().collect()
       forEvery(dims) { r =>
         r.cols should be <= NOMINAL_TILE_SIZE
         r.rows should be <= NOMINAL_TILE_SIZE
       }
     }
   }
+
+  describe("RealizeTile") {
+    it("should pass through basic Tile") {
+      val t = TestData.randomTile(5, 5, ByteConstantNoDataCellType)
+      val result = Seq(t).toDF("tile").select(rf_tile($"tile")).first()
+      assertEqual(result, t)
+    }
+
+    it("should simplify ProjectedRasterTile") {
+      val t = TestData.randNDPRT
+      val result = Seq(t).toDF("tile").select(rf_tile($"tile")).first()
+      result.isInstanceOf[ProjectedRasterLike] should be (false)
+      assertEqual(result, t.toArrayTile())
+    }
+
+    it("should resolve a RasterRef") {
+      new Fixture {
+        import RasterRef.rrEncoder // This shouldn't be required, but product encoder gets choosen.
+        val r: RasterRef = subRaster
+        val result = Seq(r).toDF("ref").select(rf_tile($"ref")).first()
+        result.isInstanceOf[RasterRefTile] should be(false)
+        assertEqual(r.tile.toArrayTile(), result)
+      }
+    }
+
+    it("should resolve a RasterRefTile") {
+      new Fixture {
+        val t: ProjectedRasterTile = RasterRefTile(subRaster)
+        val result = Seq(t).toDF("tile").select(rf_tile($"tile")).first()
+        result.isInstanceOf[RasterRefTile] should be(false)
+        assertEqual(t.toArrayTile(), result)
+      }
+    }
+  }
 }

datasource/src/test/scala/org/locationtech/rasterframes/datasource/geotrellis/GeoTrellisDataSourceSpec.scala

@@ -51,8 +51,9 @@ import org.scalatest.{BeforeAndAfterAll, Inspectors}
 import scala.math.{max, min}
 
 class GeoTrellisDataSourceSpec
-    extends TestEnvironment with TestData with BeforeAndAfterAll with Inspectors
+    extends TestEnvironment with BeforeAndAfterAll with Inspectors
     with RasterMatchers with DataSourceOptions {
+  import TestData._
 
   val tileSize = 12
   lazy val layer = Layer(new File(outputLocalPath).toURI, LayerId("test-layer", 4))