Switched Explode tiles to use UnsafeRow for slight improvement on memory pressure.

Reworked TileExplodeBench

Author: metasim

bench/src/main/scala/org/locationtech/rasterframes/bench/TileExplodeBench.scala

@@ -22,11 +22,12 @@ package org.locationtech.rasterframes.bench
 
 import java.util.concurrent.TimeUnit
 
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.BoundReference
+import org.apache.spark.sql.rf.TileUDT
 import org.locationtech.rasterframes._
-import org.apache.spark.sql._
-import org.apache.spark.sql.functions._
+import org.locationtech.rasterframes.expressions.generators.ExplodeTiles
 import org.openjdk.jmh.annotations._
-
 /**
  *
  * @author sfitch
@@ -36,33 +37,32 @@ import org.openjdk.jmh.annotations._
 @State(Scope.Benchmark)
 @OutputTimeUnit(TimeUnit.MILLISECONDS)
 class TileExplodeBench extends SparkEnv {
-  import spark.implicits._
 
-  @Param(Array("uint8", "uint16ud255", "float32", "float64"))
+  //@Param(Array("uint8", "uint16ud255", "float32", "float64"))
+  @Param(Array("uint16ud255"))
   var cellTypeName: String = _
 
   @Param(Array("256"))
   var tileSize: Int = _
 
-  @Param(Array("100"))
+  @Param(Array("2000"))
   var numTiles: Int = _
 
   @transient
-  var tiles: DataFrame = _
+  var tiles: Array[InternalRow] = _
+
+  var exploder: ExplodeTiles = _
 
   @Setup(Level.Trial)
   def setupData(): Unit = {
-    tiles = Seq.fill(numTiles)(randomTile(tileSize, tileSize, cellTypeName))
-      .toDF("tile").repartition(10)
-  }
-
-  @Benchmark
-  def arrayExplode() = {
-    tiles.select(posexplode(rf_tile_to_array_double($"tile"))).count()
+    tiles = Array.fill(numTiles)(randomTile(tileSize, tileSize, cellTypeName))
+        .map(t => InternalRow(TileUDT.tileSerializer.toInternalRow(t)))
+    val expr = BoundReference(0, TileType, true)
+    exploder = new ExplodeTiles(1.0, None, Seq(expr))
   }
-
   @Benchmark
   def tileExplode() = {
-    tiles.select(rf_explode_tiles($"tile")).count()
+    for(t <- tiles)
+      exploder.eval(t)
   }
 }

core/src/main/scala/org/locationtech/rasterframes/expressions/generators/ExplodeTiles.scala

@@ -24,8 +24,8 @@ package org.locationtech.rasterframes.expressions.generators
 import geotrellis.raster._
 import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
-import org.apache.spark.sql.catalyst.expressions.{Expression, Generator, GenericInternalRow}
+import org.apache.spark.sql.catalyst.expressions.codegen.{BufferHolder, CodegenFallback, UnsafeRowWriter}
+import org.apache.spark.sql.catalyst.expressions.{Expression, Generator, GenericInternalRow, UnsafeRow}
 import org.apache.spark.sql.types._
 import org.locationtech.rasterframes._
 import org.locationtech.rasterframes.expressions.DynamicExtractors
@@ -87,14 +87,17 @@ case class ExplodeTiles(
       cfor(0)(_ < rows, _ + 1) { row =>
         cfor(0)(_ < cols, _ + 1) { col =>
           val rowIndex = row * cols + col
-          val outCols = Array.ofDim[Any](numOutCols)
-          outCols(0) = col
-          outCols(1) = row
+          val outRow = new UnsafeRow(numOutCols)
+          val buffer = new BufferHolder(outRow)
+          val writer = new UnsafeRowWriter(buffer, numOutCols)
+          writer.write(0, col)
+          writer.write(1, row)
           cfor(0)(_ < tiles.length, _ + 1) { index =>
             val tile = tiles(index)
-            outCols(index + 2) = if(tile == null) doubleNODATA else tile.getDouble(col, row)
+            val cell: Double = if (tile == null) doubleNODATA else tile.getDouble(col, row)
+            writer.write(index + 2, cell)
           }
-          retval(rowIndex) = new GenericInternalRow(outCols)
+          retval(rowIndex) = outRow
         }
       }
       if(sampleFraction > 0.0 && sampleFraction < 1.0) sample(retval)

project/plugins.sbt

@@ -8,7 +8,7 @@ addSbtPlugin("com.typesafe.sbt" % "sbt-ghpages" % "0.6.2")
 addSbtPlugin("com.typesafe.sbt" % "sbt-site" % "1.3.2")
 addSbtPlugin("com.lightbend.paradox" % "sbt-paradox" % "0.5.5")
 addSbtPlugin("io.github.jonas" % "sbt-paradox-material-theme" % "0.6.0")
-addSbtPlugin("pl.project13.scala" % "sbt-jmh" % "0.3.6")
+addSbtPlugin("pl.project13.scala" % "sbt-jmh" % "0.3.3")
 addSbtPlugin("org.xerial.sbt" % "sbt-sonatype" % "2.1")
 addSbtPlugin("com.jsuereth" % "sbt-pgp" % "1.1.1")
 addSbtPlugin("com.eed3si9n" % "sbt-unidoc" % "0.4.1")