Merge pull request #109 from jpolchlo/optimization/checkpoint-and-persist

Minor fixes and optimizations

Author: jpolchlo

README.md

@@ -129,6 +129,26 @@ calculation from geometric types.  See
 [here](https://www.geomesa.org/documentation/user/spark/sparksql_functions.html)
 for a list of functions that operate on geometries.
 
+#### A Note on Geocoding ####
+
+VectorPipe provides the means to tag geometries with the country codes of the
+countries they interact with, but it does not provide the boundaries used to
+do the coding.  That gives the user the option to select geometries
+appropriate to the task at hand—low resolution geometries for less fussy
+applications, high resolution when precision is important.
+
+In order for an application to make use of `vectorpipe.util.Geocode`, it must
+supply a `countries.geojson` in in the root of its project's `resources`
+directory.  That GeoJSON file must contain a `FeatureCollection`, with each
+entry having an `ADM0_A3` entry in its `properties` list.
+
+One may employ the [Natural Earth Admin
+0](https://www.naturalearthdata.com/downloads/10m-cultural-vectors/10m-admin-0-boundary-lines/)
+resource for low-precision tasks, or use something like the [Global LSIB
+Polygons](http://geonode.state.gov/layers/geonode%3AGlobal_LSIB_Polygons_Detailed)
+for more precise tasks (though the latter resource does not tag its elements
+with the `ADM0_A3` three-letter codes, so some preprocessing would be required).
+
 ## The `internal` package ##
 
 While most users will rely solely on the features exposed by the `OSM` object,

project/Version.scala

@@ -1,6 +1,6 @@
 object Version {
   val awscala = "0.8.1"
-  val vectorpipe = "1.1.0"
+  val vectorpipe = "1.1.0-SNAPSHOT"
   val scala = "2.11.12"
   val geotrellis = "2.2.0"
   val geomesa = "2.2.1"

src/main/scala/vectorpipe/VectorPipe.scala

@@ -13,6 +13,7 @@ import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types.StringType
+import org.apache.spark.storage.StorageLevel
 import org.locationtech.jts.{geom => jts}
 
 object VectorPipe {
@@ -25,6 +26,8 @@ object VectorPipe {
     * @param  srcCRS              CRS of the original geometry
     * @param  destCRS             (optional) The CRS to produce vectortiles into.  When omitted,
     *                             defaults to [[WebMercator]].
+    * @param  useCaching          Allows intermediate results to be cached to disk.  May require
+    *                             additional disk space on executor nodes.
     * @param  orderAreas          Sorts polygonal geometries in vectortiles.  In case of overlaps,
     *                             smaller geometries will draw on top of larger ones.
     * @param  tileResolution      Resolution of output tiles; i.e., the number of discretized bins
@@ -36,6 +39,7 @@ object VectorPipe {
     minZoom: Option[Int],
     srcCRS: CRS,
     destCRS: Option[CRS],
+    useCaching: Boolean = false,
     orderAreas: Boolean = false,
     tileResolution: Int = 4096
   )
@@ -141,8 +145,8 @@ object VectorPipe {
       val simplify = udf { g: jts.Geometry => pipeline.simplify(g, level.layout) }
       val reduced = pipeline
         .reduce(working, level, keyColumn)
-        .localCheckpoint()
-      val prepared = reduced
+      val persisted = if (options.useCaching) reduced.persist(StorageLevel.DISK_ONLY) else reduced
+      val prepared = persisted
         .withColumn(geomColumn, simplify(col(geomColumn)))
       val vts = generateVectorTiles(prepared, level)
       saveVectorTiles(vts, zoom, pipeline.baseOutputURI)

src/main/scala/vectorpipe/functions/osm/package.scala

@@ -277,6 +277,8 @@ package object osm {
     when(comment.isNotNull and length(comment) > 0, extractHashtags(comment))
       .otherwise(typedLit(Seq.empty[String])) as 'hashtags
 
+  def isTagged(tags: Column): Column = size(map_keys(tags)) > 0 as 'isTagged
+
   def isBuilding(tags: Column): Column =
     !lower(coalesce(tags.getItem("building"), lit("no"))).isin(FalsyValues: _*) as 'isBuilding
 

src/main/scala/vectorpipe/sources/AugmentedDiffSource.scala

@@ -3,6 +3,8 @@ package vectorpipe.sources
 import java.io.{ByteArrayInputStream, File}
 import java.net.URI
 import java.nio.charset.StandardCharsets
+import java.sql.Timestamp
+import java.time.Instant
 import java.util.zip.GZIPInputStream
 
 import cats.implicits._
@@ -16,6 +18,8 @@ import io.circe.generic.auto._
 import io.circe.{yaml, _}
 import org.apache.commons.io.IOUtils
 import org.apache.spark.internal.Logging
+import org.apache.spark.sql.Column
+import org.apache.spark.sql.functions.{floor, from_unixtime, to_timestamp, unix_timestamp}
 import org.joda.time.DateTime
 import vectorpipe.model.{AugmentedDiff, ElementWithSequence}
 
@@ -105,5 +109,17 @@ object AugmentedDiffSource extends Logging {
     }
   }
 
+  def timestampToSequence(timestamp: Timestamp): Int =
+    ((timestamp.toInstant.getEpochSecond - 1347432900) / 60).toInt
+
+  def timestampToSequence(timestamp: Column): Column =
+    floor((unix_timestamp(timestamp) - 1347432900) / 60)
+
+  def sequenceToTimestamp(sequence: Int): Timestamp =
+    Timestamp.from(Instant.ofEpochSecond(sequence.toLong * 60 + 1347432900L))
+
+  def sequenceToTimestamp(sequence: Column): Column =
+    to_timestamp(from_unixtime(sequence * 60 + 1347432900))
+
   case class State(last_run: DateTime, sequence: Int)
 }

src/main/scala/vectorpipe/sources/ChangesetMicroBatchReader.scala

@@ -20,7 +20,7 @@ class ChangesetStreamBatchReader(baseURI: URI, sequences: Seq[Int])
     extends ReplicationStreamBatchReader[Changeset](baseURI, sequences) {
 
   override def getSequence(baseURI: URI, sequence: Int): Seq[Changeset] =
-    ChangesetSource.getSequence(baseURI, sequence)
+    ChangesetSource.getChangeset(baseURI, sequence)
 }
 
 class ChangesetMicroBatchReader(options: DataSourceOptions, checkpointLocation: String)
@@ -32,7 +32,7 @@ class ChangesetMicroBatchReader(options: DataSourceOptions, checkpointLocation:
   )
 
   override def getCurrentSequence: Option[Int] =
-    ChangesetSource.getCurrentSequence(baseURI)
+    ChangesetSource.getCurrentSequence(baseURI).map(_.sequence.toInt)
 
   override def createDataReaderFactories(): util.List[DataReaderFactory[Row]] =
     sequenceRange

src/main/scala/vectorpipe/sources/ChangesetReader.scala

@@ -28,7 +28,7 @@ case class ChangesetReader(options: DataSourceOptions)
   }
 
   override protected def getCurrentSequence: Option[Int] =
-    ChangesetSource.getCurrentSequence(baseURI)
+    ChangesetSource.getCurrentSequence(baseURI).map(_.sequence.toInt)
 
   private def baseURI =
     new URI(

src/main/scala/vectorpipe/sources/ChangesetSource.scala

@@ -3,6 +3,7 @@ package vectorpipe.sources
 import java.io.{ByteArrayInputStream, IOException}
 import java.net.URI
 import java.nio.charset.StandardCharsets
+import java.time.Instant
 import java.util.zip.GZIPInputStream
 
 import cats.implicits._
@@ -27,7 +28,7 @@ object ChangesetSource extends Logging {
   private implicit val dateTimeDecoder: Decoder[DateTime] =
     Decoder.instance(a => a.as[String].map(DateTime.parse(_, formatter)))
 
-  def getSequence(baseURI: URI, sequence: Int): Seq[Changeset] = {
+  def getChangeset(baseURI: URI, sequence: Int): Seq[Changeset] = {
     val s = f"$sequence%09d"
     val path = s"${s.slice(0, 3)}/${s.slice(3, 6)}/${s.slice(6, 9)}.osm.gz"
 
@@ -40,7 +41,7 @@ object ChangesetSource extends Logging {
       if (response.code == 404) {
         logDebug(s"$sequence is not yet available, sleeping.")
         Thread.sleep(Delay.toMillis)
-        getSequence(baseURI, sequence)
+        getChangeset(baseURI, sequence)
       } else {
         // NOTE: if diff bodies get really large, switch to a SAX parser to help with the memory footprint
         val bais = new ByteArrayInputStream(response.body)
@@ -62,25 +63,27 @@ object ChangesetSource extends Logging {
       case e: IOException =>
         logWarning(s"Error fetching changeset $sequence", e)
         Thread.sleep(Delay.toMillis)
-        getSequence(baseURI, sequence)
+        getChangeset(baseURI, sequence)
     }
   }
 
+  case class Sequence(last_run: DateTime, sequence: Long)
+
   @memoize(maxSize = 1, expiresAfter = 30 seconds)
-  def getCurrentSequence(baseURI: URI): Option[Int] = {
+  def getCurrentSequence(baseURI: URI): Option[Sequence] = {
     try {
       val response =
         Http(baseURI.resolve("state.yaml").toString).asString
 
       val state = yaml.parser
         .parse(response.body)
         .leftMap(err => err: Error)
-        .flatMap(_.as[State])
+        .flatMap(_.as[Sequence])
         .valueOr(throw _)
 
       logDebug(s"$baseURI state: ${state.sequence} @ ${state.last_run}")
 
-      Some(state.sequence)
+      Some(state)
     } catch {
       case err: Throwable =>
         logError("Error fetching / parsing changeset state.", err)
@@ -89,5 +92,42 @@ object ChangesetSource extends Logging {
     }
   }
 
-  case class State(last_run: DateTime, sequence: Int)
+  def getSequence(baseURI: URI, sequence: Long): Option[Sequence] = {
+    val s = f"${sequence+1}%09d"
+    val path = s"${s.slice(0, 3)}/${s.slice(3, 6)}/${s.slice(6, 9)}.state.txt"
+
+    try {
+      val response =
+        Http(baseURI.resolve(path).toString).asString
+
+      val state = yaml.parser
+        .parse(response.body)
+        .leftMap(err => err: Error)
+        .flatMap(_.as[Sequence])
+        .valueOr(throw _)
+
+      Some(state)
+    } catch {
+      case err: Throwable =>
+        logError("Error fetching / parsing changeset state.", err)
+
+        None
+    }
+  }
+
+  def estimateSequenceNumber(modifiedTime: Instant, baseURI: URI): Long = {
+    val current = getCurrentSequence(baseURI)
+    val diffMinutes = (current.get.last_run.toInstant.getMillis/1000 - modifiedTime.getEpochSecond) / 60
+    current.get.sequence - diffMinutes
+  }
+
+  def findSequenceFor(modifiedTime: Instant, baseURI: URI): Long = {
+    var guess = estimateSequenceNumber(modifiedTime, baseURI)
+    val target = org.joda.time.Instant.parse(modifiedTime.toString)
+
+    while (getSequence(baseURI, guess).get.last_run.isAfter(target)) { guess -= 1 }
+    while (getSequence(baseURI, guess).get.last_run.isBefore(target)) { guess += 1 }
+
+    getSequence(baseURI, guess).get.sequence
+  }
 }

src/main/scala/vectorpipe/vectortile/Pipeline.scala

@@ -3,10 +3,8 @@ package vectorpipe.vectortile
 import geotrellis.spark.SpatialKey
 import geotrellis.spark.tiling._
 import geotrellis.vector.{Feature, Geometry}
-import geotrellis.vectortile.Value
 import org.apache.spark.sql.{DataFrame, Row}
 import org.locationtech.jts.{geom => jts}
-import org.locationtech.jts.simplify
 
 /**
  * The interface governing the transformation from processed OSM dataframes to

src/main/scala/vectorpipe/vectortile/package.scala

@@ -4,7 +4,6 @@ import geotrellis.proj4._
 import geotrellis.spark.SpatialKey
 import geotrellis.spark.tiling.LayoutDefinition
 import geotrellis.vector._
-import geotrellis.vector.reproject._
 import geotrellis.vectortile._
 import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema
@@ -27,15 +26,19 @@ package object vectortile {
   @transient lazy val st_reprojectGeom = udf { (g: jts.Geometry, srcProj: String, destProj: String) =>
     val trans = Proj4Transform(CRS.fromString(srcProj), CRS.fromString(destProj))
     if (Option(g).isDefined) {
-      val gt = Geometry(g)
-      gt.reproject(trans).jtsGeom
+      if (g.isEmpty)
+        g
+      else {
+        val gt = Geometry(g)
+        gt.reproject(trans).jtsGeom
+      }
     } else {
       null
     }
   }
 
   def keyTo(layout: LayoutDefinition) = udf { g: jts.Geometry =>
-    if (Option(g).isDefined) {
+    if (Option(g).isDefined && !g.isEmpty) {
       layout.mapTransform.keysForGeometry(geotrellis.vector.Geometry(g)).toArray
     } else {
       Array.empty[SpatialKey]