Add buffer operation (#198)

build.sbt

@@ -29,7 +29,7 @@ libraryDependencies ++= Seq(
   "com.lihaoyi" % "fastparse_2.11" % "0.4.3" % "provided",
   "org.scalatest" %% "scalatest" % "2.2.1" % "test",
   "com.vividsolutions" % "jts" % "1.13" % "test",
-  "com.esri.geometry" % "esri-geometry-api" % "1.2.1" % "test"
+  "com.esri.geometry" % "esri-geometry-api" % "1.2.1"
 )
 
 libraryDependencies ++= Seq(

src/main/scala/magellan/Shape.scala

@@ -16,7 +16,9 @@
 
 package magellan
 
+import com.esri.core.geometry.OperatorBuffer
 import com.fasterxml.jackson.annotation.{JsonIgnore, JsonProperty}
+import magellan.esri.ESRIUtil
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.types._
 
@@ -173,6 +175,14 @@ trait Shape extends DataType with Serializable {
    */
   @JsonIgnore
   def isEmpty(): Boolean
+
+  def buffer(distance: Double): Polygon = {
+    val esriGeometry = ESRIUtil.toESRIGeometry(this)
+    val bufferedEsriGeometry = OperatorBuffer.local()
+      .execute(esriGeometry, null, distance, null)
+
+    ESRIUtil.fromESRIGeometry(bufferedEsriGeometry).asInstanceOf[Polygon]
+  }
 }
 
 /**

src/main/scala/magellan/dsl/package.scala

@@ -49,6 +49,8 @@ package object dsl {
       def withinRange(origin: Point, radius: Double): Column = Column(WithinCircleRange(c.expr, origin, radius))
 
       def asGeoJSON(): Column = Column(AsGeoJSON(c.expr))
+
+      def buffer(distance: Double): Column = Column(Buffer(c.expr, distance))
     }
 
     implicit def point(x: Column, y: Column) = Column(PointConverter(x.expr, y.expr))
@@ -57,6 +59,8 @@ package object dsl {
 
     implicit def asGeoJSON(x: Column) = Column(AsGeoJSON(x.expr))
 
+    implicit def buffer(x: Column, distance: Double) = Column(Buffer(x.expr, distance))
+
     implicit class DslDataset[T](c: Dataset[T]) {
       def df: Dataset[T] = c
 

src/main/scala/magellan/esri/ESRIUtil.scala

@@ -0,0 +1,206 @@
+package magellan.esri
+
+import com.esri.core.geometry.{Geometry => ESRIGeometry, Point => ESRIPoint, Polygon => ESRIPolygon, Polyline => ESRIPolyLine}
+import magellan._
+import scala.collection.mutable.ArrayBuffer
+
+object ESRIUtil {
+
+  def fromESRIGeometry(geometry: ESRIGeometry): Shape = {
+    geometry match {
+      case esriPoint: ESRIPoint => fromESRI(esriPoint)
+      case esriPolyline: ESRIPolyLine => fromESRI(esriPolyline)
+      case esriPolygon: ESRIPolygon => fromESRI(esriPolygon)
+    }
+  }
+
+  def toESRIGeometry(shape: Shape): ESRIGeometry = {
+    shape match {
+      case point: Point => toESRI(point)
+      case polyline: PolyLine => toESRI(polyline)
+      case polygon: Polygon => toESRI(polygon)
+    }
+  }
+
+  /**
+   * Convert ESRI 2D Point to Magellan Point.
+   *
+   * @param esriPoint
+   * @return
+   */
+  def fromESRI(esriPoint: ESRIPoint): Point = {
+    Point(esriPoint.getX, esriPoint.getY)
+  }
+
+  /**
+   * Convert Magellan Point to ESRI 2D Point
+   *
+   * @param point
+   * @return
+   */
+  def toESRI(point: Point): ESRIPoint = {
+    val esriPoint = new ESRIPoint()
+    esriPoint.setXY(point.getX(), point.getY())
+    esriPoint
+  }
+
+  /**
+   * Convert ESRI PolyLine to Magellan PolyLine.
+   *
+   * @param esriPolyLine
+   * @return
+   */
+  def fromESRI(esriPolyLine: ESRIPolyLine): PolyLine = {
+    val length = esriPolyLine.getPointCount
+    if (length == 0) {
+      PolyLine(Array[Int](), Array[Point]())
+    } else {
+      val indices = ArrayBuffer[Int]()
+      indices.+=(0)
+      val points = ArrayBuffer[Point]()
+      var start = esriPolyLine.getPoint(0)
+      var currentRingIndex = 0
+      points.+=(Point(start.getX(), start.getY()))
+
+      for (i <- (1 until length)) {
+        val p = esriPolyLine.getPoint(i)
+        val j = esriPolyLine.getPathEnd(currentRingIndex)
+        if (j < length) {
+          val end = esriPolyLine.getPoint(j)
+          if (p.getX == end.getX && p.getY == end.getY) {
+            indices.+=(i)
+            currentRingIndex += 1
+            // add start point
+            points.+= (Point(start.getX(), start.getY()))
+            start = end
+          }
+        }
+        points.+=(Point(p.getX(), p.getY()))
+      }
+      PolyLine(indices.toArray, points.toArray)
+    }
+  }
+
+  /**
+   * Convert Magellan PolyLine to ESRI PolyLine.
+   *
+   * @param polyline
+   * @return
+   */
+  def toESRI(polyline: PolyLine): ESRIPolyLine = {
+    val l = new ESRIPolyLine()
+    val indices = polyline.getRings()
+    val length = polyline.length
+    if (length > 0) {
+      var startIndex = 0
+      var endIndex = 1
+      var currentRingIndex = 0
+      val startVertex = polyline.getVertex(startIndex)
+      l.startPath(
+        startVertex.getX(),
+        startVertex.getY())
+
+      while (endIndex < length) {
+        val endVertex = polyline.getVertex(endIndex)
+        l.lineTo(endVertex.getX(), endVertex.getY())
+        startIndex += 1
+        endIndex += 1
+        // if we reach a ring boundary skip it
+        val nextRingIndex = currentRingIndex + 1
+        if (nextRingIndex < indices.length) {
+          val nextRing = indices(nextRingIndex)
+          if (endIndex == nextRing) {
+            startIndex += 1
+            endIndex += 1
+            currentRingIndex = nextRingIndex
+            val startVertex = polyline.getVertex(startIndex)
+            l.startPath(
+              startVertex.getX(),
+              startVertex.getY())
+          }
+        }
+      }
+    }
+    l
+  }
+
+  /**
+   * Convert ESRI Polygon to Magellan Polygon.
+   *
+   * @param esriPolygon
+   * @return
+   */
+  def fromESRI(esriPolygon: ESRIPolygon): Polygon = {
+    val length = esriPolygon.getPointCount
+    if (length == 0) {
+      Polygon(Array[Int](), Array[Point]())
+    } else {
+      val indices = ArrayBuffer[Int]()
+      indices.+=(0)
+      val points = ArrayBuffer[Point]()
+      var start = esriPolygon.getPoint(0)
+      var currentRingIndex = 0
+      points.+=(Point(start.getX(), start.getY()))
+
+      for (i <- (1 until length)) {
+        val p = esriPolygon.getPoint(i)
+        val j = esriPolygon.getPathEnd(currentRingIndex)
+        if (j < length) {
+          val end = esriPolygon.getPoint(j)
+          if (p.getX == end.getX && p.getY == end.getY) {
+            indices.+=(i)
+            currentRingIndex += 1
+            // add start point
+            points.+= (Point(start.getX(), start.getY()))
+            start = end
+          }
+        }
+        points.+=(Point(p.getX(), p.getY()))
+      }
+      Polygon(indices.toArray, points.toArray)
+    }
+  }
+
+  /**
+   * Convert Magellan Polygon to ESRI Polygon.
+   *
+   * @param polygon
+   * @return
+   */
+  def toESRI(polygon: Polygon): ESRIPolygon = {
+    val p = new ESRIPolygon()
+    val indices = polygon.getRings()
+    val length = polygon.length
+    if (length > 0) {
+      var startIndex = 0
+      var endIndex = 1
+      var currentRingIndex = 0
+      val startVertex = polygon.getVertex(startIndex)
+      p.startPath(
+        startVertex.getX(),
+        startVertex.getY())
+
+      while (endIndex < length) {
+        val endVertex = polygon.getVertex(endIndex)
+        p.lineTo(endVertex.getX(), endVertex.getY())
+        startIndex += 1
+        endIndex += 1
+        // if we reach a ring boundary skip it
+        val nextRingIndex = currentRingIndex + 1
+        if (nextRingIndex < indices.length) {
+          val nextRing = indices(nextRingIndex)
+          if (endIndex == nextRing) {
+            startIndex += 1
+            endIndex += 1
+            currentRingIndex = nextRingIndex
+            val startVertex = polygon.getVertex(startIndex)
+            p.startPath(
+              startVertex.getX(),
+              startVertex.getY())
+          }
+        }
+      }
+    }
+    p
+  }
+}

src/main/scala/org/apache/spark/sql/catalyst/expressions/functions.scala

@@ -285,6 +285,59 @@ case class AsGeoJSON(override val child: Expression)
   }
 }
 
+case class Buffer(override val child: Expression, distance: Double)
+  extends UnaryExpression with MagellanExpression {
+
+  override def nullable: Boolean = false
+
+  override def dataType: DataType = new PolygonUDT()
+
+  override protected def nullSafeEval(input: Any): Any = {
+    val shape = newInstance(input.asInstanceOf[InternalRow])
+    val bufferedShape = shape.buffer(distance)
+    serialize(bufferedShape)
+  }
+
+  override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    val childTypeVar = ctx.freshName("childType")
+    val childShapeVar = ctx.freshName("childShape")
+    val shapeSerializerVar = ctx.freshName("shapeSerializer")
+    val resultSerializerVar = ctx.freshName("resultSerializer")
+    val distanceVar = ctx.freshName("distance")
+    val resultVar = ctx.freshName("result")
+    val resultTypeVar = ctx.freshName("resultType")
+    val serializersVar = ctx.freshName("serializers")
+    val idx = ctx.references.length
+    ctx.addReferenceObj("distance", distance);
+
+    ctx.addMutableState(classOf[java.util.HashMap[Integer, UserDefinedType[Shape]]].getName, s"$serializersVar",
+      s"$serializersVar = new java.util.HashMap<Integer, org.apache.spark.sql.types.UserDefinedType<magellan.Shape>>() ;" +
+        s"$serializersVar.put(1, new org.apache.spark.sql.types.PointUDT());" +
+        s"$serializersVar.put(2, new org.apache.spark.sql.types.LineUDT());" +
+        s"$serializersVar.put(3, new org.apache.spark.sql.types.PolyLineUDT());" +
+        s"$serializersVar.put(5, new org.apache.spark.sql.types.PolygonUDT());" +
+        "")
+
+    nullSafeCodeGen(ctx, ev, (c1) => {
+      s"" +
+        s"Integer $childTypeVar = $c1.getInt(0); \n" +
+        s"Double $distanceVar = (Double) references[$idx]; \n" +
+        s"org.apache.spark.sql.types.UserDefinedType<magellan.Shape> $shapeSerializerVar = " +
+        s"((org.apache.spark.sql.types.UserDefinedType<magellan.Shape>)" +
+        s"$serializersVar.get($childTypeVar)); \n" +
+        s"magellan.Shape $childShapeVar = (magellan.Shape)" +
+        s"$shapeSerializerVar.deserialize($c1); \n" +
+        s"magellan.Shape $resultVar = $childShapeVar.buffer($distanceVar); \n" +
+        s"Integer $resultTypeVar = $resultVar.getType(); \n" +
+        s"org.apache.spark.sql.types.UserDefinedType<magellan.Shape> $resultSerializerVar = " +
+        s"((org.apache.spark.sql.types.UserDefinedType<magellan.Shape>)" +
+        s"$serializersVar.get($resultTypeVar)); \n" +
+        s"${ev.value} = (org.apache.spark.sql.catalyst.InternalRow)$resultSerializerVar.serialize($resultVar); \n"
+    })
+
+  }
+}
+
 object Indexer {
 
   val indexUDT = new ZOrderCurveUDT()

src/test/scala/magellan/PointSuite.scala

@@ -69,6 +69,16 @@ class PointSuite extends FunSuite with TestSparkContext {
   test("within circle") {
     assert(Point(0.0, 0.0) withinCircle (Point(0.5, 0.5), 0.75))
     assert(!(Point(0.0, 0.0) withinCircle (Point(0.5, 0.5), 0.5)))
+  }
 
+  test("buffer point") {
+    val polygon = Point(0.0, 1.0).buffer(0.5)
+    assert(polygon.getNumRings() === 1)
+    // check that [0.0, 0.75] is within this polygon
+    assert(polygon.contains(Point(0.0, 0.75)))
+    // check that [0.4, 1.0] is within this polygon
+    assert(polygon.contains(Point(0.4, 1.0)))
+    // check that [0.6, 1.0] is outside this polygon
+    assert(!polygon.contains(Point(0.6, 1.0)))
   }
 }

src/test/scala/magellan/PolyLineSuite.scala

@@ -115,4 +115,24 @@ class PolyLineSuite extends FunSuite with TestSparkContext {
     val y = PolyLine(Array(0), Array(Point(0.5, 0.0), Point(0.0, 0.5)))
     assert(y.contains(Point(0.5, 0.0)))
   }
+
+  test("buffer polyline") {
+    /**
+     *    +-------+ 1.5,1.5
+     *    +----+  +
+     *         +  +
+     *    +----+  +
+     *    +-------+
+     *
+     */
+
+    val ring = Array(Point(-1.0, 1.0), Point(1.0, 1.0),
+      Point(1.0, -1.0), Point(-1.0, -1.0))
+    val polyline = PolyLine(Array(0), ring)
+
+    val bufferedPolygon = polyline.buffer(0.5)
+    assert(bufferedPolygon.getNumRings() === 1)
+    assert(bufferedPolygon.contains(Point(1.3, 1.3)))
+    assert(!bufferedPolygon.contains(Point(0.5, 0.5)))
+  }
 }