Reworked TestEnvironment to make it easier to override SparkContext settings.

Author: metasim

core/src/main/scala/org/locationtech/rasterframes/ref/RasterSource.scala

@@ -93,9 +93,12 @@ object RasterSource extends LazyLogging {
   val cacheTimeout: Duration = Duration.fromNanos(rfConfig.getDuration("raster-source-cache-timeout").toNanos)
 
   private val rsCache = Scaffeine()
+    .recordStats()
     .expireAfterAccess(RasterSource.cacheTimeout)
     .build[String, RasterSource]
 
+  def cacheStats = rsCache.stats()
+
   implicit def rsEncoder: ExpressionEncoder[RasterSource] = {
     RasterSourceUDT // Makes sure UDT is registered first
     ExpressionEncoder()

core/src/main/scala/org/locationtech/rasterframes/ref/SimpleRasterInfo.scala

@@ -76,7 +76,10 @@ object SimpleRasterInfo {
     )
   }
 
+  private[rasterframes]
   lazy val cache = Scaffeine()
-    //.recordStats()
+    .recordStats()
     .build[String, SimpleRasterInfo]
+
+  def cacheStats = cache.stats()
 }

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

@@ -33,7 +33,7 @@ import geotrellis.raster.resample.NearestNeighbor
  */
 class ExplodeSpec extends TestEnvironment with TestData {
   describe("conversion to/from exploded representation of tiles") {
-    import sqlContext.implicits._
+    import spark.implicits._
 
     it("should explode tiles") {
       val query = sql(

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

@@ -56,7 +56,7 @@ class RasterFrameSpec extends TestEnvironment with MetadataKeys
     }
     it("should provide Spark initialization methods") {
       assert(spark.withRasterFrames.isInstanceOf[SparkSession])
-      assert(sqlContext.withRasterFrames.isInstanceOf[SQLContext])
+      assert(spark.sqlContext.withRasterFrames.isInstanceOf[SQLContext])
     }
   }
 

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

@@ -20,46 +20,55 @@
  */
 package org.locationtech.rasterframes
 
-import java.nio.file.{Files, Paths}
+import java.io.File
+import java.nio.file.{Files, Path}
 
 import com.typesafe.scalalogging.LazyLogging
-import geotrellis.spark.testkit.{TestEnvironment => GeoTrellisTestEnvironment}
-import org.apache.spark.SparkContext
+import geotrellis.raster.testkit.RasterMatchers
+import org.apache.hadoop.fs.FileUtil
 import org.apache.spark.sql._
 import org.apache.spark.sql.functions.col
 import org.apache.spark.sql.types.StructType
+import org.apache.spark.{SparkConf, SparkContext}
 import org.locationtech.jts.geom.Geometry
+import org.locationtech.rasterframes.util._
 import org.scalactic.Tolerance
 import org.scalatest._
 import org.scalatest.matchers.{MatchResult, Matcher}
-import org.locationtech.rasterframes.util._
 
-trait TestEnvironment extends FunSpec with GeoTrellisTestEnvironment
-  with Matchers with Inspectors with Tolerance with LazyLogging {
+trait TestEnvironment extends FunSpec
+  with Matchers with Inspectors with Tolerance with RasterMatchers with LazyLogging {
 
-  override def sparkMaster: String = "local[*]"
+  lazy val scratchDir: Path = {
+    val outputDir = Files.createTempDirectory("rf-scratch-")
+    outputDir.toFile.deleteOnExit()
+    outputDir
+  }
 
-  override implicit def sc: SparkContext = { _sc.setLogLevel("ERROR"); _sc }
+  def sparkMaster: String = "local[*]"
 
-  lazy val sqlContext: SQLContext = {
+  def additionalConf = new SparkConf(false)
+
+  implicit lazy val spark: SparkSession = {
     val session = SparkSession.builder
-      .config(_sc.getConf)
-      .config("spark.sql.crossJoin.enabled", true)
+      .master(sparkMaster)
       .withKryoSerialization
+      .config(additionalConf)
       .getOrCreate()
-    session.sqlContext.withRasterFrames
+    session.withRasterFrames
   }
 
-  lazy val sql: String ⇒ DataFrame = sqlContext.sql
-  implicit lazy val spark: SparkSession = sqlContext.sparkSession
+  implicit def sc: SparkContext = spark.sparkContext
+
+  lazy val sql: String ⇒ DataFrame = spark.sql
 
   def isCI: Boolean = sys.env.get("CI").contains("true")
 
   /** This is here so we can test writing UDF generated/modified GeoTrellis types to ensure they are Parquet compliant. */
   def write(df: Dataset[_]): Boolean = {
     val sanitized = df.select(df.columns.map(c ⇒ col(c).as(toParquetFriendlyColumnName(c))): _*)
     val inRows = sanitized.count()
-    val dest = Files.createTempFile(Paths.get(outputLocalPath), "rf", ".parquet")
+    val dest = Files.createTempFile("rf", ".parquet")
     logger.trace(s"Writing '${sanitized.columns.mkString(", ")}' to '$dest'...")
     sanitized.write.mode(SaveMode.Overwrite).parquet(dest.toString)
     val in = df.sparkSession.read.parquet(dest.toString)

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

@@ -37,7 +37,7 @@ import org.locationtech.rasterframes.ref.{InMemoryRasterSource, RasterSource}
 class TileAssemblerSpec extends TestEnvironment {
   import TileAssemblerSpec._
   describe("TileAssembler") {
-    import sqlContext.implicits._
+    import spark.implicits._
 
     it("should reassemble a small scene") {
       val raster = TestData.l8Sample(8).projectedRaster

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

@@ -37,7 +37,7 @@ import org.locationtech.rasterframes.stats.CellHistogram
  */
 class TileStatsSpec extends TestEnvironment with TestData {
   import TestData.injectND
-  import sqlContext.implicits._
+  import spark.implicits._
 
   describe("computing statistics over tiles") {
     //import org.apache.spark.sql.execution.debug._
@@ -97,7 +97,7 @@ class TileStatsSpec extends TestEnvironment with TestData {
     }
 
     it("should support local min/max") {
-      import sqlContext.implicits._
+      import spark.implicits._
       val ds = Seq[Tile](byteArrayTile, byteConstantTile).toDF("tiles")
       ds.createOrReplaceTempView("tmp")
 
@@ -124,7 +124,7 @@ class TileStatsSpec extends TestEnvironment with TestData {
     }
 
     it("should compute tile statistics") {
-      import sqlContext.implicits._
+      import spark.implicits._
       withClue("mean") {
 
         val ds = Seq.fill[Tile](3)(randomTile(5, 5, FloatConstantNoDataCellType)).toDS()
@@ -229,7 +229,7 @@ class TileStatsSpec extends TestEnvironment with TestData {
     }
 
     it("should compute aggregate local stats") {
-      import sqlContext.implicits._
+      import spark.implicits._
       val ave = (nums: Array[Double]) => nums.sum / nums.length
 
       val ds = (Seq

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

@@ -43,7 +43,7 @@ import org.locationtech.rasterframes.tiles.ProjectedRasterTile
  */
 class EncodingSpec extends TestEnvironment with TestData {
 
-  import sqlContext.implicits._
+  import spark.implicits._
 
   describe("Spark encoding on standard types") {
 

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

@@ -20,8 +20,6 @@
  */
 package org.locationtech.rasterframes.datasource.geotrellis
 
-import java.io.File
-
 import org.locationtech.rasterframes._
 import geotrellis.proj4.LatLng
 import geotrellis.spark._
@@ -39,29 +37,27 @@ class GeoTrellisCatalogSpec
 
   lazy val testRdd = TestData.randomSpatioTemporalTileLayerRDD(10, 12, 5, 6)
 
-  import sqlContext.implicits._
+  import spark.implicits._
 
   before {
-    val outputDir = new File(outputLocalPath)
-    FileUtil.fullyDelete(outputDir)
-    outputDir.deleteOnExit()
-    lazy val writer = LayerWriter(outputDir.toURI)
+    FileUtil.fullyDelete(scratchDir.toFile)
+    lazy val writer = LayerWriter(scratchDir.toUri)
     val index =  ZCurveKeyIndexMethod.byDay()
     writer.write(LayerId("layer-1", 0), testRdd, index)
     writer.write(LayerId("layer-2", 0), testRdd, index)
   }
 
   describe("Catalog reading") {
     it("should show two zoom levels") {
-      val cat = sqlContext.read
-        .geotrellisCatalog(outputLocal.toUri)
+      val cat = spark.read
+        .geotrellisCatalog(scratchDir.toUri)
       assert(cat.schema.length > 4)
       assert(cat.count() === 2)
     }
 
     it("should support loading a layer in a nice way") {
-      val cat = sqlContext.read
-        .geotrellisCatalog(outputLocal.toUri)
+      val cat = spark.read
+        .geotrellisCatalog(scratchDir.toUri)
 
       // Select two layers.
       val layer = cat
@@ -70,7 +66,7 @@ class GeoTrellisCatalogSpec
         .collect
       assert(layer.length === 2)
 
-      val lots = layer.map(sqlContext.read.geotrellis.loadLayer).map(_.toDF).reduce(_ union _)
+      val lots = layer.map(spark.read.geotrellis.loadLayer).map(_.toDF).reduce(_ union _)
       assert(lots.count === 60)
     }
   }

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

@@ -56,9 +56,9 @@ class GeoTrellisDataSourceSpec
   import TestData._
 
   val tileSize = 12
-  lazy val layer = Layer(new File(outputLocalPath).toURI, LayerId("test-layer", 4))
-  lazy val tfLayer = Layer(new File(outputLocalPath).toURI, LayerId("test-tf-layer", 4))
-  lazy val sampleImageLayer = Layer(new File(outputLocalPath).toURI, LayerId("sample", 0))
+  lazy val layer = Layer(scratchDir.toUri, LayerId("test-layer", 4))
+  lazy val tfLayer = Layer(scratchDir.toUri, LayerId("test-tf-layer", 4))
+  lazy val sampleImageLayer = Layer(scratchDir.toUri, LayerId("sample", 0))
   val now = ZonedDateTime.now()
   val tileCoordRange = 2 to 5